drpcmanager: replace streamBuffer with streamRegistry

Replace the single-stream streamBuffer with a stream registry that maps
stream IDs to stream objects. The registry currently holds at most one
active stream (two briefly during handoff), but provides the foundation
for stream multiplexing where callers will look up streams by ID directly.

Author: shubhamdhama

drpcmanager/manager.go

@@ -11,6 +11,7 @@ import (
 	"net"
 	"strings"
 	"sync"
+	"sync/atomic"
 	"syscall"
 	"time"
 
@@ -81,8 +82,12 @@ type Manager struct {
 
 	wg sync.WaitGroup // tracks active manageStream goroutines
 
-	sem  drpcsignal.Chan // held by the active stream
-	sbuf streamBuffer    // largest stream id created
+	// reg tracks active streams. Currently holds at most one active stream;
+	// a second may briefly coexist during stream handoff (old stream's
+	// Unregister races with new stream's Register).
+	reg *streamRegistry
+
+	sem drpcsignal.Chan // held by the active stream
 
 	pdone   drpcsignal.Chan // signals when NewServerStream has registered the new stream
 	invokes chan invokeInfo // completed invoke info from manageReader to NewServerStream
@@ -123,9 +128,6 @@ func NewWithOptions(tr drpc.Transport, opts Options) *Manager {
 		invokes: make(chan invokeInfo),
 	}
 
-	// initialize the stream buffer
-	m.sbuf.init()
-
 	// this semaphore controls the number of concurrent streams. it MUST be 1.
 	m.sem.Make(1)
 
@@ -134,6 +136,7 @@ func NewWithOptions(tr drpc.Transport, opts Options) *Manager {
 	m.pdone.Make(1)
 
 	m.pa = drpcwire.NewPacketAssembler()
+	m.reg = newStreamRegistry()
 
 	// set the internal stream options
 	drpcopts.SetStreamTransport(&m.opts.Stream.Internal, m.tr)
@@ -186,7 +189,7 @@ func (m *Manager) acquireSemaphore(ctx context.Context) error {
 // longer make any reads or writes on the transport. It exits early if the
 // context is canceled or the manager is terminated.
 func (m *Manager) waitForPreviousStream(ctx context.Context) (err error) {
-	prev := m.sbuf.Get()
+	prev := m.reg.GetLatest()
 	if prev == nil {
 		return nil
 	}
@@ -217,7 +220,8 @@ func (m *Manager) terminate(err error) {
 	if m.sigs.term.Set(err) {
 		m.log("TERM", func() string { return fmt.Sprint(err) })
 		m.sigs.tport.Set(m.tr.Close())
-		m.sbuf.Close()
+		m.reg.ForEach(func(s *drpcstream.Stream) { s.Close() })
+		m.reg.Close()
 	}
 }
 
@@ -249,7 +253,7 @@ func (m *Manager) manageReader() {
 			return
 		}
 
-		switch curr := m.sbuf.Get(); {
+		switch curr := m.reg.GetLatest(); {
 		// If the frame is for the current stream, deliver it.
 		case curr != nil && incomingFrame.ID.Stream == curr.ID():
 			if err := curr.HandleFrame(incomingFrame); err != nil {
@@ -319,9 +323,9 @@ func (m *Manager) handleInvokeFrame(fr drpcwire.Frame) error {
 	// Invoke packet completes the sequence. Send to NewServerStream.
 	select {
 	case m.invokes <- invokeInfo{sid: pkt.ID.Stream, data: pkt.Data, metadata: m.metadata}:
-		// Wait for NewServerStream to finish stream creation (including
-		// sbuf.Set) before reading the next frame. This guarantees curr
-		// is set for subsequent non-invoke packets.
+		// Wait for NewServerStream to finish stream creation before reading the
+		// next frame. This guarantees curr is set for subsequent non-invoke
+		// packets.
 		m.pdone.Recv()
 
 		m.pa.Reset()
@@ -346,10 +350,13 @@ func (m *Manager) newStream(ctx context.Context, sid uint64, kind drpc.StreamKin
 
 	stream := drpcstream.NewWithOptions(ctx, sid, m.wr, opts)
 
+	if err := m.reg.Register(sid, stream); err != nil {
+		return nil, err
+	}
+
 	m.wg.Add(1)
 	go m.manageStream(ctx, stream)
 
-	m.sbuf.Set(stream)
 	m.log("STREAM", stream.String)
 
 	return stream, nil
@@ -359,6 +366,7 @@ func (m *Manager) newStream(ctx context.Context, sid uint64, kind drpc.StreamKin
 // is finished, canceling the stream if the context is canceled.
 func (m *Manager) manageStream(ctx context.Context, stream *drpcstream.Stream) {
 	defer m.wg.Done()
+	defer m.reg.Unregister(stream.ID())
 	select {
 	case <-m.sigs.term.Signal():
 		err := m.sigs.term.Err()
@@ -429,7 +437,7 @@ func (m *Manager) Closed() <-chan struct{} {
 // the return result is only valid until the next call to NewClientStream or
 // NewServerStream.
 func (m *Manager) Unblocked() <-chan struct{} {
-	if prev := m.sbuf.Get(); prev != nil {
+	if prev := m.reg.GetLatest(); prev != nil {
 		return prev.Context().Done()
 	}
 	return closedCh
@@ -506,9 +514,8 @@ func (m *Manager) NewServerStream(ctx context.Context) (stream *drpcstream.Strea
 			}
 		}
 		stream, err := m.newStream(ctx, pkt.sid, drpc.StreamKindServer, rpc)
-		// Signal pdone only after stream registration so that
-		// manageReader sees the new stream via sbuf.Get() when it reads
-		// the next frame.
+		// Signal pdone only after stream registration so that manageReader sees
+		// the new stream in the registry when it reads the next frame.
 		m.pdone.Send()
 		return stream, rpc, err
 	}

drpcmanager/registry.go

@@ -0,0 +1,109 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcmanager
+
+import (
+	"sync"
+
+	"storj.io/drpc/drpcstream"
+)
+
+// streamRegistry is a thread-safe map of stream IDs to stream objects.
+// It is used by the Manager to track active streams for lifecycle management.
+type streamRegistry struct {
+	mu      sync.RWMutex
+	streams map[uint64]*drpcstream.Stream
+	closed  bool
+}
+
+func newStreamRegistry() *streamRegistry {
+	return &streamRegistry{
+		streams: make(map[uint64]*drpcstream.Stream),
+	}
+}
+
+// Register adds a stream to the registry. It returns an error if the registry
+// is closed or if a stream with the same ID is already registered.
+func (r *streamRegistry) Register(id uint64, stream *drpcstream.Stream) error {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if stream == nil {
+		return managerClosed.New("stream can't be nil")
+	}
+
+	if r.closed {
+		return managerClosed.New("register")
+	}
+	if _, ok := r.streams[id]; ok {
+		return managerClosed.New("duplicate stream id")
+	}
+	r.streams[id] = stream
+	return nil
+}
+
+// Unregister removes a stream from the registry. It is a no-op if the stream
+// is not registered or if the registry has been closed.
+func (r *streamRegistry) Unregister(id uint64) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if r.streams != nil {
+		delete(r.streams, id)
+	}
+}
+
+// Get returns the stream for the given ID and whether it was found.
+func (r *streamRegistry) Get(id uint64) (*drpcstream.Stream, bool) {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	s, ok := r.streams[id]
+	return s, ok
+}
+
+// GetLatest returns the stream with the highest ID, or nil if the registry is
+// empty. It iterates the map because the registry may briefly hold two streams
+// during stream handoff. This method should be removed once multiplexing is
+// supported and callers look up streams by ID directly.
+func (r *streamRegistry) GetLatest() *drpcstream.Stream {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	var latest *drpcstream.Stream
+	for _, s := range r.streams {
+		if latest == nil || latest.ID() < s.ID() {
+			latest = s
+		}
+	}
+	return latest
+}
+
+// Close marks the registry as closed, preventing future Register calls.
+// It does not cancel any streams.
+func (r *streamRegistry) Close() {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.closed = true
+}
+
+// ForEach calls fn for each registered stream. The registry is read-locked
+// during iteration.
+func (r *streamRegistry) ForEach(fn func(*drpcstream.Stream)) {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	for _, s := range r.streams {
+		fn(s)
+	}
+}
+
+// Len returns the number of registered streams.
+func (r *streamRegistry) Len() int {
+	r.mu.RLock()
+	defer r.mu.RUnlock()
+
+	return len(r.streams)
+}

drpcmanager/registry_test.go

@@ -0,0 +1,134 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcmanager
+
+import (
+	"context"
+	"testing"
+
+	"github.com/zeebo/assert"
+
+	"storj.io/drpc/drpcstream"
+	"storj.io/drpc/drpcwire"
+)
+
+func testStream(id uint64) *drpcstream.Stream {
+	return drpcstream.New(context.Background(), id, &drpcwire.Writer{})
+}
+
+func TestStreamRegistry_RegisterAndGet(t *testing.T) {
+	reg := newStreamRegistry()
+	s := testStream(1)
+
+	assert.NoError(t, reg.Register(1, s))
+
+	got, ok := reg.Get(1)
+	assert.That(t, ok)
+	assert.Equal(t, got, s)
+}
+
+func TestStreamRegistry_GetMissing(t *testing.T) {
+	reg := newStreamRegistry()
+
+	got, ok := reg.Get(42)
+	assert.That(t, !ok)
+	assert.Nil(t, got)
+}
+
+func TestStreamRegistry_Unregister(t *testing.T) {
+	reg := newStreamRegistry()
+	s := testStream(1)
+
+	assert.NoError(t, reg.Register(1, s))
+	assert.Equal(t, reg.Len(), 1)
+
+	reg.Unregister(1)
+
+	_, ok := reg.Get(1)
+	assert.That(t, !ok)
+	assert.Equal(t, reg.Len(), 0)
+}
+
+func TestStreamRegistry_UnregisterIdempotent(t *testing.T) {
+	reg := newStreamRegistry()
+
+	// must not panic when unregistering a non-existent ID
+	reg.Unregister(99)
+}
+
+func TestStreamRegistry_DuplicateRegister(t *testing.T) {
+	reg := newStreamRegistry()
+	s1 := testStream(1)
+	s2 := testStream(1)
+
+	assert.NoError(t, reg.Register(1, s1))
+	assert.Error(t, reg.Register(1, s2))
+
+	// original stream is still registered
+	got, ok := reg.Get(1)
+	assert.That(t, ok)
+	assert.Equal(t, got, s1)
+}
+
+func TestStreamRegistry_RegisterAfterClose(t *testing.T) {
+	reg := newStreamRegistry()
+	reg.Close()
+
+	err := reg.Register(1, testStream(1))
+	assert.Error(t, err)
+}
+
+func TestStreamRegistry_UnregisterAfterClose(t *testing.T) {
+	reg := newStreamRegistry()
+	s := testStream(1)
+	assert.NoError(t, reg.Register(1, s))
+
+	reg.Close()
+
+	// must not panic
+	reg.Unregister(1)
+}
+
+func TestStreamRegistry_Len(t *testing.T) {
+	reg := newStreamRegistry()
+	assert.Equal(t, reg.Len(), 0)
+
+	assert.NoError(t, reg.Register(1, testStream(1)))
+	assert.Equal(t, reg.Len(), 1)
+
+	assert.NoError(t, reg.Register(2, testStream(2)))
+	assert.Equal(t, reg.Len(), 2)
+
+	reg.Unregister(1)
+	assert.Equal(t, reg.Len(), 1)
+}
+
+func TestStreamRegistry_ForEach(t *testing.T) {
+	reg := newStreamRegistry()
+	s1 := testStream(1)
+	s2 := testStream(2)
+	s3 := testStream(3)
+
+	assert.NoError(t, reg.Register(1, s1))
+	assert.NoError(t, reg.Register(2, s2))
+	assert.NoError(t, reg.Register(3, s3))
+
+	seen := make(map[uint64]*drpcstream.Stream)
+	reg.ForEach(func(s *drpcstream.Stream) {
+		seen[s.ID()] = s
+	})
+
+	assert.Equal(t, len(seen), 3)
+	assert.Equal(t, seen[1], s1)
+	assert.Equal(t, seen[2], s2)
+	assert.Equal(t, seen[3], s3)
+}
+
+func TestStreamRegistry_ForEach_Empty(t *testing.T) {
+	reg := newStreamRegistry()
+
+	count := 0
+	reg.ForEach(func(_ *drpcstream.Stream) { count++ })
+	assert.Equal(t, count, 0)
+}