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  }