From 8052df4ef5f8eb74744d052b9bc8bade310bef8b Mon Sep 17 00:00:00 2001
From: Joachim Bauch <bauch@struktur.de>
Date: Fri, 23 Apr 2021 15:12:09 +0200
Subject: [PATCH] Wait for clients to be terminated before destroying hub in
 tests.

This should help with the flaky test failures on CI.
---
 hub.go      | 18 ++++++++++++++----
 hub_test.go |  5 ++++-
 2 files changed, 18 insertions(+), 5 deletions(-)

diff --git a/hub.go b/hub.go
index 5ff34cc..7c23dcc 100644
--- a/hub.go
+++ b/hub.go
@@ -105,8 +105,10 @@ type Hub struct {
 	info         *HelloServerMessageServer
 	infoInternal *HelloServerMessageServer
 
-	stopped  int32
-	stopChan chan bool
+	stopped         int32
+	stopChan        chan bool
+	readPumpActive  uint32
+	writePumpActive uint32
 
 	roomUpdated      chan *BackendServerRoomRequest
 	roomDeleted      chan *BackendServerRoomRequest
@@ -1858,6 +1860,14 @@ func (h *Hub) serveWs(w http.ResponseWriter, r *http.Request) {
 	}
 
 	h.processNewClient(client)
-	go client.WritePump()
-	go client.ReadPump()
+	go func(h *Hub) {
+		atomic.AddUint32(&h.writePumpActive, 1)
+		defer atomic.AddUint32(&h.writePumpActive, ^uint32(0))
+		client.WritePump()
+	}(h)
+	go func(h *Hub) {
+		atomic.AddUint32(&h.readPumpActive, 1)
+		defer atomic.AddUint32(&h.readPumpActive, ^uint32(0))
+		client.ReadPump()
+	}(h)
 }
diff --git a/hub_test.go b/hub_test.go
index a7569a5..85c800f 100644
--- a/hub_test.go
+++ b/hub_test.go
@@ -147,7 +147,9 @@ func WaitForHub(ctx context.Context, t *testing.T, h *Hub) {
 		h.ru.Lock()
 		rooms := len(h.rooms)
 		h.ru.Unlock()
-		if clients == 0 && rooms == 0 && sessions == 0 {
+		readActive := atomic.LoadUint32(&h.readPumpActive)
+		writeActive := atomic.LoadUint32(&h.writePumpActive)
+		if clients == 0 && rooms == 0 && sessions == 0 && readActive == 0 && writeActive == 0 {
 			break
 		}
 
@@ -163,6 +165,7 @@ func WaitForHub(ctx context.Context, t *testing.T, h *Hub) {
 			time.Sleep(time.Millisecond)
 		}
 	}
+
 }
 
 func validateBackendChecksum(t *testing.T, f func(http.ResponseWriter, *http.Request, *BackendClientRequest) *BackendClientResponse) func(http.ResponseWriter, *http.Request) {