Merge pull request #215 from uber/weiqing.addCtx

Update tchannel version and add a context option

.gitignore

@@ -7,3 +7,4 @@ lint.log
 /_venv/
 README.md.err
 tick-cluster.log
+.idea/

examples/keyvalue/gen-go/keyvalue/tchan-keyvalue.go

@@ -45,6 +45,10 @@ func (c *tchanKeyValueServiceClient) Get(ctx thrift.Context, key string) (string
 	}
 	success, err := c.client.Call(ctx, c.thriftService, "Get", &args, &resp)
 	if err == nil && !success {
+		switch {
+		default:
+			err = fmt.Errorf("received no result or unknown exception for Get")
+		}
 	}
 
 	return resp.GetSuccess(), err
@@ -57,6 +61,10 @@ func (c *tchanKeyValueServiceClient) GetAll(ctx thrift.Context, keys []string) (
 	}
 	success, err := c.client.Call(ctx, c.thriftService, "GetAll", &args, &resp)
 	if err == nil && !success {
+		switch {
+		default:
+			err = fmt.Errorf("received no result or unknown exception for GetAll")
+		}
 	}
 
 	return resp.GetSuccess(), err
@@ -70,6 +78,10 @@ func (c *tchanKeyValueServiceClient) Set(ctx thrift.Context, key string, value s
 	}
 	success, err := c.client.Call(ctx, c.thriftService, "Set", &args, &resp)
 	if err == nil && !success {
+		switch {
+		default:
+			err = fmt.Errorf("received no result or unknown exception for Set")
+		}
 	}
 
 	return err

examples/ping-thrift-gen/gen-go/ping/tchan-ping.go

@@ -43,6 +43,10 @@ func (c *tchanPingPongServiceClient) Ping(ctx thrift.Context, request *Ping) (*P
 	}
 	success, err := c.client.Call(ctx, c.thriftService, "Ping", &args, &resp)
 	if err == nil && !success {
+		switch {
+		default:
+			err = fmt.Errorf("received no result or unknown exception for Ping")
+		}
 	}
 
 	return resp.GetSuccess(), err

examples/role-labels/gen-go/role/tchan-role.go

@@ -44,6 +44,10 @@ func (c *tchanRoleServiceClient) GetMembers(ctx thrift.Context, role string) ([]
 	}
 	success, err := c.client.Call(ctx, c.thriftService, "GetMembers", &args, &resp)
 	if err == nil && !success {
+		switch {
+		default:
+			err = fmt.Errorf("received no result or unknown exception for GetMembers")
+		}
 	}
 
 	return resp.GetSuccess(), err
@@ -56,6 +60,10 @@ func (c *tchanRoleServiceClient) SetRole(ctx thrift.Context, role string) error
 	}
 	success, err := c.client.Call(ctx, c.thriftService, "SetRole", &args, &resp)
 	if err == nil && !success {
+		switch {
+		default:
+			err = fmt.Errorf("received no result or unknown exception for SetRole")
+		}
 	}
 
 	return err

forward/forwarder.go

@@ -47,6 +47,7 @@ type Sender interface {
 
 // Options for the creation of a forwarder
 type Options struct {
+	Ctx            thrift.Context
 	MaxRetries     int
 	RerouteRetries bool
 	RetrySchedule  []time.Duration
@@ -80,6 +81,7 @@ func (f *Forwarder) mergeDefaultOptions(opts *Options) *Options {
 		merged.RetrySchedule = def.RetrySchedule
 	}
 	merged.Headers = opts.Headers
+	merged.Ctx = opts.Ctx
 
 	return &merged
 }

forward/forwarder_test.go

@@ -41,6 +41,8 @@ import (
 	"golang.org/x/net/context"
 )
 
+type ContextKey string
+
 type ForwarderTestSuite struct {
 	suite.Suite
 	sender    *MockSender
@@ -77,6 +79,16 @@ func (s *ForwarderTestSuite) registerPong(address string, channel *tchannel.Chan
 
 	thriftHandler := &pingpong.MockTChanPingPong{}
 
+	// successful request with context
+	thriftHandler.On("Ping", mock.MatchedBy(
+		func(c thrift.Context) bool {
+			return true
+		}), &pingpong.Ping{
+		Key: "ctxTest",
+	}).Return(&pingpong.Pong{
+		Source: address,
+	}, nil)
+
 	// successful request
 	thriftHandler.On("Ping", mock.Anything, &pingpong.Ping{
 		Key: "success",
@@ -194,6 +206,37 @@ func (s *ForwarderTestSuite) TestForwardThrift() {
 	s.Equal("correct pinging host", response.Success.Source)
 }
 
+func (s *ForwarderTestSuite) TestForwardThriftWithCtxOption() {
+	dest, err := s.sender.Lookup("reachable")
+	s.NoError(err)
+
+	request := &pingpong.PingPongPingArgs{
+		Request: &pingpong.Ping{
+			Key: "ctxTest",
+		},
+	}
+
+	bytes1, err := SerializeThrift(request)
+	s.NoError(err, "expected ping to be serialized")
+
+
+	k := ContextKey("key")
+	ctx := thrift.Wrap(context.WithValue(context.Background(), k, "val"))
+
+	res, err := s.forwarder.ForwardRequest(bytes1, dest, "test", "PingPong::Ping", []string{"reachable"},
+		tchannel.Thrift, &Options{
+			Ctx: ctx,
+		})
+	s.NoError(err, "expected request to be forwarded")
+
+	var response pingpong.PingPongPingResult
+
+	err = DeserializeThrift(res, &response)
+	s.NoError(err)
+
+	s.Equal("correct pinging host", response.Success.Source)
+}
+
 func (s *ForwarderTestSuite) TestForwardThriftErrorResponse() {
 	dest, err := s.sender.Lookup("reachable")
 	s.NoError(err)

forward/request_sender.go

@@ -33,6 +33,7 @@ import (
 	"github.com/uber/ringpop-go/shared"
 	"github.com/uber/tchannel-go"
 	"github.com/uber/tchannel-go/raw"
+	"github.com/uber/tchannel-go/thrift"
 )
 
 // errDestinationsDiverged is an error that is returned from AttemptRetry
@@ -60,6 +61,7 @@ type requestSender struct {
 	rerouteRetries      bool
 
 	headers []byte
+	ctx     thrift.Context
 
 	startTime, retryStartTime time.Time
 
@@ -90,12 +92,21 @@ func newRequestSender(sender Sender, emitter events.EventEmitter, channel shared
 		retrySchedule:  opts.RetrySchedule,
 		rerouteRetries: opts.RerouteRetries,
 		headers:        opts.Headers,
+		ctx:            opts.Ctx,
 		logger:         logger,
 	}
 }
 
 func (s *requestSender) Send() (res []byte, err error) {
-	ctx, cancel := shared.NewTChannelContext(s.timeout)
+	var cancel context.CancelFunc
+	var ctx thrift.Context
+	if s.ctx == nil {
+		ctx, cancel = shared.NewTChannelContext(s.timeout)
+	} else {
+		var c context.Context
+		c, cancel = context.WithTimeout(s.ctx, s.timeout)
+		ctx = tchannel.Wrap(c)
+	}
 	defer cancel()
 
 	var forwardError, applicationError error

glide.yaml

@@ -14,6 +14,7 @@ import:
   - require
 - package: github.com/uber-common/bark
 - package: github.com/uber/tchannel-go
+  version: 1.8.0
   subpackages:
   - json
   - raw