Merge pull request cloudwego#170 from ppzqh/hotfix/disable_timeout_st…

…reaming hotfix: disable timeout for streaming
starswiss · Oct 14, 2021 · b804b15 · b804b15
2 parents fb93b6b + b820843
commit b804b15
Show file tree

Hide file tree

Showing 8 changed files with 41 additions and 3 deletions.
diff --git a/client/rpctimeout.go b/client/rpctimeout.go
@@ -81,10 +81,13 @@ func rpcTimeoutMW(mwCtx context.Context) endpoint.Middleware {
 	logger := mwCtx.Value(endpoint.CtxLoggerKey).(klog.FormatLogger)
 	return func(next endpoint.Endpoint) endpoint.Endpoint {
 		return func(ctx context.Context, request, response interface{}) error {
-			var err error
 			ri := rpcinfo.GetRPCInfo(ctx)
-			tm := ri.Config().RPCTimeout()
+			if ri.Config().InteractionMode() == rpcinfo.Streaming {
+				return next(ctx, request, response)
+			}
 
+			var err error
+			tm := ri.Config().RPCTimeout()
 			start := time.Now()
 			ctx, cancel := context.WithTimeout(ctx, tm+moreTimeout)
 			defer cancel()

diff --git a/client/stream.go b/client/stream.go
@@ -45,6 +45,9 @@ func (kc *kClient) Stream(ctx context.Context, method string, request, response
 	var ri rpcinfo.RPCInfo
 	ctx, ri = kc.initRPCInfo(ctx, method)
 
+	rpcinfo.AsMutableRPCConfig(ri.Config()).SetInteractionMode(rpcinfo.Streaming)
+	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
+
 	ctx = kc.opt.TracerCtl.DoStart(ctx, ri, kc.opt.Logger)
 	return kc.sEps(ctx, request, response)
 }

diff --git a/pkg/rpcinfo/convert_test.go b/pkg/rpcinfo/convert_test.go
@@ -128,12 +128,14 @@ func TestAsMutableRPCConfig(t *testing.T) {
 	test.Assert(t, mc.SetReadWriteTimeout(time.Hour*3) == nil)
 	test.Assert(t, mc.SetIOBufferSize(9999) == nil)
 	test.Assert(t, mc.SetTransportProtocol(transport.HTTP) == nil)
+	test.Assert(t, mc.SetInteractionMode(rpcinfo.Unary) == nil)
 
 	test.Assert(t, c1.RPCTimeout() == time.Hour)
 	test.Assert(t, c1.ConnectTimeout() == time.Hour*2)
 	test.Assert(t, c1.ReadWriteTimeout() == time.Hour*3)
 	test.Assert(t, c1.IOBufferSize() == 9999)
 	test.Assert(t, c1.TransportProtocol() == transport.HTTP)
+	test.Assert(t, c1.InteractionMode() == rpcinfo.Unary)
 
 	mc.LockConfig(rpcinfo.BitRPCTimeout)
 	test.Assert(t, mc.IsRPCTimeoutLocked())

diff --git a/pkg/rpcinfo/interface.go b/pkg/rpcinfo/interface.go
@@ -73,6 +73,7 @@ type RPCConfig interface {
 	Timeouts
 	IOBufferSize() int
 	TransportProtocol() transport.Protocol
+	InteractionMode() InteractionMode
 }
 
 // Invocation contains specific information about the call.

diff --git a/pkg/rpcinfo/mocks_test.go b/pkg/rpcinfo/mocks_test.go
@@ -34,6 +34,14 @@ type MockRPCConfig struct {
 	ReadWriteTimeoutFunc  func() (r time.Duration)
 	IOBufferSizeFunc      func() (r int)
 	TransportProtocolFunc func() transport.Protocol
+	InteractionModeFunc   func() (r rpcinfo.InteractionMode)
+}
+
+func (m *MockRPCConfig) InteractionMode() (r rpcinfo.InteractionMode) {
+	if m.InteractionModeFunc != nil {
+		return m.InteractionModeFunc()
+	}
+	return
 }
 
 // RPCTimeout implements the rpcinfo.RPCConfig interface.

diff --git a/pkg/rpcinfo/mutable.go b/pkg/rpcinfo/mutable.go
@@ -43,6 +43,7 @@ type MutableRPCConfig interface {
 	IsReadWriteTimeoutLocked() bool
 	SetIOBufferSize(sz int) error
 	SetTransportProtocol(tp transport.Protocol) error
+	SetInteractionMode(mode InteractionMode) error
 	LockConfig(bits int)
 	Clone() MutableRPCConfig
 	ImmutableView() RPCConfig

diff --git a/pkg/rpcinfo/rpcconfig.go b/pkg/rpcinfo/rpcconfig.go
@@ -36,6 +36,7 @@ var (
 	defaultConnectTimeout   = time.Millisecond * 50
 	defaultReadWriteTimeout = time.Second * 5
 	defaultBufferSize       = 4096
+	defaultInteractionMode  = Unary
 )
 
 // Mask bits.
@@ -46,6 +47,14 @@ const (
 	BitIOBufferSize
 )
 
+type InteractionMode int32
+
+const (
+	Unary     InteractionMode = 0
+	Oneway    InteractionMode = 1
+	Streaming InteractionMode = 2
+)
+
 // rpcConfig is a set of configurations used during RPC calls.
 type rpcConfig struct {
 	readOnlyMask      int
@@ -54,6 +63,7 @@ type rpcConfig struct {
 	readWriteTimeout  time.Duration
 	ioBufferSize      int
 	transportProtocol transport.Protocol
+	interactionMode   InteractionMode
 }
 
 func init() {
@@ -156,6 +166,15 @@ func (r *rpcConfig) SetTransportProtocol(tp transport.Protocol) error {
 	return nil
 }
 
+func (r *rpcConfig) SetInteractionMode(mode InteractionMode) error {
+	r.interactionMode = mode
+	return nil
+}
+
+func (r *rpcConfig) InteractionMode() InteractionMode {
+	return r.interactionMode
+}
+
 // Clone returns a copy of the current rpcConfig.
 func (r *rpcConfig) Clone() MutableRPCConfig {
 	r2 := rpcConfigPool.Get().(*rpcConfig)
@@ -185,5 +204,6 @@ func NewRPCConfig() RPCConfig {
 	r.connectTimeout = defaultConnectTimeout
 	r.readWriteTimeout = defaultReadWriteTimeout
 	r.ioBufferSize = defaultBufferSize
+	r.interactionMode = defaultInteractionMode
 	return r
 }
diff --git a/version.go b/version.go
@@ -19,5 +19,5 @@ package kitex
 // Name and Version info of this framework, used for statistics and debug
 const (
 	Name    = "Kitex"
-	Version = "v0.0.5"
+	Version = "v0.0.6"
 )