drpcmanager: fix data race between newStream and Close on WaitGroup

Move `wg.Add(1)` inside `activeStreams.Add` under its mutex so the
increment and the closed check are atomic. This prevents `wg.Add`
racing with `wg.Wait` in `Manager.Close` when terminate closes streams
before `newStream` finishes registering.

Author: Nukitt

drpcmanager/active_streams.go

@@ -25,8 +25,11 @@ func newActiveStreams() *activeStreams {
 }
 
 // Add adds a stream. It returns an error if the collection is closed or if a
-// stream with the same ID already exists.
-func (r *activeStreams) Add(id uint64, stream *drpcstream.Stream) error {
+// stream with the same ID already exists. If wg is non-nil, it is incremented
+// under the same lock that checks closed, so wg.Add and the closed check are
+// atomic with respect to Close (which sets closed before Manager.Close calls
+// wg.Wait).
+func (r *activeStreams) Add(id uint64, stream *drpcstream.Stream, wg *sync.WaitGroup) error {
 	if stream == nil {
 		return managerClosed.New("stream can't be nil")
 	}
@@ -41,6 +44,9 @@ func (r *activeStreams) Add(id uint64, stream *drpcstream.Stream) error {
 		return managerClosed.New("duplicate stream id")
 	}
 	r.streams[id] = stream
+	if wg != nil {
+		wg.Add(1)
+	}
 	return nil
 }
 

drpcmanager/active_streams_test.go

@@ -29,7 +29,7 @@ func TestActiveStreams_AddAndGet(t *testing.T) {
 	streams := newActiveStreams()
 	s := testStream(t, 1)
 
-	assert.NoError(t, streams.Add(1, s))
+	assert.NoError(t, streams.Add(1, s, nil))
 
 	got, ok := streams.Get(1)
 	assert.That(t, ok)
@@ -48,7 +48,7 @@ func TestActiveStreams_Remove(t *testing.T) {
 	streams := newActiveStreams()
 	s := testStream(t, 1)
 
-	assert.NoError(t, streams.Add(1, s))
+	assert.NoError(t, streams.Add(1, s, nil))
 	assert.Equal(t, streams.Len(), 1)
 
 	streams.Remove(1)
@@ -70,8 +70,8 @@ func TestActiveStreams_DuplicateAdd(t *testing.T) {
 	s1 := testStream(t, 1)
 	s2 := testStream(t, 1)
 
-	assert.NoError(t, streams.Add(1, s1))
-	assert.Error(t, streams.Add(1, s2))
+	assert.NoError(t, streams.Add(1, s1, nil))
+	assert.Error(t, streams.Add(1, s2, nil))
 
 	// original stream is still present
 	got, ok := streams.Get(1)
@@ -83,14 +83,14 @@ func TestActiveStreams_AddAfterClose(t *testing.T) {
 	streams := newActiveStreams()
 	streams.Close(errors.New("closed"))
 
-	err := streams.Add(1, testStream(t, 1))
+	err := streams.Add(1, testStream(t, 1), nil)
 	assert.Error(t, err)
 }
 
 func TestActiveStreams_RemoveAfterClose(t *testing.T) {
 	streams := newActiveStreams()
 	s := testStream(t, 1)
-	assert.NoError(t, streams.Add(1, s))
+	assert.NoError(t, streams.Add(1, s, nil))
 
 	streams.Close(errors.New("closed"))
 
@@ -102,10 +102,10 @@ func TestActiveStreams_Len(t *testing.T) {
 	streams := newActiveStreams()
 	assert.Equal(t, streams.Len(), 0)
 
-	assert.NoError(t, streams.Add(1, testStream(t, 1)))
+	assert.NoError(t, streams.Add(1, testStream(t, 1), nil))
 	assert.Equal(t, streams.Len(), 1)
 
-	assert.NoError(t, streams.Add(2, testStream(t, 2)))
+	assert.NoError(t, streams.Add(2, testStream(t, 2), nil))
 	assert.Equal(t, streams.Len(), 2)
 
 	streams.Remove(1)

drpcmanager/manager.go

@@ -324,15 +324,14 @@ func (m *Manager) newStream(ctx context.Context, sid uint64, kind drpc.StreamKin
 
 	stream := drpcstream.NewWithOptions(ctx, sid, m.wr, m.recvPool, opts)
 
-	if err := m.streams.Add(sid, stream); err != nil {
+	if err := m.streams.Add(sid, stream, &m.wg); err != nil {
 		return nil, err
 	}
 
 	if m.metrics.ShouldRecord() {
 		m.metrics.StreamsStarted.Inc(1)
 	}
 
-	m.wg.Add(1)
 	go m.manageStream(ctx, stream)
 
 	m.log("STREAM", stream.String)

drpcmanager/manager_test.go

@@ -588,3 +588,41 @@ func TestManager_ServerClientHangupCancels(t *testing.T) {
 	assert.That(t, errors.Is(err, context.Canceled))
 	assert.Equal(t, status.Code(drpc.ToRPCErr(err)), codes.Canceled)
 }
+
+// TestManager_ConcurrentCloseAndNewClientStream exercises the race between
+// Manager.Close (terminate → stop writer, close transport, close streams) and
+// Manager.NewClientStream (newStream → create stream, add to streams, wg.Add).
+// Under -race this would fail without wg.Add being atomic with the closed
+// check in activeStreams.Add.
+func TestManager_ConcurrentCloseAndNewClientStream(t *testing.T) {
+	for i := 0; i < 100; i++ {
+		cconn, sconn := net.Pipe()
+
+		cman := New(cconn, Client)
+
+		start := make(chan struct{})
+		done := make(chan struct{}, 2)
+
+		go func() {
+			<-start
+			_ = cman.Close()
+			done <- struct{}{}
+		}()
+
+		go func() {
+			<-start
+			stream, err := cman.NewClientStream(context.Background(), "rpc", 0)
+			if err == nil {
+				_ = stream.Close()
+			}
+			done <- struct{}{}
+		}()
+
+		close(start)
+		<-done
+		<-done
+
+		_ = cconn.Close()
+		_ = sconn.Close()
+	}
+}