Removed wg from worker start. Fixed various races.

inbound/pool.go

@@ -278,7 +278,8 @@ func (p *Pool) addLocked(id string, socket types.Socket) error {
 	go func() {
 		defer cancel()
 		defer close(peer.done)
-		worker.Start(pool, &p.wg)
+		worker.Start(pool)
+		p.wg.Done()
 	}()
 
 	p.peers[id] = peer

inbound/worker.go

@@ -10,7 +10,7 @@ import (
 )
 
 type Worker interface {
-	Start(pool PoolPlugin, wg *sync.WaitGroup)
+	Start(pool PoolPlugin)
 	Stop()
 	Send(data []byte) error
 }
@@ -61,29 +61,28 @@ func NewWorker(
 	}, nil
 }
 
-func (w *DefaultWorker) Start(pool PoolPlugin, wg *sync.WaitGroup) {
+func (w *DefaultWorker) Start(pool PoolPlugin) {
 	messages := make(chan ReceivedMessage, 256)
 
-	var owg sync.WaitGroup
-	owg.Add(3)
+	var wg sync.WaitGroup
+	wg.Add(3)
 
 	go func() {
-		defer owg.Done()
+		defer wg.Done()
 		RunReader(w.ctx, pool.OnExit, w.conn, messages, w.heartbeat)
 	}()
 
 	go func() {
-		defer owg.Done()
+		defer wg.Done()
 		RunForwarder(w.id, w.ctx, messages, pool.Inbox, w.config.MaxQueueSize)
 	}()
 
 	go func() {
-		defer owg.Done()
+		defer wg.Done()
 		RunWatchdog(w.ctx, pool.OnExit, w.heartbeat, w.config.InactivityTimeout)
 	}()
 
-	owg.Wait()
-	wg.Done()
+	wg.Wait()
 }
 
 func (w *DefaultWorker) Stop() {

inbound/worker_test.go

@@ -70,7 +70,7 @@ func TestWorkerStart(t *testing.T) {
 		v := setupWorkerTest(t)
 		defer v.worker.Stop()
 
-		go v.worker.Start(v.pool, v.wg)
+		go v.worker.Start(v.pool)
 
 		v.incoming <- honeybeetest.MockIncomingData{
 			MsgType: websocket.TextMessage,
@@ -91,7 +91,7 @@ func TestWorkerStart(t *testing.T) {
 		v := setupWorkerTest(t)
 		defer v.worker.Stop()
 
-		go v.worker.Start(v.pool, v.wg)
+		go v.worker.Start(v.pool)
 
 		v.incoming <- honeybeetest.MockIncomingData{
 			Err: &websocket.CloseError{Code: websocket.CloseNormalClosure},
@@ -107,7 +107,7 @@ func TestWorkerStart(t *testing.T) {
 		v := setupWorkerTest(t)
 		defer v.worker.Stop()
 
-		go v.worker.Start(v.pool, v.wg)
+		go v.worker.Start(v.pool)
 
 		v.incoming <- honeybeetest.MockIncomingData{
 			Err: &websocket.CloseError{Code: websocket.CloseProtocolError},
@@ -142,7 +142,10 @@ func TestWorkerStart(t *testing.T) {
 
 		var wg sync.WaitGroup
 		wg.Add(1)
-		go worker.Start(pool, &wg)
+		go func() {
+			worker.Start(pool)
+			wg.Done()
+		}()
 
 		honeybeetest.Eventually(t, func() bool {
 			val := exitKind.Load()
@@ -154,7 +157,7 @@ func TestWorkerStart(t *testing.T) {
 func TestWorkerStop(t *testing.T) {
 	v := setupWorkerTest(t)
 
-	go v.worker.Start(v.pool, v.wg)
+	go func() { v.worker.Start(v.pool); v.wg.Done() }()
 
 	v.worker.Stop()
 
@@ -179,7 +182,7 @@ func TestWorkerSend(t *testing.T) {
 		v := setupWorkerTest(t)
 		defer v.worker.Stop()
 
-		go v.worker.Start(v.pool, v.wg)
+		go v.worker.Start(v.pool)
 
 		err := v.worker.Send([]byte("hello"))
 		assert.NoError(t, err)
@@ -214,7 +217,7 @@ func TestWorkerSend(t *testing.T) {
 		v := setupWorkerTest(t)
 		defer v.worker.Stop()
 
-		go v.worker.Start(v.pool, v.wg)
+		go v.worker.Start(v.pool)
 
 		v.conn.Close()