From f45cdba9ba0491e157b3e813f057771245ac94a4 Mon Sep 17 00:00:00 2001 From: ben-toogood Date: Tue, 2 Jun 2020 17:56:26 +0100 Subject: [PATCH] Apply wrappers to gRPC streams (#1675) * Add wrappers to grpc streams * Fix typo --- client/grpc/grpc.go | 50 ++++++++++++++++++++++++++++++--------------- 1 file changed, 33 insertions(+), 17 deletions(-) diff --git a/client/grpc/grpc.go b/client/grpc/grpc.go index 1932f9ce..7571e5bf 100644 --- a/client/grpc/grpc.go +++ b/client/grpc/grpc.go @@ -6,6 +6,7 @@ import ( "crypto/tls" "fmt" "net" + "reflect" "strings" "sync/atomic" "time" @@ -173,7 +174,7 @@ func (g *grpcClient) call(ctx context.Context, node *registry.Node, req client.R return grr } -func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client.Request, opts client.CallOptions) (client.Stream, error) { +func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client.Request, rsp interface{}, opts client.CallOptions) error { var header map[string]string address := node.Address @@ -199,7 +200,7 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client cf, err := g.newGRPCCodec(req.ContentType()) if err != nil { - return nil, errors.InternalServerError("go.micro.client", err.Error()) + return errors.InternalServerError("go.micro.client", err.Error()) } var dialCtx context.Context @@ -224,7 +225,7 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client cc, err := grpc.DialContext(dialCtx, address, grpcDialOptions...) if err != nil { - return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err)) + return errors.InternalServerError("go.micro.client", fmt.Sprintf("Error sending request: %v", err)) } desc := &grpc.StreamDesc{ @@ -252,7 +253,7 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client // close the connection cc.Close() // now return the error - return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error creating stream: %v", err)) + return errors.InternalServerError("go.micro.client", fmt.Sprintf("Error creating stream: %v", err)) } codec := &grpcCodec{ @@ -265,21 +266,25 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client r.codec = codec } - rsp := &response{ - conn: cc, + // setup the stream response + stream := &grpcStream{ + context: ctx, + request: req, + response: &response{ + conn: cc, + stream: st, + codec: cf, + gcodec: codec, + }, stream: st, - codec: cf, - gcodec: codec, + conn: cc, + cancel: cancel, } - return &grpcStream{ - context: ctx, - request: req, - response: rsp, - stream: st, - conn: cc, - cancel: cancel, - }, nil + // set the stream as the response + val := reflect.ValueOf(rsp).Elem() + val.Set(reflect.ValueOf(stream).Elem()) + return nil } func (g *grpcClient) poolMaxStreams() int { @@ -506,6 +511,14 @@ func (g *grpcClient) Stream(ctx context.Context, req client.Request, opts ...cli default: } + // make a copy of stream + gstream := g.stream + + // wrap the call in reverse + for i := len(callOpts.CallWrappers); i > 0; i-- { + gstream = callOpts.CallWrappers[i-1](gstream) + } + call := func(i int) (client.Stream, error) { // call backoff first. Someone may want an initial start delay t, err := callOpts.Backoff(ctx, req, i) @@ -527,7 +540,10 @@ func (g *grpcClient) Stream(ctx context.Context, req client.Request, opts ...cli return nil, errors.InternalServerError("go.micro.client", "error selecting %s node: %s", service, err.Error()) } - stream, err := g.stream(ctx, node, req, callOpts) + // make the call + stream := &grpcStream{} + err = g.stream(ctx, node, req, stream, callOpts) + g.opts.Selector.Mark(service, node, err) return stream, err }