Repository: go-kratos/kratos
Pull request #1770 (2022-01-16 11:23:49)

feat:add stream interceptor use ctx encapsulation
<!--
🎉 Thanks for sending a pull request to Kratos! Here are some tips for you:

1. If this is your first time contributing to Kratos, please read our contribution guide: https://go-kratos.dev/en/docs/community/contribution/
2. Ensure you have added or ran the appropriate tests and lint for your PR, please use `make lint` and `make test` before filing your PR, use `make clean` to tidy your go mod.
3. If the PR is unfinished, you may need mark it as a WIP(Work In Progress) PR or draft PR
4. Please use a conventional commits format title: `<type>[optional scope]: <description>`
    
    Some suggestion on <type>:

    fix: A bug fix
    feat: A new feature
    test: Adding missing tests or correcting existing tests
    refactor: A code change that neither fixes a bug nor adds a feature
    break: Changes has break change

    docs: Documentation only changes
    deps: Changes external dependencies
    style: Changes that do not affect the meaning of the code (white-space, formatting, etc)
    chore Daily work, examples, etc.
    ci: Changes to our CI configuration files and scripts
-->

#### Description (what this PR does / why we need it):
<!--
* The description should include the motivation for this PR or contrast this with previous behavior
-->


#### Which issue(s) this PR fixes (resolves / be part of):
<!--
* Automatically closes linked issue when PR is merged.
* If you PR is not fully resolved issue, please use `part of #<issue number>` instead.

Usage: `fixes/resolves #<issue number>`, or `fixes/resolves (paste link of issue)`.
-->
fixes #

#### Other special notes for reviewer:
<!--
* Somethings that need extra attention for reviewer
* Some additional notes, TODO list, etc.
-->


