[VOL-4291] Rw-core updates for gRPC migration

Change-Id: I8d5a554409115b29318089671ca4e1ab3fa98810
diff --git a/vendor/github.com/uber/jaeger-client-go/thrift/binary_protocol.go b/vendor/github.com/uber/jaeger-client-go/thrift/binary_protocol.go
index 690d341..45c880d 100644
--- a/vendor/github.com/uber/jaeger-client-go/thrift/binary_protocol.go
+++ b/vendor/github.com/uber/jaeger-client-go/thrift/binary_protocol.go
@@ -21,6 +21,7 @@
 
 import (
 	"bytes"
+	"context"
 	"encoding/binary"
 	"errors"
 	"fmt"
@@ -31,190 +32,220 @@
 type TBinaryProtocol struct {
 	trans         TRichTransport
 	origTransport TTransport
-	reader        io.Reader
-	writer        io.Writer
-	strictRead    bool
-	strictWrite   bool
+	cfg           *TConfiguration
 	buffer        [64]byte
 }
 
 type TBinaryProtocolFactory struct {
-	strictRead  bool
-	strictWrite bool
+	cfg *TConfiguration
 }
 
+// Deprecated: Use NewTBinaryProtocolConf instead.
 func NewTBinaryProtocolTransport(t TTransport) *TBinaryProtocol {
-	return NewTBinaryProtocol(t, false, true)
+	return NewTBinaryProtocolConf(t, &TConfiguration{
+		noPropagation: true,
+	})
 }
 
+// Deprecated: Use NewTBinaryProtocolConf instead.
 func NewTBinaryProtocol(t TTransport, strictRead, strictWrite bool) *TBinaryProtocol {
-	p := &TBinaryProtocol{origTransport: t, strictRead: strictRead, strictWrite: strictWrite}
+	return NewTBinaryProtocolConf(t, &TConfiguration{
+		TBinaryStrictRead:  &strictRead,
+		TBinaryStrictWrite: &strictWrite,
+
+		noPropagation: true,
+	})
+}
+
+func NewTBinaryProtocolConf(t TTransport, conf *TConfiguration) *TBinaryProtocol {
+	PropagateTConfiguration(t, conf)
+	p := &TBinaryProtocol{
+		origTransport: t,
+		cfg:           conf,
+	}
 	if et, ok := t.(TRichTransport); ok {
 		p.trans = et
 	} else {
 		p.trans = NewTRichTransport(t)
 	}
-	p.reader = p.trans
-	p.writer = p.trans
 	return p
 }
 
+// Deprecated: Use NewTBinaryProtocolFactoryConf instead.
 func NewTBinaryProtocolFactoryDefault() *TBinaryProtocolFactory {
-	return NewTBinaryProtocolFactory(false, true)
+	return NewTBinaryProtocolFactoryConf(&TConfiguration{
+		noPropagation: true,
+	})
 }
 
+// Deprecated: Use NewTBinaryProtocolFactoryConf instead.
 func NewTBinaryProtocolFactory(strictRead, strictWrite bool) *TBinaryProtocolFactory {
-	return &TBinaryProtocolFactory{strictRead: strictRead, strictWrite: strictWrite}
+	return NewTBinaryProtocolFactoryConf(&TConfiguration{
+		TBinaryStrictRead:  &strictRead,
+		TBinaryStrictWrite: &strictWrite,
+
+		noPropagation: true,
+	})
+}
+
+func NewTBinaryProtocolFactoryConf(conf *TConfiguration) *TBinaryProtocolFactory {
+	return &TBinaryProtocolFactory{
+		cfg: conf,
+	}
 }
 
 func (p *TBinaryProtocolFactory) GetProtocol(t TTransport) TProtocol {
-	return NewTBinaryProtocol(t, p.strictRead, p.strictWrite)
+	return NewTBinaryProtocolConf(t, p.cfg)
+}
+
+func (p *TBinaryProtocolFactory) SetTConfiguration(conf *TConfiguration) {
+	p.cfg = conf
 }
 
 /**
  * Writing Methods
  */
 
-func (p *TBinaryProtocol) WriteMessageBegin(name string, typeId TMessageType, seqId int32) error {
-	if p.strictWrite {
+func (p *TBinaryProtocol) WriteMessageBegin(ctx context.Context, name string, typeId TMessageType, seqId int32) error {
+	if p.cfg.GetTBinaryStrictWrite() {
 		version := uint32(VERSION_1) | uint32(typeId)
-		e := p.WriteI32(int32(version))
+		e := p.WriteI32(ctx, int32(version))
 		if e != nil {
 			return e
 		}
-		e = p.WriteString(name)
+		e = p.WriteString(ctx, name)
 		if e != nil {
 			return e
 		}
-		e = p.WriteI32(seqId)
+		e = p.WriteI32(ctx, seqId)
 		return e
 	} else {
-		e := p.WriteString(name)
+		e := p.WriteString(ctx, name)
 		if e != nil {
 			return e
 		}
-		e = p.WriteByte(int8(typeId))
+		e = p.WriteByte(ctx, int8(typeId))
 		if e != nil {
 			return e
 		}
-		e = p.WriteI32(seqId)
+		e = p.WriteI32(ctx, seqId)
 		return e
 	}
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteMessageEnd() error {
+func (p *TBinaryProtocol) WriteMessageEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteStructBegin(name string) error {
+func (p *TBinaryProtocol) WriteStructBegin(ctx context.Context, name string) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteStructEnd() error {
+func (p *TBinaryProtocol) WriteStructEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteFieldBegin(name string, typeId TType, id int16) error {
-	e := p.WriteByte(int8(typeId))
+func (p *TBinaryProtocol) WriteFieldBegin(ctx context.Context, name string, typeId TType, id int16) error {
+	e := p.WriteByte(ctx, int8(typeId))
 	if e != nil {
 		return e
 	}
-	e = p.WriteI16(id)
+	e = p.WriteI16(ctx, id)
 	return e
 }
 
-func (p *TBinaryProtocol) WriteFieldEnd() error {
+func (p *TBinaryProtocol) WriteFieldEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteFieldStop() error {
-	e := p.WriteByte(STOP)
+func (p *TBinaryProtocol) WriteFieldStop(ctx context.Context) error {
+	e := p.WriteByte(ctx, STOP)
 	return e
 }
 
-func (p *TBinaryProtocol) WriteMapBegin(keyType TType, valueType TType, size int) error {
-	e := p.WriteByte(int8(keyType))
+func (p *TBinaryProtocol) WriteMapBegin(ctx context.Context, keyType TType, valueType TType, size int) error {
+	e := p.WriteByte(ctx, int8(keyType))
 	if e != nil {
 		return e
 	}
-	e = p.WriteByte(int8(valueType))
+	e = p.WriteByte(ctx, int8(valueType))
 	if e != nil {
 		return e
 	}
-	e = p.WriteI32(int32(size))
+	e = p.WriteI32(ctx, int32(size))
 	return e
 }
 
-func (p *TBinaryProtocol) WriteMapEnd() error {
+func (p *TBinaryProtocol) WriteMapEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteListBegin(elemType TType, size int) error {
-	e := p.WriteByte(int8(elemType))
+func (p *TBinaryProtocol) WriteListBegin(ctx context.Context, elemType TType, size int) error {
+	e := p.WriteByte(ctx, int8(elemType))
 	if e != nil {
 		return e
 	}
-	e = p.WriteI32(int32(size))
+	e = p.WriteI32(ctx, int32(size))
 	return e
 }
 
-func (p *TBinaryProtocol) WriteListEnd() error {
+func (p *TBinaryProtocol) WriteListEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteSetBegin(elemType TType, size int) error {
-	e := p.WriteByte(int8(elemType))
+func (p *TBinaryProtocol) WriteSetBegin(ctx context.Context, elemType TType, size int) error {
+	e := p.WriteByte(ctx, int8(elemType))
 	if e != nil {
 		return e
 	}
-	e = p.WriteI32(int32(size))
+	e = p.WriteI32(ctx, int32(size))
 	return e
 }
 
-func (p *TBinaryProtocol) WriteSetEnd() error {
+func (p *TBinaryProtocol) WriteSetEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) WriteBool(value bool) error {
+func (p *TBinaryProtocol) WriteBool(ctx context.Context, value bool) error {
 	if value {
-		return p.WriteByte(1)
+		return p.WriteByte(ctx, 1)
 	}
-	return p.WriteByte(0)
+	return p.WriteByte(ctx, 0)
 }
 
-func (p *TBinaryProtocol) WriteByte(value int8) error {
+func (p *TBinaryProtocol) WriteByte(ctx context.Context, value int8) error {
 	e := p.trans.WriteByte(byte(value))
 	return NewTProtocolException(e)
 }
 
-func (p *TBinaryProtocol) WriteI16(value int16) error {
+func (p *TBinaryProtocol) WriteI16(ctx context.Context, value int16) error {
 	v := p.buffer[0:2]
 	binary.BigEndian.PutUint16(v, uint16(value))
-	_, e := p.writer.Write(v)
+	_, e := p.trans.Write(v)
 	return NewTProtocolException(e)
 }
 
-func (p *TBinaryProtocol) WriteI32(value int32) error {
+func (p *TBinaryProtocol) WriteI32(ctx context.Context, value int32) error {
 	v := p.buffer[0:4]
 	binary.BigEndian.PutUint32(v, uint32(value))
-	_, e := p.writer.Write(v)
+	_, e := p.trans.Write(v)
 	return NewTProtocolException(e)
 }
 
-func (p *TBinaryProtocol) WriteI64(value int64) error {
+func (p *TBinaryProtocol) WriteI64(ctx context.Context, value int64) error {
 	v := p.buffer[0:8]
 	binary.BigEndian.PutUint64(v, uint64(value))
-	_, err := p.writer.Write(v)
+	_, err := p.trans.Write(v)
 	return NewTProtocolException(err)
 }
 
-func (p *TBinaryProtocol) WriteDouble(value float64) error {
-	return p.WriteI64(int64(math.Float64bits(value)))
+func (p *TBinaryProtocol) WriteDouble(ctx context.Context, value float64) error {
+	return p.WriteI64(ctx, int64(math.Float64bits(value)))
 }
 
-func (p *TBinaryProtocol) WriteString(value string) error {
-	e := p.WriteI32(int32(len(value)))
+func (p *TBinaryProtocol) WriteString(ctx context.Context, value string) error {
+	e := p.WriteI32(ctx, int32(len(value)))
 	if e != nil {
 		return e
 	}
@@ -222,12 +253,12 @@
 	return NewTProtocolException(err)
 }
 
-func (p *TBinaryProtocol) WriteBinary(value []byte) error {
-	e := p.WriteI32(int32(len(value)))
+func (p *TBinaryProtocol) WriteBinary(ctx context.Context, value []byte) error {
+	e := p.WriteI32(ctx, int32(len(value)))
 	if e != nil {
 		return e
 	}
-	_, err := p.writer.Write(value)
+	_, err := p.trans.Write(value)
 	return NewTProtocolException(err)
 }
 
@@ -235,8 +266,8 @@
  * Reading methods
  */
 
-func (p *TBinaryProtocol) ReadMessageBegin() (name string, typeId TMessageType, seqId int32, err error) {
-	size, e := p.ReadI32()
+func (p *TBinaryProtocol) ReadMessageBegin(ctx context.Context) (name string, typeId TMessageType, seqId int32, err error) {
+	size, e := p.ReadI32(ctx)
 	if e != nil {
 		return "", typeId, 0, NewTProtocolException(e)
 	}
@@ -246,79 +277,79 @@
 		if version != VERSION_1 {
 			return name, typeId, seqId, NewTProtocolExceptionWithType(BAD_VERSION, fmt.Errorf("Bad version in ReadMessageBegin"))
 		}
-		name, e = p.ReadString()
+		name, e = p.ReadString(ctx)
 		if e != nil {
 			return name, typeId, seqId, NewTProtocolException(e)
 		}
-		seqId, e = p.ReadI32()
+		seqId, e = p.ReadI32(ctx)
 		if e != nil {
 			return name, typeId, seqId, NewTProtocolException(e)
 		}
 		return name, typeId, seqId, nil
 	}
-	if p.strictRead {
+	if p.cfg.GetTBinaryStrictRead() {
 		return name, typeId, seqId, NewTProtocolExceptionWithType(BAD_VERSION, fmt.Errorf("Missing version in ReadMessageBegin"))
 	}
 	name, e2 := p.readStringBody(size)
 	if e2 != nil {
 		return name, typeId, seqId, e2
 	}
-	b, e3 := p.ReadByte()
+	b, e3 := p.ReadByte(ctx)
 	if e3 != nil {
 		return name, typeId, seqId, e3
 	}
 	typeId = TMessageType(b)
-	seqId, e4 := p.ReadI32()
+	seqId, e4 := p.ReadI32(ctx)
 	if e4 != nil {
 		return name, typeId, seqId, e4
 	}
 	return name, typeId, seqId, nil
 }
 
-func (p *TBinaryProtocol) ReadMessageEnd() error {
+func (p *TBinaryProtocol) ReadMessageEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) ReadStructBegin() (name string, err error) {
+func (p *TBinaryProtocol) ReadStructBegin(ctx context.Context) (name string, err error) {
 	return
 }
 
-func (p *TBinaryProtocol) ReadStructEnd() error {
+func (p *TBinaryProtocol) ReadStructEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) ReadFieldBegin() (name string, typeId TType, seqId int16, err error) {
-	t, err := p.ReadByte()
+func (p *TBinaryProtocol) ReadFieldBegin(ctx context.Context) (name string, typeId TType, seqId int16, err error) {
+	t, err := p.ReadByte(ctx)
 	typeId = TType(t)
 	if err != nil {
 		return name, typeId, seqId, err
 	}
 	if t != STOP {
-		seqId, err = p.ReadI16()
+		seqId, err = p.ReadI16(ctx)
 	}
 	return name, typeId, seqId, err
 }
 
-func (p *TBinaryProtocol) ReadFieldEnd() error {
+func (p *TBinaryProtocol) ReadFieldEnd(ctx context.Context) error {
 	return nil
 }
 
 var invalidDataLength = NewTProtocolExceptionWithType(INVALID_DATA, errors.New("Invalid data length"))
 
-func (p *TBinaryProtocol) ReadMapBegin() (kType, vType TType, size int, err error) {
-	k, e := p.ReadByte()
+func (p *TBinaryProtocol) ReadMapBegin(ctx context.Context) (kType, vType TType, size int, err error) {
+	k, e := p.ReadByte(ctx)
 	if e != nil {
 		err = NewTProtocolException(e)
 		return
 	}
 	kType = TType(k)
-	v, e := p.ReadByte()
+	v, e := p.ReadByte(ctx)
 	if e != nil {
 		err = NewTProtocolException(e)
 		return
 	}
 	vType = TType(v)
-	size32, e := p.ReadI32()
+	size32, e := p.ReadI32(ctx)
 	if e != nil {
 		err = NewTProtocolException(e)
 		return
@@ -331,18 +362,18 @@
 	return kType, vType, size, nil
 }
 
-func (p *TBinaryProtocol) ReadMapEnd() error {
+func (p *TBinaryProtocol) ReadMapEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) ReadListBegin() (elemType TType, size int, err error) {
-	b, e := p.ReadByte()
+func (p *TBinaryProtocol) ReadListBegin(ctx context.Context) (elemType TType, size int, err error) {
+	b, e := p.ReadByte(ctx)
 	if e != nil {
 		err = NewTProtocolException(e)
 		return
 	}
 	elemType = TType(b)
-	size32, e := p.ReadI32()
+	size32, e := p.ReadI32(ctx)
 	if e != nil {
 		err = NewTProtocolException(e)
 		return
@@ -356,18 +387,18 @@
 	return
 }
 
-func (p *TBinaryProtocol) ReadListEnd() error {
+func (p *TBinaryProtocol) ReadListEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) ReadSetBegin() (elemType TType, size int, err error) {
-	b, e := p.ReadByte()
+func (p *TBinaryProtocol) ReadSetBegin(ctx context.Context) (elemType TType, size int, err error) {
+	b, e := p.ReadByte(ctx)
 	if e != nil {
 		err = NewTProtocolException(e)
 		return
 	}
 	elemType = TType(b)
-	size32, e := p.ReadI32()
+	size32, e := p.ReadI32(ctx)
 	if e != nil {
 		err = NewTProtocolException(e)
 		return
@@ -380,12 +411,12 @@
 	return elemType, size, nil
 }
 
-func (p *TBinaryProtocol) ReadSetEnd() error {
+func (p *TBinaryProtocol) ReadSetEnd(ctx context.Context) error {
 	return nil
 }
 
-func (p *TBinaryProtocol) ReadBool() (bool, error) {
-	b, e := p.ReadByte()
+func (p *TBinaryProtocol) ReadBool(ctx context.Context) (bool, error) {
+	b, e := p.ReadByte(ctx)
 	v := true
 	if b != 1 {
 		v = false
@@ -393,122 +424,132 @@
 	return v, e
 }
 
-func (p *TBinaryProtocol) ReadByte() (int8, error) {
+func (p *TBinaryProtocol) ReadByte(ctx context.Context) (int8, error) {
 	v, err := p.trans.ReadByte()
 	return int8(v), err
 }
 
-func (p *TBinaryProtocol) ReadI16() (value int16, err error) {
+func (p *TBinaryProtocol) ReadI16(ctx context.Context) (value int16, err error) {
 	buf := p.buffer[0:2]
-	err = p.readAll(buf)
+	err = p.readAll(ctx, buf)
 	value = int16(binary.BigEndian.Uint16(buf))
 	return value, err
 }
 
-func (p *TBinaryProtocol) ReadI32() (value int32, err error) {
+func (p *TBinaryProtocol) ReadI32(ctx context.Context) (value int32, err error) {
 	buf := p.buffer[0:4]
-	err = p.readAll(buf)
+	err = p.readAll(ctx, buf)
 	value = int32(binary.BigEndian.Uint32(buf))
 	return value, err
 }
 
-func (p *TBinaryProtocol) ReadI64() (value int64, err error) {
+func (p *TBinaryProtocol) ReadI64(ctx context.Context) (value int64, err error) {
 	buf := p.buffer[0:8]
-	err = p.readAll(buf)
+	err = p.readAll(ctx, buf)
 	value = int64(binary.BigEndian.Uint64(buf))
 	return value, err
 }
 
-func (p *TBinaryProtocol) ReadDouble() (value float64, err error) {
+func (p *TBinaryProtocol) ReadDouble(ctx context.Context) (value float64, err error) {
 	buf := p.buffer[0:8]
-	err = p.readAll(buf)
+	err = p.readAll(ctx, buf)
 	value = math.Float64frombits(binary.BigEndian.Uint64(buf))
 	return value, err
 }
 
-func (p *TBinaryProtocol) ReadString() (value string, err error) {
-	size, e := p.ReadI32()
+func (p *TBinaryProtocol) ReadString(ctx context.Context) (value string, err error) {
+	size, e := p.ReadI32(ctx)
 	if e != nil {
 		return "", e
 	}
+	err = checkSizeForProtocol(size, p.cfg)
+	if err != nil {
+		return
+	}
 	if size < 0 {
 		err = invalidDataLength
 		return
 	}
+	if size == 0 {
+		return "", nil
+	}
+	if size < int32(len(p.buffer)) {
+		// Avoid allocation on small reads
+		buf := p.buffer[:size]
+		read, e := io.ReadFull(p.trans, buf)
+		return string(buf[:read]), NewTProtocolException(e)
+	}
 
 	return p.readStringBody(size)
 }
 
-func (p *TBinaryProtocol) ReadBinary() ([]byte, error) {
-	size, e := p.ReadI32()
+func (p *TBinaryProtocol) ReadBinary(ctx context.Context) ([]byte, error) {
+	size, e := p.ReadI32(ctx)
 	if e != nil {
 		return nil, e
 	}
-	if size < 0 {
-		return nil, invalidDataLength
-	}
-	if uint64(size) > p.trans.RemainingBytes() {
-		return nil, invalidDataLength
+	if err := checkSizeForProtocol(size, p.cfg); err != nil {
+		return nil, err
 	}
 
-	isize := int(size)
-	buf := make([]byte, isize)
-	_, err := io.ReadFull(p.trans, buf)
+	buf, err := safeReadBytes(size, p.trans)
 	return buf, NewTProtocolException(err)
 }
 
-func (p *TBinaryProtocol) Flush() (err error) {
-	return NewTProtocolException(p.trans.Flush())
+func (p *TBinaryProtocol) Flush(ctx context.Context) (err error) {
+	return NewTProtocolException(p.trans.Flush(ctx))
 }
 
-func (p *TBinaryProtocol) Skip(fieldType TType) (err error) {
-	return SkipDefaultDepth(p, fieldType)
+func (p *TBinaryProtocol) Skip(ctx context.Context, fieldType TType) (err error) {
+	return SkipDefaultDepth(ctx, p, fieldType)
 }
 
 func (p *TBinaryProtocol) Transport() TTransport {
 	return p.origTransport
 }
 
-func (p *TBinaryProtocol) readAll(buf []byte) error {
-	_, err := io.ReadFull(p.reader, buf)
+func (p *TBinaryProtocol) readAll(ctx context.Context, buf []byte) (err error) {
+	var read int
+	_, deadlineSet := ctx.Deadline()
+	for {
+		read, err = io.ReadFull(p.trans, buf)
+		if deadlineSet && read == 0 && isTimeoutError(err) && ctx.Err() == nil {
+			// This is I/O timeout without anything read,
+			// and we still have time left, keep retrying.
+			continue
+		}
+		// For anything else, don't retry
+		break
+	}
 	return NewTProtocolException(err)
 }
 
-const readLimit = 32768
-
 func (p *TBinaryProtocol) readStringBody(size int32) (value string, err error) {
+	buf, err := safeReadBytes(size, p.trans)
+	return string(buf), NewTProtocolException(err)
+}
+
+func (p *TBinaryProtocol) SetTConfiguration(conf *TConfiguration) {
+	PropagateTConfiguration(p.trans, conf)
+	PropagateTConfiguration(p.origTransport, conf)
+	p.cfg = conf
+}
+
+var (
+	_ TConfigurationSetter = (*TBinaryProtocolFactory)(nil)
+	_ TConfigurationSetter = (*TBinaryProtocol)(nil)
+)
+
+// This function is shared between TBinaryProtocol and TCompactProtocol.
+//
+// It tries to read size bytes from trans, in a way that prevents large
+// allocations when size is insanely large (mostly caused by malformed message).
+func safeReadBytes(size int32, trans io.Reader) ([]byte, error) {
 	if size < 0 {
-		return "", nil
-	}
-	if uint64(size) > p.trans.RemainingBytes() {
-		return "", invalidDataLength
+		return nil, nil
 	}
 
-	var (
-		buf bytes.Buffer
-		e   error
-		b   []byte
-	)
-
-	switch {
-	case int(size) <= len(p.buffer):
-		b = p.buffer[:size] // avoids allocation for small reads
-	case int(size) < readLimit:
-		b = make([]byte, size)
-	default:
-		b = make([]byte, readLimit)
-	}
-
-	for size > 0 {
-		_, e = io.ReadFull(p.trans, b)
-		buf.Write(b)
-		if e != nil {
-			break
-		}
-		size -= readLimit
-		if size < readLimit && size > 0 {
-			b = b[:size]
-		}
-	}
-	return buf.String(), NewTProtocolException(e)
+	buf := new(bytes.Buffer)
+	_, err := io.CopyN(buf, trans, int64(size))
+	return buf.Bytes(), err
 }