github.com/batchcorp/thrift-iterator@v0.0.0-20220918180557-4c4a158fc6e9/protocol/binary/stream.go (about) 1 package binary 2 3 import ( 4 "fmt" 5 "github.com/batchcorp/thrift-iterator/protocol" 6 "github.com/batchcorp/thrift-iterator/spi" 7 "io" 8 "math" 9 ) 10 11 type Stream struct { 12 spi.ValEncoderProvider 13 writer io.Writer 14 buf []byte 15 err error 16 } 17 18 func NewStream(provider spi.ValEncoderProvider, writer io.Writer, buf []byte) *Stream { 19 return &Stream{ 20 ValEncoderProvider: provider, 21 writer: writer, 22 buf: buf, 23 } 24 } 25 26 func (stream *Stream) Spawn() spi.Stream { 27 return &Stream{ 28 ValEncoderProvider: stream.ValEncoderProvider, 29 } 30 } 31 32 func (stream *Stream) Error() error { 33 return stream.err 34 } 35 36 func (stream *Stream) ReportError(operation string, err string) { 37 if stream.err == nil { 38 stream.err = fmt.Errorf("%s: %s", operation, err) 39 } 40 } 41 42 func (stream *Stream) Buffer() []byte { 43 return stream.buf 44 } 45 46 func (stream *Stream) Reset(writer io.Writer) { 47 stream.writer = writer 48 stream.err = nil 49 stream.buf = stream.buf[:0] 50 } 51 52 func (stream *Stream) Flush() { 53 if stream.writer == nil { 54 return 55 } 56 _, err := stream.writer.Write(stream.buf) 57 if err != nil { 58 stream.ReportError("Flush", err.Error()) 59 return 60 } 61 if f, ok := stream.writer.(protocol.Flusher); ok { 62 if err = f.Flush(); err != nil { 63 stream.ReportError("Flush", err.Error()) 64 } 65 } 66 stream.buf = stream.buf[:0] 67 } 68 69 func (stream *Stream) Write(buf []byte) error { 70 stream.buf = append(stream.buf, buf...) 71 stream.Flush() 72 return stream.Error() 73 } 74 75 func (stream *Stream) WriteMessageHeader(header protocol.MessageHeader) { 76 versionAndMessageType := uint32(protocol.BINARY_VERSION_1) | uint32(header.MessageType) 77 stream.WriteUint32(versionAndMessageType) 78 stream.WriteString(header.MessageName) 79 stream.WriteInt32(int32(header.SeqId)) 80 } 81 82 func (stream *Stream) WriteListHeader(elemType protocol.TType, length int) { 83 stream.buf = append(stream.buf, byte(elemType), 84 byte(length>>24), byte(length>>16), byte(length>>8), byte(length)) 85 } 86 87 func (stream *Stream) WriteStructHeader() { 88 } 89 90 func (stream *Stream) WriteStructField(fieldType protocol.TType, fieldId protocol.FieldId) { 91 stream.buf = append(stream.buf, byte(fieldType), byte(fieldId>>8), byte(fieldId)) 92 } 93 94 func (stream *Stream) WriteStructFieldStop() { 95 stream.buf = append(stream.buf, byte(protocol.TypeStop)) 96 } 97 98 func (stream *Stream) WriteMapHeader(keyType protocol.TType, elemType protocol.TType, length int) { 99 stream.buf = append(stream.buf, byte(keyType), byte(elemType), 100 byte(length>>24), byte(length>>16), byte(length>>8), byte(length)) 101 } 102 103 func (stream *Stream) WriteBool(val bool) { 104 if val { 105 stream.WriteUint8(1) 106 } else { 107 stream.WriteUint8(0) 108 } 109 } 110 111 func (stream *Stream) WriteInt8(val int8) { 112 stream.WriteUint8(uint8(val)) 113 } 114 115 func (stream *Stream) WriteUint8(val uint8) { 116 stream.buf = append(stream.buf, byte(val)) 117 } 118 119 func (stream *Stream) WriteInt16(val int16) { 120 stream.WriteUint16(uint16(val)) 121 } 122 123 func (stream *Stream) WriteUint16(val uint16) { 124 stream.buf = append(stream.buf, byte(val>>8), byte(val)) 125 } 126 127 func (stream *Stream) WriteInt32(val int32) { 128 stream.WriteUint32(uint32(val)) 129 } 130 131 func (stream *Stream) WriteUint32(val uint32) { 132 stream.buf = append(stream.buf, byte(val>>24), byte(val>>16), byte(val>>8), byte(val)) 133 } 134 135 func (stream *Stream) WriteInt64(val int64) { 136 stream.WriteUint64(uint64(val)) 137 } 138 139 func (stream *Stream) WriteUint64(val uint64) { 140 stream.buf = append(stream.buf, 141 byte(val>>56), byte(val>>48), byte(val>>40), byte(val>>32), 142 byte(val>>24), byte(val>>16), byte(val>>8), byte(val)) 143 } 144 145 func (stream *Stream) WriteInt(val int) { 146 stream.WriteInt64(int64(val)) 147 } 148 149 func (stream *Stream) WriteUint(val uint) { 150 stream.WriteUint64(uint64(val)) 151 } 152 153 func (stream *Stream) WriteFloat64(val float64) { 154 stream.WriteUint64(math.Float64bits(val)) 155 } 156 157 func (stream *Stream) WriteBinary(val []byte) { 158 stream.WriteUint32(uint32(len(val))) 159 stream.buf = append(stream.buf, val...) 160 } 161 162 func (stream *Stream) WriteString(val string) { 163 stream.WriteUint32(uint32(len(val))) 164 stream.buf = append(stream.buf, val...) 165 }