[FABG-954] Cancel stream context when event client closes (hyperledge…

…r#55) This update ensures that the GRPC stream context is cancelled when the Deliver client connection is closed. Signed-off-by: Bob Stasyszyn <[email protected]>
Massad · Feb 22, 2020 · ff3bdd7 · ff3bdd7
1 parent ae164c9
commit ff3bdd7
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 12 deletions.
diff --git a/pkg/fab/comm/streamconnection.go b/pkg/fab/comm/streamconnection.go
@@ -21,13 +21,14 @@ import (
 )
 
 // StreamProvider creates a GRPC stream
-type StreamProvider func(conn *grpc.ClientConn) (grpc.ClientStream, error)
+type StreamProvider func(conn *grpc.ClientConn) (grpc.ClientStream, func(), error)
 
 // StreamConnection manages the GRPC connection and client stream
 type StreamConnection struct {
 	*GRPCConnection
 	chConfig fab.ChannelCfg
 	stream   grpc.ClientStream
+	cancel   func()
 	lock     sync.Mutex
 }
 
@@ -38,7 +39,7 @@ func NewStreamConnection(ctx fabcontext.Client, chConfig fab.ChannelCfg, streamP
 		return nil, err
 	}
 
-	stream, err := streamProvider(conn.conn)
+	stream, cancel, err := streamProvider(conn.conn)
 	if err != nil {
 		conn.commManager.ReleaseConn(conn.conn)
 		return nil, errors.Wrapf(err, "could not create stream to %s", url)
@@ -70,6 +71,7 @@ func NewStreamConnection(ctx fabcontext.Client, chConfig fab.ChannelCfg, streamP
 		GRPCConnection: conn,
 		chConfig:       chConfig,
 		stream:         stream,
+		cancel:         cancel,
 	}, nil
 }
 
@@ -88,6 +90,9 @@ func (c *StreamConnection) Close() {
 	}
 
 	logger.Debug("Closing stream....")
+
+	c.cancel()
+
 	if err := c.stream.CloseSend(); err != nil {
 		logger.Warnf("error closing GRPC stream: %s", err)
 	}

diff --git a/pkg/fab/comm/streamconnection_test.go b/pkg/fab/comm/streamconnection_test.go
@@ -20,12 +20,14 @@ import (
 	"google.golang.org/grpc"
 )
 
-var testStream = func(grpcconn *grpc.ClientConn) (grpc.ClientStream, error) {
-	return pb.NewDeliverClient(grpcconn).Deliver(context.Background())
+var testStream = func(grpcconn *grpc.ClientConn) (grpc.ClientStream, func(), error) {
+	ctx, cancel := context.WithCancel(context.Background())
+	stream, err := pb.NewDeliverClient(grpcconn).Deliver(ctx)
+	return stream, cancel, err
 }
 
-var invalidStream = func(grpcconn *grpc.ClientConn) (grpc.ClientStream, error) {
-	return nil, errors.New("simulated error creating stream")
+var invalidStream = func(grpcconn *grpc.ClientConn) (grpc.ClientStream, func(), error) {
+	return nil, nil, errors.New("simulated error creating stream")
 }
 
 func TestStreamConnection(t *testing.T) {

diff --git a/pkg/fab/events/deliverclient/connection/connection.go b/pkg/fab/events/deliverclient/connection/connection.go
@@ -42,17 +42,21 @@ type DeliverConnection struct {
 }
 
 // StreamProvider creates a deliver stream
-type StreamProvider func(pb.DeliverClient) (deliverStream, error)
+type StreamProvider func(pb.DeliverClient) (stream deliverStream, cancel func(), err error)
 
 var (
 	// Deliver creates a Deliver stream
-	Deliver = func(client pb.DeliverClient) (deliverStream, error) {
-		return client.Deliver(context.Background())
+	Deliver = func(client pb.DeliverClient) (deliverStream, func(), error) {
+		ctx, cancel := context.WithCancel(context.Background())
+		stream, err := client.Deliver(ctx)
+		return stream, cancel, err
 	}
 
 	// DeliverFiltered creates a DeliverFiltered stream
-	DeliverFiltered = func(client pb.DeliverClient) (deliverStream, error) {
-		return client.DeliverFiltered(context.Background())
+	DeliverFiltered = func(client pb.DeliverClient) (deliverStream, func(), error) {
+		ctx, cancel := context.WithCancel(context.Background())
+		stream, err := client.DeliverFiltered(ctx)
+		return stream, cancel, err
 	}
 )
 
@@ -61,7 +65,7 @@ func New(ctx fabcontext.Client, chConfig fab.ChannelCfg, streamProvider StreamPr
 	logger.Debugf("Connecting to %s...", url)
 	connect, err := comm.NewStreamConnection(
 		ctx, chConfig,
-		func(grpcconn *grpc.ClientConn) (grpc.ClientStream, error) {
+		func(grpcconn *grpc.ClientConn) (grpc.ClientStream, func(), error) {
 			return streamProvider(pb.NewDeliverClient(grpcconn))
 		},
 		url, opts...,