unistack-org/micro
4
0
Fork 3
Code Issues 19 Pull Requests 2 Actions 2 Packages Projects Releases 262 Wiki Activity

cancel stream

Browse Source
This commit is contained in:
Asim Aslam 2020-01-19 22:53:56 +00:00
parent 3f3c1919f4
commit 9df19e826e
2 changed files with 8 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
client/grpc/grpc.go
View File

@ -212,7 +212,10 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client
grpcCallOptions = append(grpcCallOptions, opts...) grpcCallOptions = append(grpcCallOptions, opts...)
} }
st, err := cc.NewStream(ctx, desc, methodToGRPC(req.Service(), req.Endpoint()), grpcCallOptions...) // create a new cancelling context
newCtx, cancel := context.WithCancel(ctx)
st, err := cc.NewStream(newCtx, desc, methodToGRPC(req.Service(), req.Endpoint()), grpcCallOptions...)
if err != nil { if err != nil {
return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error creating stream: %v", err)) return nil, errors.InternalServerError("go.micro.client", fmt.Sprintf("Error creating stream: %v", err))
} }
@ -240,6 +243,7 @@ func (g *grpcClient) stream(ctx context.Context, node *registry.Node, req client
response: rsp, response: rsp,
stream: st, stream: st,
conn: cc, conn: cc,
cancel: cancel,
}, nil }, nil
} }

4
client/grpc/stream.go
View File

@ -19,6 +19,7 @@ type grpcStream struct {
request client.Request request client.Request
response client.Response response client.Response
context context.Context context context.Context
cancel func()
} }
func (g *grpcStream) Context() context.Context { func (g *grpcStream) Context() context.Context {
@ -79,7 +80,8 @@ func (g *grpcStream) Close() error {
if g.closed { if g.closed {
return nil return nil
} }
// cancel the context
g.cancel()
g.closed = true g.closed = true
g.stream.CloseSend() g.stream.CloseSend()
return g.conn.Close() return g.conn.Close()