codecov-commenter: # [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1770?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) Report
> Merging [#1770](https://codecov.io/gh/go-kratos/kratos/pull/1770?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (8d60b5d) into [main](https://codecov.io/gh/go-kratos/kratos/commit/83fad75c0f642092619a8d4d4927f88bf200a79e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) (83fad75) will **decrease** coverage by `0.51%`.
> The diff coverage is `59.32%`.

[![Impacted file tree graph](https://codecov.io/gh/go-kratos/kratos/pull/1770/graphs/tree.svg?width=650&height=150&src=pr&token=XwnmBTz78x&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)](https://codecov.io/gh/go-kratos/kratos/pull/1770?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)

```diff
@@            Coverage Diff             @@
##             main    #1770      +/-   ##
==========================================
- Coverage   81.30%   80.79%   -0.52%     
==========================================
  Files          79       80       +1     
  Lines        3525     3556      +31     
==========================================
+ Hits         2866     2873       +7     
- Misses        455      477      +22     
- Partials      204      206       +2     
```


| [Impacted Files](https://codecov.io/gh/go-kratos/kratos/pull/1770?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos) | Coverage Δ | |
|---|---|---|
| [transport/grpc/interceptor.go](https://codecov.io/gh/go-kratos/kratos/pull/1770/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-dHJhbnNwb3J0L2dycGMvaW50ZXJjZXB0b3IuZ28=) | `52.17% <52.17%> (ø)` | |
| [transport/grpc/server.go](https://codecov.io/gh/go-kratos/kratos/pull/1770/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos#diff-dHJhbnNwb3J0L2dycGMvc2VydmVyLmdv) | `90.00% <84.61%> (-3.34%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1770?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/go-kratos/kratos/pull/1770?src=pr&el=footer&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Last update [83fad75...8d60b5d](https://codecov.io/gh/go-kratos/kratos/pull/1770?src=pr&el=lastupdated&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=go-kratos).


Patch:
diff --git a/cmd/kratos/internal/upgrade/upgrade.go b/cmd/kratos/internal/upgrade/upgrade.go
index 5c5f8818..a22eeca4 100644
--- a/cmd/kratos/internal/upgrade/upgrade.go
+++ b/cmd/kratos/internal/upgrade/upgrade.go
@@ -19,13 +19,12 @@ var CmdUpgrade = &cobra.Command{
 // Run upgrade the kratos tools.
 func Run(cmd *cobra.Command, args []string) {
 	err := base.GoInstall(
-		"github.com/go-kratos/kratos/cmd/kratos/v2",
-		"github.com/go-kratos/kratos/cmd/protoc-gen-go-http/v2",
-		"github.com/go-kratos/kratos/cmd/protoc-gen-go-errors/v2",
-		"google.golang.org/protobuf/cmd/protoc-gen-go",
-		"google.golang.org/grpc/cmd/protoc-gen-go-grpc",
-		"github.com/envoyproxy/protoc-gen-validate",
-		"github.com/google/gnostic",
+		"github.com/go-kratos/kratos/cmd/kratos/v2@latest",
+		"github.com/go-kratos/kratos/cmd/protoc-gen-go-http/v2@latest",
+		"github.com/go-kratos/kratos/cmd/protoc-gen-go-errors/v2@latest",
+		"google.golang.org/protobuf/cmd/protoc-gen-go@latest",
+		"google.golang.org/grpc/cmd/protoc-gen-go-grpc@latest",
+		"github.com/envoyproxy/protoc-gen-validate@latest",
 		"github.com/google/gnostic/cmd/protoc-gen-openapi@v0.6.2",
 	)
 	if err != nil {
diff --git a/transport/grpc/interceptor.go b/transport/grpc/interceptor.go
new file mode 100644
index 00000000..3f3827f6
--- /dev/null
+++ b/transport/grpc/interceptor.go
@@ -0,0 +1,83 @@
+package grpc
+
+import (
+	"context"
+
+	ic "github.com/go-kratos/kratos/v2/internal/context"
+	"github.com/go-kratos/kratos/v2/middleware"
+	"github.com/go-kratos/kratos/v2/transport"
+	"google.golang.org/grpc"
+	grpcmd "google.golang.org/grpc/metadata"
+)
+
+// unaryServerInterceptor is a gRPC unary server interceptor
+func (s *Server) unaryServerInterceptor() grpc.UnaryServerInterceptor {
+	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
+		ctx, cancel := ic.Merge(ctx, s.baseCtx)
+		defer cancel()
+		md, _ := grpcmd.FromIncomingContext(ctx)
+		replyHeader := grpcmd.MD{}
+		ctx = transport.NewServerContext(ctx, &Transport{
+			endpoint:    s.endpoint.String(),
+			operation:   info.FullMethod,
+			reqHeader:   headerCarrier(md),
+			replyHeader: headerCarrier(replyHeader),
+		})
+		if s.timeout > 0 {
+			ctx, cancel = context.WithTimeout(ctx, s.timeout)
+			defer cancel()
+		}
+		h := func(ctx context.Context, req interface{}) (interface{}, error) {
+			return handler(ctx, req)
+		}
+		if len(s.middleware) > 0 {
+			h = middleware.Chain(s.middleware...)(h)
+		}
+		reply, err := h(ctx, req)
+		if len(replyHeader) > 0 {
+			_ = grpc.SetHeader(ctx, replyHeader)
+		}
+		return reply, err
+	}
+}
+
+// wrappedStream is rewrite grpc stream's context
+type wrappedStream struct {
+	grpc.ServerStream
+	ctx context.Context
+}
+
+func NewWrappedStream(ctx context.Context, stream grpc.ServerStream) grpc.ServerStream {
+	return &wrappedStream{
+		ServerStream: stream,
+		ctx:          ctx,
+	}
+}
+
+func (w *wrappedStream) Context() context.Context {
+	return w.ctx
+}
+
+// streamServerInterceptor is a gRPC stream server interceptor
+func (s *Server) streamServerInterceptor() grpc.StreamServerInterceptor {
+	return func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+		ctx, cancel := ic.Merge(ss.Context(), s.baseCtx)
+		defer cancel()
+		md, _ := grpcmd.FromIncomingContext(ctx)
+		replyHeader := grpcmd.MD{}
+		ctx = transport.NewServerContext(ctx, &Transport{
+			endpoint:    s.endpoint.String(),
+			operation:   info.FullMethod,
+			reqHeader:   headerCarrier(md),
+			replyHeader: headerCarrier(replyHeader),
+		})
+
+		ws := NewWrappedStream(ctx, ss)
+
+		err := handler(srv, ws)
+		if len(replyHeader) > 0 {
+			_ = grpc.SetHeader(ctx, replyHeader)
+		}
+		return err
+	}
+}
diff --git a/transport/grpc/server.go b/transport/grpc/server.go
index c7ebbaa1..caf4e694 100644
--- a/transport/grpc/server.go
+++ b/transport/grpc/server.go
@@ -10,7 +10,7 @@ import (
 	"github.com/go-kratos/kratos/v2/internal/endpoint"
 
 	apimd "github.com/go-kratos/kratos/v2/api/metadata"
-	ic "github.com/go-kratos/kratos/v2/internal/context"
+
 	"github.com/go-kratos/kratos/v2/internal/host"
 	"github.com/go-kratos/kratos/v2/log"
 	"github.com/go-kratos/kratos/v2/middleware"
@@ -20,7 +20,7 @@ import (
 	"google.golang.org/grpc/credentials"
 	"google.golang.org/grpc/health"
 	"google.golang.org/grpc/health/grpc_health_v1"
-	grpcmd "google.golang.org/grpc/metadata"
+
 	"google.golang.org/grpc/reflection"
 )
 
@@ -84,7 +84,14 @@ func Listener(lis net.Listener) ServerOption {
 // UnaryInterceptor returns a ServerOption that sets the UnaryServerInterceptor for the server.
 func UnaryInterceptor(in ...grpc.UnaryServerInterceptor) ServerOption {
 	return func(s *Server) {
-		s.ints = in
+		s.unaryInts = in
+	}
+}
+
+// StreamInterceptor returns a ServerOption that sets the StreamServerInterceptor for the server.
+func StreamInterceptor(in ...grpc.StreamServerInterceptor) ServerOption {
+	return func(s *Server) {
+		s.streamInts = in
 	}
 }
 
@@ -108,7 +115,8 @@ type Server struct {
 	timeout    time.Duration
 	log        *log.Helper
 	middleware []middleware.Middleware
-	ints       []grpc.UnaryServerInterceptor
+	unaryInts  []grpc.UnaryServerInterceptor
+	streamInts []grpc.StreamServerInterceptor
 	grpcOpts   []grpc.ServerOption
 	health     *health.Server
 	metadata   *apimd.Server
@@ -127,14 +135,21 @@ func NewServer(opts ...ServerOption) *Server {
 	for _, o := range opts {
 		o(srv)
 	}
-	ints := []grpc.UnaryServerInterceptor{
+	unaryInts := []grpc.UnaryServerInterceptor{
 		srv.unaryServerInterceptor(),
 	}
-	if len(srv.ints) > 0 {
-		ints = append(ints, srv.ints...)
+	streamInts := []grpc.StreamServerInterceptor{
+		srv.streamServerInterceptor(),
+	}
+	if len(srv.unaryInts) > 0 {
+		unaryInts = append(unaryInts, srv.unaryInts...)
+	}
+	if len(srv.streamInts) > 0 {
+		streamInts = append(streamInts, srv.streamInts...)
 	}
 	grpcOpts := []grpc.ServerOption{
-		grpc.ChainUnaryInterceptor(ints...),
+		grpc.ChainUnaryInterceptor(unaryInts...),
+		grpc.ChainStreamInterceptor(streamInts...),
 	}
 	if srv.tlsConf != nil {
 		grpcOpts = append(grpcOpts, grpc.Creds(credentials.NewTLS(srv.tlsConf)))
@@ -198,33 +213,3 @@ func (s *Server) listenAndEndpoint() error {
 	s.endpoint = endpoint.NewEndpoint("grpc", addr, s.tlsConf != nil)
 	return nil
 }
-
-func (s *Server) unaryServerInterceptor() grpc.UnaryServerInterceptor {
-	return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
-		ctx, cancel := ic.Merge(ctx, s.baseCtx)
-		defer cancel()
-		md, _ := grpcmd.FromIncomingContext(ctx)
-		replyHeader := grpcmd.MD{}
-		ctx = transport.NewServerContext(ctx, &Transport{
-			endpoint:    s.endpoint.String(),
-			operation:   info.FullMethod,
-			reqHeader:   headerCarrier(md),
-			replyHeader: headerCarrier(replyHeader),
-		})
-		if s.timeout > 0 {
-			ctx, cancel = context.WithTimeout(ctx, s.timeout)
-			defer cancel()
-		}
-		h := func(ctx context.Context, req interface{}) (interface{}, error) {
-			return handler(ctx, req)
-		}
-		if len(s.middleware) > 0 {
-			h = middleware.Chain(s.middleware...)(h)
-		}
-		reply, err := h(ctx, req)
-		if len(replyHeader) > 0 {
-			_ = grpc.SetHeader(ctx, replyHeader)
-		}
-		return reply, err
-	}
-}


Test patch:
diff --git a/transport/grpc/server_test.go b/transport/grpc/server_test.go
index f8063913..aaa6aca1 100644
--- a/transport/grpc/server_test.go
+++ b/transport/grpc/server_test.go
@@ -207,8 +207,24 @@ func TestUnaryInterceptor(t *testing.T) {
 		},
 	}
 	UnaryInterceptor(v...)(o)
-	if !reflect.DeepEqual(v, o.ints) {
-		t.Errorf("expect %v, got %v", v, o.ints)
+	if !reflect.DeepEqual(v, o.unaryInts) {
+		t.Errorf("expect %v, got %v", v, o.unaryInts)
+	}
+}
+
+func TestStramInterceptor(t *testing.T) {
+	o := &Server{}
+	v := []grpc.StreamServerInterceptor{
+		func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+			return nil
+		},
+		func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
+			return nil
+		},
+	}
+	StreamInterceptor(v...)(o)
+	if !reflect.DeepEqual(v, o.streamInts) {
+		t.Errorf("expect %v, got %v", v, o.streamInts)
 	}
 }