feat: simplify client-side Stream using

Author: DMwangnima

client/callopt/streamcall/call_options.go

@@ -17,10 +17,12 @@
 package streamcall
 
 import (
+	"fmt"
 	"strings"
 	"time"
 
 	"github.com/cloudwego/kitex/client/callopt"
+	"github.com/cloudwego/kitex/pkg/streaming"
 )
 
 // These options are directly translated from callopt.Option(s). If you can't find the option with the
@@ -79,3 +81,13 @@ func WithStreamTimeout(d time.Duration) Option {
 		o.StreamOptions.StreamTimeout = d
 	}}
 }
+
+// WithCallbackConfig adds and calls FinishCallback when client-side stream
+// has finished.
+func WithCallbackConfig(cfg *streaming.FinishCallback) Option {
+	return Option{f: func(o *callopt.CallOptions, di *strings.Builder) {
+		fmt.Fprintf(di, "WithCallbackConfig(%+v)", cfg)
+
+		o.StreamOptions.FinishCallback = cfg
+	}}
+}

client/client.go

@@ -833,6 +833,9 @@ func initRPCInfo(ctx context.Context, method string, opt *client.Options, svcInf
 	if sopt.StreamTimeout > 0 {
 		cfg.SetStreamTimeout(sopt.StreamTimeout)
 	}
+	if sopt.FinishCallback != nil {
+		cfg.SetStreamCallbackConfig(sopt.FinishCallback)
+	}
 
 	ctx = rpcinfo.NewCtxWithRPCInfo(ctx, ri)
 
@@ -850,6 +853,9 @@ func initRPCInfo(ctx context.Context, method string, opt *client.Options, svcInf
 		if callOpts.StreamOptions.StreamTimeout != 0 {
 			cfg.SetStreamTimeout(callOpts.StreamOptions.StreamTimeout)
 		}
+		if callOpts.StreamOptions.FinishCallback != nil {
+			cfg.SetStreamCallbackConfig(callOpts.StreamOptions.FinishCallback)
+		}
 	}
 
 	return ctx, ri, callOpts

client/option_stream.go

@@ -22,6 +22,7 @@ import (
 
 	"github.com/cloudwego/kitex/internal/client"
 	"github.com/cloudwego/kitex/pkg/endpoint/cep"
+	"github.com/cloudwego/kitex/pkg/streaming"
 	"github.com/cloudwego/kitex/pkg/utils"
 )
 
@@ -66,6 +67,16 @@ func WithStreamTimeout(d time.Duration) StreamOption {
 	}}
 }
 
+// WithStreamCallbackConfig adds and calls FinishCallback when client-side stream
+// has finished.
+func WithStreamCallbackConfig(cfg *streaming.FinishCallback) StreamOption {
+	return StreamOption{F: func(o *StreamOptions, di *utils.Slice) {
+		di.Push(fmt.Sprintf("WithStreamCallbackConfig(%+v)", cfg))
+
+		o.FinishCallback = cfg
+	}}
+}
+
 // WithStreamMiddleware add middleware for stream.
 func WithStreamMiddleware(mw cep.StreamMiddleware) StreamOption {
 	return StreamOption{F: func(o *StreamOptions, di *utils.Slice) {

client/stream.go

@@ -20,6 +20,7 @@ import (
 	"context"
 	"fmt"
 	"io"
+	"runtime/debug"
 	"sync/atomic"
 	"time"
 
@@ -29,6 +30,7 @@ import (
 	"github.com/cloudwego/kitex/pkg/endpoint"
 	"github.com/cloudwego/kitex/pkg/endpoint/cep"
 	"github.com/cloudwego/kitex/pkg/kerrors"
+	"github.com/cloudwego/kitex/pkg/klog"
 	"github.com/cloudwego/kitex/pkg/remote"
 	"github.com/cloudwego/kitex/pkg/remote/remotecli"
 	"github.com/cloudwego/kitex/pkg/remote/trans/nphttp2/codes"
@@ -262,9 +264,6 @@ func newStream(ctx context.Context, cancel context.CancelFunc, s streaming.Clien
 			st.grpcStream.st = st
 		}
 	}
-	if register, ok := s.(streaming.CloseCallbackRegister); ok {
-		register.RegisterCloseCallback(st.DoFinish)
-	}
 	return st
 }
 
@@ -363,18 +362,44 @@ func (s *stream) DoFinish(err error) {
 		// already called
 		return
 	}
+
 	// release stream timeout cancel
 	if s.cancelFunc != nil {
 		s.cancelFunc()
 	}
-	if !isRPCError(err) {
-		// only rpc errors are reported
-		err = nil
+
+	// reporting and release connection
+	reportErr := err
+	if err == nil || err == io.EOF {
+		reportErr = nil
+	}
+	// If the client-side callback returns a biz error or manually calls streaming.FinishStream/streaming.FinishClientStream with a biz error,
+	// it needs to be set.
+	if bizErr, bizOk := err.(kerrors.BizStatusErrorIface); bizOk {
+		if setter, ok := s.ri.Invocation().(rpcinfo.InvocationSetter); ok {
+			setter.SetBizStatusErr(bizErr)
+		}
 	}
 	if s.scm != nil {
-		s.scm.ReleaseConn(err, s.ri)
+		s.scm.ReleaseConn(reportErr, s.ri)
+	}
+	s.kc.opt.TracerCtl.DoFinish(s.ctx, s.ri, reportErr)
+
+	// processing callback with original err
+	stCfg := s.ri.Config().StreamCallbackConfig()
+	if err == nil || err == io.EOF || stCfg == nil {
+		return
+	}
+	defer func() {
+		if r := recover(); r != nil {
+			klog.CtxWarnf(s.ctx, "Panic happened during stream DoFinish. This may caused by injected stream Callback: error=%v, stack=%s", r, string(debug.Stack()))
+		}
+	}()
+	if s.isGRPC {
+		handleGRPC(s.ctx, s.ri, err, stCfg)
+	} else {
+		handleTTStream(s.ctx, s.ri, err, stCfg)
 	}
-	s.kc.opt.TracerCtl.DoFinish(s.ctx, s.ri, err)
 }
 
 func (s *stream) GetGRPCStream() streaming.Stream {
@@ -478,18 +503,6 @@ func (s *grpcStream) DoFinish(err error) {
 	s.st.DoFinish(err)
 }
 
-func isRPCError(err error) bool {
-	if err == nil {
-		return false
-	}
-	if err == io.EOF {
-		return false
-	}
-	_, isBizStatusError := err.(kerrors.BizStatusErrorIface)
-	// if a tracer needs to get the BizStatusError, it should read from rpcinfo.invocation.bizStatusErr
-	return !isBizStatusError
-}
-
 func callWithTimeout(tm time.Duration, call func() error, buildTmErr func(time.Duration) error, cancel func(error)) error {
 	if tm <= 0 {
 		return call()