Webhub race condition (#9863)

* fix webconn close semantics

Avoid race conditions in WebConn on shutdown by closing channels to guarantee all readers are notified. Wrap this with sync.Once to avoid closing the channel more than once.

* web_hub_test.go

* webhub: fix race condition on shutdown

Ensure that if the webhub shuts down in the process of sending, the caller unblocks given that the webhub will no longer consume incoming events.

* panic if app shutdown takes >30 seconds

* simplify WebConn::Pump channel semantics too

app/apptestlib.go

@@ -391,8 +391,25 @@ func (me *TestHelper) CreateScheme() (*model.Scheme, []*model.Role) {
 	return scheme, roles
 }
 
+func (me *TestHelper) ShutdownApp() {
+	done := make(chan bool)
+	go func() {
+		me.App.Shutdown()
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(30 * time.Second):
+		// panic instead of t.Fatal to terminate all tests in this package, otherwise the
+		// still running App could spuriously fail subsequent tests.
+		panic("failed to shutdown App within 30 seconds")
+	}
+}
+
 func (me *TestHelper) TearDown() {
-	me.App.Shutdown()
+	me.ShutdownApp()
+
 	os.Remove(me.tempConfigPath)
 	if err := recover(); err != nil {
 		StopTestStore()

app/web_conn.go

@@ -5,6 +5,7 @@ package app
 
 import (
 	"fmt"
+	"sync"
 	"sync/atomic"
 	"time"
 
@@ -39,6 +40,7 @@ type WebConn struct {
 	AllChannelMembers         map[string]string
 	LastAllChannelMembersTime int64
 	Sequence                  int64
+	closeOnce                 sync.Once
 	endWritePump              chan struct{}
 	pumpFinished              chan struct{}
 }
@@ -59,8 +61,8 @@ func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t goi18n.Tra
 		UserId:             session.UserId,
 		T:                  t,
 		Locale:             locale,
-		endWritePump:       make(chan struct{}, 2),
-		pumpFinished:       make(chan struct{}, 1),
+		endWritePump:       make(chan struct{}),
+		pumpFinished:       make(chan struct{}),
 	}
 
 	wc.SetSession(&session)
@@ -72,7 +74,9 @@ func (a *App) NewWebConn(ws *websocket.Conn, session model.Session, t goi18n.Tra
 
 func (wc *WebConn) Close() {
 	wc.WebSocket.Close()
-	wc.endWritePump <- struct{}{}
+	wc.closeOnce.Do(func() {
+		close(wc.endWritePump)
+	})
 	<-wc.pumpFinished
 }
 
@@ -105,16 +109,18 @@ func (c *WebConn) SetSession(v *model.Session) {
 }
 
 func (c *WebConn) Pump() {
-	ch := make(chan struct{}, 1)
+	ch := make(chan struct{})
 	go func() {
 		c.writePump()
-		ch <- struct{}{}
+		close(ch)
 	}()
 	c.readPump()
-	c.endWritePump <- struct{}{}
+	c.closeOnce.Do(func() {
+		close(c.endWritePump)
+	})
 	<-ch
 	c.App.HubUnregister(c)
-	c.pumpFinished <- struct{}{}
+	close(c.pumpFinished)
 }
 
 func (c *WebConn) readPump() {

app/web_hub.go

@@ -346,7 +346,10 @@ func (a *App) UpdateWebConnUserActivity(session model.Session, activityAt int64)
 }
 
 func (h *Hub) Register(webConn *WebConn) {
-	h.register <- webConn
+	select {
+	case h.register <- webConn:
+	case <-h.didStop:
+	}
 
 	if webConn.IsAuthenticated() {
 		webConn.SendHello()
@@ -356,22 +359,31 @@ func (h *Hub) Register(webConn *WebConn) {
 func (h *Hub) Unregister(webConn *WebConn) {
 	select {
 	case h.unregister <- webConn:
-	case <-h.stop:
+	case <-h.didStop:
 	}
 }
 
 func (h *Hub) Broadcast(message *model.WebSocketEvent) {
 	if h != nil && h.broadcast != nil && message != nil {
-		h.broadcast <- message
+		select {
+		case h.broadcast <- message:
+		case <-h.didStop:
+		}
 	}
 }
 
 func (h *Hub) InvalidateUser(userId string) {
-	h.invalidateUser <- userId
+	select {
+	case h.invalidateUser <- userId:
+	case <-h.didStop:
+	}
 }
 
 func (h *Hub) UpdateActivity(userId, sessionToken string, activityAt int64) {
-	h.activity <- &WebConnActivityMessage{UserId: userId, SessionToken: sessionToken, ActivityAt: activityAt}
+	select {
+	case h.activity <- &WebConnActivityMessage{UserId: userId, SessionToken: sessionToken, ActivityAt: activityAt}:
+	case <-h.didStop:
+	}
 }
 
 func getGoroutineId() int {

app/web_hub_test.go

@@ -5,6 +5,7 @@ import (
 	"net/http"
 	"net/http/httptest"
 	"testing"
+	"time"
 
 	"github.com/gorilla/websocket"
 	goi18n "github.com/nicksnyder/go-i18n/i18n"
@@ -60,3 +61,45 @@ func TestHubStopWithMultipleConnections(t *testing.T) {
 	defer wc2.Close()
 	defer wc3.Close()
 }
+
+// TestHubStopRaceCondition verifies that attempts to use the hub after it has shutdown does not
+// block the caller indefinitely.
+func TestHubStopRaceCondition(t *testing.T) {
+	th := Setup().InitBasic()
+	defer th.TearDown()
+
+	s := httptest.NewServer(http.HandlerFunc(dummyWebsocketHandler(t)))
+
+	th.App.HubStart()
+	wc1 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
+	defer wc1.Close()
+
+	hub := th.App.Srv.Hubs[0]
+	th.App.HubStop()
+	time.Sleep(5 * time.Second)
+
+	done := make(chan bool)
+	go func() {
+		wc4 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
+		wc5 := registerDummyWebConn(t, th.App, s.Listener.Addr(), th.BasicUser.Id)
+		hub.Register(wc4)
+		hub.Register(wc5)
+
+		hub.UpdateActivity("userId", "sessionToken", 0)
+
+		for i := 0; i <= BROADCAST_QUEUE_SIZE; i++ {
+			hub.Broadcast(&model.WebSocketEvent{})
+		}
+
+		hub.InvalidateUser("userId")
+		hub.Unregister(wc4)
+		hub.Unregister(wc5)
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(15 * time.Second):
+		t.Fatalf("hub call did not return within 15 seconds after stop")
+	}
+}