write_buffer.go
2.4 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
package pool
import (
"encoding/binary"
"io"
"sync"
)
var pool = sync.Pool{
New: func() interface{} {
return NewWriteBuffer()
},
}
func GetWriteBuffer() *WriteBuffer {
wb := pool.Get().(*WriteBuffer)
wb.Reset()
return wb
}
func PutWriteBuffer(wb *WriteBuffer) {
pool.Put(wb)
}
type WriteBuffer struct {
Bytes []byte
msgStart int
paramStart int
}
func NewWriteBuffer() *WriteBuffer {
return &WriteBuffer{
Bytes: make([]byte, 0, defaultBufSize),
}
}
func (buf *WriteBuffer) Reset() {
buf.Bytes = buf.Bytes[:0]
}
func (buf *WriteBuffer) ResetBuffer(b []byte) {
buf.Bytes = b[:0]
}
func (buf *WriteBuffer) StartMessage(c byte) {
if c == 0 {
buf.msgStart = len(buf.Bytes)
buf.Bytes = append(buf.Bytes, 0, 0, 0, 0)
} else {
buf.msgStart = len(buf.Bytes) + 1
buf.Bytes = append(buf.Bytes, c, 0, 0, 0, 0)
}
}
func (buf *WriteBuffer) FinishMessage() {
binary.BigEndian.PutUint32(
buf.Bytes[buf.msgStart:], uint32(len(buf.Bytes)-buf.msgStart))
}
func (buf *WriteBuffer) Query() []byte {
return buf.Bytes[buf.msgStart+4 : len(buf.Bytes)-1]
}
func (buf *WriteBuffer) StartParam() {
buf.paramStart = len(buf.Bytes)
buf.Bytes = append(buf.Bytes, 0, 0, 0, 0)
}
func (buf *WriteBuffer) FinishParam() {
binary.BigEndian.PutUint32(
buf.Bytes[buf.paramStart:], uint32(len(buf.Bytes)-buf.paramStart-4))
}
var nullParamLength = int32(-1)
func (buf *WriteBuffer) FinishNullParam() {
binary.BigEndian.PutUint32(
buf.Bytes[buf.paramStart:], uint32(nullParamLength))
}
func (buf *WriteBuffer) Write(b []byte) (int, error) {
buf.Bytes = append(buf.Bytes, b...)
return len(b), nil
}
func (buf *WriteBuffer) WriteInt16(num int16) {
buf.Bytes = append(buf.Bytes, 0, 0)
binary.BigEndian.PutUint16(buf.Bytes[len(buf.Bytes)-2:], uint16(num))
}
func (buf *WriteBuffer) WriteInt32(num int32) {
buf.Bytes = append(buf.Bytes, 0, 0, 0, 0)
binary.BigEndian.PutUint32(buf.Bytes[len(buf.Bytes)-4:], uint32(num))
}
func (buf *WriteBuffer) WriteString(s string) {
buf.Bytes = append(buf.Bytes, s...)
buf.Bytes = append(buf.Bytes, 0)
}
func (buf *WriteBuffer) WriteBytes(b []byte) {
buf.Bytes = append(buf.Bytes, b...)
buf.Bytes = append(buf.Bytes, 0)
}
func (buf *WriteBuffer) WriteByte(c byte) error {
buf.Bytes = append(buf.Bytes, c)
return nil
}
func (buf *WriteBuffer) ReadFrom(r io.Reader) (int64, error) {
n, err := r.Read(buf.Bytes[len(buf.Bytes):cap(buf.Bytes)])
buf.Bytes = buf.Bytes[:len(buf.Bytes)+n]
return int64(n), err
}