Fixes shutdown deadlock in websocket client (#1221)

* Adds a failing test for deadlocks in shutting down the write pump

* Protects access to stopReadPump and isConnected with RWMutex to prevent races

* Makes the send channel buffered to allow re-queueing messages for transmission during writePump shutdown

pkg/websocket/client.go

@@ -79,16 +79,30 @@ type Client struct {
 	// Optional configuration parameters
 	cfg *Config
 
-	conn        *ws.Conn
-	done        chan struct{}
-	isConnected bool
-
-	NotifyExpired chan struct{}
-	notifyClose   chan error
-	send          chan *OutgoingMessage
-	stopReadPump  chan struct{}
-	stopWritePump chan struct{}
-	wg            *sync.WaitGroup
+	conn             *ws.Conn
+	done             chan struct{}
+	isConnected      bool
+	isConnectedMutex sync.RWMutex
+
+	NotifyExpired     chan struct{}
+	notifyClose       chan error
+	send              chan *OutgoingMessage
+	stopReadPumpMutex sync.RWMutex
+	stopReadPump      chan struct{}
+	stopWritePump     chan struct{}
+	wg                *sync.WaitGroup
+}
+
+func (c *Client) setIsConnected(newValue bool) {
+	c.isConnectedMutex.Lock()
+	defer c.isConnectedMutex.Unlock()
+	c.isConnected = newValue
+}
+
+func (c *Client) getIsConnected() bool {
+	c.isConnectedMutex.RLock()
+	defer c.isConnectedMutex.RUnlock()
+	return c.isConnected
 }
 
 // Connected returns a channel that's closed when the client has finished
@@ -97,7 +111,7 @@ func (c *Client) Connected() <-chan struct{} {
 	d := make(chan struct{})
 
 	go func() {
-		for !c.isConnected {
+		for !c.getIsConnected() {
 			time.Sleep(100 * time.Millisecond)
 		}
 		close(d)
@@ -109,7 +123,7 @@ func (c *Client) Connected() <-chan struct{} {
 // Run starts listening for incoming webhook requests from Stripe.
 func (c *Client) Run(ctx context.Context) {
 	for {
-		c.isConnected = false
+		c.setIsConnected(false)
 		c.cfg.Log.WithFields(log.Fields{
 			"prefix": "websocket.client.Run",
 		}).Debug("Attempting to connect to Stripe")
@@ -171,6 +185,8 @@ func (c *Client) Run(ctx context.Context) {
 // Close executes a proper closure handshake then closes the connection
 // list of close codes: https://datatracker.ietf.org/doc/html/rfc6455#section-7.4
 func (c *Client) Close(closeCode int, text string) {
+	c.stopReadPumpMutex.Lock()
+	defer c.stopReadPumpMutex.Unlock()
 	close(c.stopReadPump)
 	close(c.stopWritePump)
 	if c.conn != nil {
@@ -271,7 +287,7 @@ func (c *Client) connect(ctx context.Context) error {
 	defer resp.Body.Close()
 
 	c.changeConnection(conn)
-	c.isConnected = true
+	c.setIsConnected(true)
 
 	c.wg = &sync.WaitGroup{}
 	c.wg.Add(2)
@@ -289,6 +305,8 @@ func (c *Client) connect(ctx context.Context) error {
 
 // changeConnection takes a new connection and recreates the channels.
 func (c *Client) changeConnection(conn *ws.Conn) {
+	c.stopReadPumpMutex.Lock()
+	defer c.stopReadPumpMutex.Unlock()
 	c.conn = conn
 	c.notifyClose = make(chan error)
 	c.stopReadPump = make(chan struct{})
@@ -461,6 +479,12 @@ func (c *Client) writePump() {
 	}
 }
 
+func (c *Client) terminateReadPump() {
+	c.stopReadPumpMutex.Lock()
+	defer c.stopReadPumpMutex.Unlock()
+	c.stopReadPump <- struct{}{}
+}
+
 //
 // Public functions
 //
@@ -513,7 +537,7 @@ func NewClient(url string, webSocketID string, websocketAuthorizedFeature string
 		WebSocketAuthorizedFeature: websocketAuthorizedFeature,
 		cfg:                        cfg,
 		done:                       make(chan struct{}),
-		send:                       make(chan *OutgoingMessage),
+		send:                       make(chan *OutgoingMessage, 10),
 		NotifyExpired:              make(chan struct{}),
 	}
 }

pkg/websocket/client_test.go

@@ -3,6 +3,7 @@ package websocket
 import (
 	"context"
 	"encoding/json"
+	"fmt"
 	"net/http"
 	"net/http/httptest"
 	"strings"
@@ -11,7 +12,8 @@ import (
 	"time"
 
 	ws "github.com/gorilla/websocket"
-	// log "github.com/sirupsen/logrus"
+
+	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
 
@@ -176,55 +178,96 @@ func TestClientExpiredError(t *testing.T) {
 	}
 }
 
-/* func TestClientWebhookReconnect(t *testing.T) {
-	log.SetLevel(log.DebugLevel)
-	wg := &sync.WaitGroup{}
-	wg.Add(20)
+// This test is a regression test for deadlocks that can be encountered
+// when the write pump is interrupted by closed connections at inopportune
+// times.
+//
+// The goal is to simulate a scenario where the read pump is shut down but the
+// client still has messages to send. The read pump should be shut down because
+// in the majority of cases it is how the client ends up stopped. However, there's
+// no hard synchronization between the read and write pumps so we have to defend
+// against race conditions where the read side is shut down, hence this test.
+func TestWritePumpInterruptionRequeued(t *testing.T) {
+	serverReceivedMessages := make(chan string, 10)
+	wg := sync.WaitGroup{}
+
 	upgrader := ws.Upgrader{}
 	ts := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		wg.Add(1)
+
+		require.NotEmpty(t, r.UserAgent())
+		require.NotEmpty(t, r.Header.Get("X-Stripe-Client-User-Agent"))
+		require.Equal(t, "websocket-random-id", r.Header.Get("Websocket-Id"))
 		c, err := upgrader.Upgrade(w, r, nil)
 		require.NoError(t, err)
 
-		defer c.Close()
+		require.Equal(t, "websocket_feature=webhook-payloads", r.URL.RawQuery)
 
-		swg := &sync.WaitGroup{}
-		swg.Add(1)
+		defer c.Close()
 
-		go func() {
-			for {
-				if _, _, err := c.ReadMessage(); err != nil {
-					swg.Done()
-					return
-				}
-			}
-		}()
+		msgType, msg, err := c.ReadMessage()
+		require.NoError(t, err)
+		require.Equal(t, msgType, ws.TextMessage)
+		serverReceivedMessages <- string(msg)
 
-		swg.Wait()
+		// To simulate a forced reconnection, the server closes the connection
+		// after receiving any messages
+		c.WriteControl(ws.CloseMessage, ws.FormatCloseMessage(ws.CloseNormalClosure, ""), time.Now().Add(5*time.Second))
+		c.Close()
 		wg.Done()
 	}))
 
 	defer ts.Close()
 
 	url := "ws" + strings.TrimPrefix(ts.URL, "http")
 
-	rcvMsgChan := make(chan WebhookEvent)
-
 	client := NewClient(
 		url,
 		"websocket-random-id",
 		"webhook-payloads",
 		&Config{
-			EventHandler: EventHandlerFunc(func(msg IncomingMessage) {
-				rcvMsgChan <- *msg.WebhookEvent
-			}),
-			Log:               log.StandardLogger(),
-			ReconnectInterval: 10 * time.Second,
+			EventHandler: EventHandlerFunc(func(msg IncomingMessage) {}),
+			WriteWait:    10 * time.Second,
+			PongWait:     60 * time.Second,
+			PingPeriod:   60 * time.Hour,
 		},
 	)
 
 	go client.Run(context.Background())
 
 	defer client.Stop()
 
+	actualMessages := []string{}
+	connectedChan := client.Connected()
+	<-connectedChan
+	go func() { client.terminateReadPump() }()
+
+	for i := 0; i < 2; i++ {
+		client.SendMessage(NewEventAck(fmt.Sprintf("event_%d", i), fmt.Sprintf("event_%d", i)))
+		// Needed to deflake the test from racing against itself
+		// Something to do with the buffering
+		time.Sleep(100 * time.Millisecond)
+
+		msg := <-serverReceivedMessages
+		actualMessages = append(actualMessages, msg)
+		wg.Wait()
+	}
+
 	wg.Wait()
-} */
+
+	for {
+		exhausted := false
+		select {
+		case msg := <-serverReceivedMessages:
+			actualMessages = append(actualMessages, msg)
+		default:
+			exhausted = true
+		}
+
+		if exhausted {
+			break
+		}
+	}
+
+	assert.Len(t, actualMessages, 2)
+}