Merge pull request #74 from hyperledger/websocket-close

Close ws client if context is canceled

Author: nguyer

pkg/wsclient/wsclient.go

@@ -115,7 +115,7 @@ type WSPreConnectHandler func(ctx context.Context) error
 type WSPostConnectHandler func(ctx context.Context, w WSClient) error
 
 func New(ctx context.Context, config *WSConfig, beforeConnect WSPreConnectHandler, afterConnect WSPostConnectHandler) (WSClient, error) {
-
+	l := log.L(ctx)
 	wsURL, err := buildWSUrl(ctx, config)
 	if err != nil {
 		return nil, err
@@ -160,6 +160,16 @@ func New(ctx context.Context, config *WSConfig, beforeConnect WSPreConnectHandle
 		w.headers.Set("Authorization", fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", authUsername, authPassword)))))
 	}
 
+	go func() {
+		select {
+		case <-ctx.Done():
+			l.Tracef("WS %s closing due to canceled context", w.url)
+			w.Close()
+		case <-w.closing:
+			l.Tracef("WS %s closing", w.url)
+		}
+	}()
+
 	return w, nil
 }
 

pkg/wsclient/wsclient_test.go

@@ -626,3 +626,14 @@ func TestTestTLSServerFailsBadCerts(t *testing.T) {
 	assert.Error(t, err)
 
 }
+
+func TestWSClientContextClosed(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	wsConfig := generateConfig()
+	wsc, err := New(ctx, wsConfig, nil, nil)
+	assert.NoError(t, err)
+	assert.NotNil(t, wsc)
+	err = wsc.Send(context.Background(), []byte{})
+	assert.Regexp(t, "FF00147", err)
+}