*: add opt-in stream multiplexing support

DRPC currently allows only one active stream per transport at a time.
This commit adds multiplexing as an opt-in mode (Options.Mux) that
enables concurrent streams over a single transport, while preserving
the original sequential behavior as the default.

Mux and non-mux paths have fundamentally different concurrency models,
so a single Manager with conditionals would add branching in hot paths.
Instead, two separate types are used: Manager (non-mux, unchanged) and
MuxManager (new, in manager_mux.go). MuxManager runs two goroutines:
manageReader routes packets to streams via a streamRegistry, and
manageWriter batches frames from a sharedWriteBuf into transport writes.

The Stream now accepts a StreamWriter interface instead of *Writer
directly, so both paths share the same Stream code. Non-mux uses *Writer
(direct transport writes), mux uses muxWriter (serializes each packet
atomically into the shared buffer to prevent frame interleaving).

Packet buffering is abstracted behind a packetStore interface with two
implementations: syncPacketBuffer (blocking single-slot, non-mux) and
queuePacketBuffer (non-blocking queue with sync.Pool recycling, mux).
RawRecv and MsgRecv branch on the mux flag for buffer lifecycle: non-mux
copies data then calls Done() to unblock the reader, mux takes ownership
and recycles after consumption.

HandlePacket now processes KindMessage before checking the term signal.
This is needed for mux mode where Put must always run to return pool
buffers. In non-mux mode, Put on a closed syncPacketBuffer returns
immediately, so there is no behavioral change. Though there is a brief
blocking window if terminate() has set the term signal but hasn't called
pbuf.Close() yet.

The reader's monotonicity check is relaxed from global to per-stream so
interleaved frames from different streams are accepted. Non-mux never
produces interleaved stream IDs, so this has no behavioral impact there.

Conn uses a streamManager interface satisfied by both Manager types,
branching once in the constructor. In mux mode, Invoke allocates a
per-call marshal buffer instead of reusing a shared one to support
concurrent calls. On the server side, ServeOne branches once to either
serveOneNonMux (sequential handleRPC) or serveOneMux (concurrent
handleRPC with WaitGroup).

Author: shubhamdhama

drpcconn/conn.go

@@ -29,10 +29,20 @@ type Options struct {
 	CollectStats bool
 }
 
+// streamManager is the interface satisfied by both drpcmanager.Manager (non-mux)
+// and drpcmanager.MuxManager (mux).
+type streamManager interface {
+	NewClientStream(ctx context.Context, rpc string) (*drpcstream.Stream, error)
+	Closed() <-chan struct{}
+	Unblocked() <-chan struct{}
+	Close() error
+}
+
 // Conn is a drpc client connection.
 type Conn struct {
 	tr   drpc.Transport
-	man  *drpcmanager.Manager
+	man  streamManager
+	mux  bool
 	mu   sync.Mutex
 	wbuf []byte
 
@@ -56,7 +66,12 @@ func NewWithOptions(tr drpc.Transport, opts Options) *Conn {
 		c.stats = make(map[string]*drpcstats.Stats)
 	}
 
-	c.man = drpcmanager.NewWithOptions(tr, opts.Manager)
+	c.mux = opts.Manager.Mux
+	if c.mux {
+		c.man = drpcmanager.NewMuxWithOptions(tr, opts.Manager)
+	} else {
+		c.man = drpcmanager.NewWithOptions(tr, opts.Manager)
+	}
 
 	return c
 }
@@ -100,8 +115,9 @@ func (c *Conn) Unblocked() <-chan struct{} { return c.man.Unblocked() }
 // Close closes the connection.
 func (c *Conn) Close() (err error) { return c.man.Close() }
 
-// Invoke issues the rpc on the transport serializing in, waits for a response, and
-// deserializes it into out. Only one Invoke or Stream may be open at a time.
+// Invoke issues the rpc on the transport serializing in, waits for a response,
+// and deserializes it into out. In non-mux mode, only one Invoke or Stream may
+// be open at a time. In mux mode, multiple calls may be open concurrently.
 func (c *Conn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, out drpc.Message) (err error) {
 	defer func() { err = drpc.ToRPCErr(err) }()
 
@@ -117,18 +133,28 @@ func (c *Conn) Invoke(ctx context.Context, rpc string, enc drpc.Encoding, in, ou
 	}
 	defer func() { err = errs.Combine(err, stream.Close()) }()
 
-	// we have to protect c.wbuf here even though the manager only allows one
-	// stream at a time because the stream may async close allowing another
-	// concurrent call to Invoke to proceed.
-	c.mu.Lock()
-	defer c.mu.Unlock()
-
-	c.wbuf, err = drpcenc.MarshalAppend(in, enc, c.wbuf[:0])
-	if err != nil {
-		return err
+	var data []byte
+	if c.mux {
+		// Per-call buffer allocation for concurrent access.
+		data, err = drpcenc.MarshalAppend(in, enc, nil)
+		if err != nil {
+			return err
+		}
+	} else {
+		// We have to protect c.wbuf here even though the manager only allows
+		// one stream at a time because the stream may async close allowing
+		// another concurrent call to Invoke to proceed.
+		c.mu.Lock()
+		defer c.mu.Unlock()
+
+		c.wbuf, err = drpcenc.MarshalAppend(in, enc, c.wbuf[:0])
+		if err != nil {
+			return err
+		}
+		data = c.wbuf
 	}
 
-	if err := c.doInvoke(stream, enc, rpc, c.wbuf, metadata, out); err != nil {
+	if err := c.doInvoke(stream, enc, rpc, data, metadata, out); err != nil {
 		return err
 	}
 	return nil

drpcconn/conn_test.go

@@ -180,7 +180,10 @@ func TestConn_NewStreamSendsGrpcAndDrpcMetadata(t *testing.T) {
 	})
 	s, err := conn.NewStream(ctx, "/com.example.Foo/Bar", testEncoding{})
 	assert.NoError(t, err)
-	_ = s.CloseSend()
+
+	assert.NoError(t, s.CloseSend())
+
+	ctx.Wait()
 }
 
 func TestConn_encodeMetadata(t *testing.T) {

drpcmanager/frame_queue_test.go

@@ -0,0 +1,82 @@
+// Copyright (C) 2026 Cockroach Labs.
+// See LICENSE for copying information.
+
+package drpcmanager
+
+import (
+	"testing"
+
+	"github.com/zeebo/assert"
+	"storj.io/drpc/drpcwire"
+)
+
+func TestSharedWriteBuf_AppendDrain(t *testing.T) {
+	sw := newSharedWriteBuf()
+
+	pkt := drpcwire.Packet{
+		Data: []byte("hello"),
+		ID:   drpcwire.ID{Stream: 1, Message: 2},
+		Kind: drpcwire.KindMessage,
+	}
+
+	assert.NoError(t, sw.Append(pkt))
+
+	// Drain should return serialized bytes.
+	data := sw.Drain(nil)
+	assert.That(t, len(data) > 0)
+
+	// Parse the frame back out to verify correctness.
+	_, got, ok, err := drpcwire.ParseFrame(data)
+	assert.NoError(t, err)
+	assert.That(t, ok)
+	assert.DeepEqual(t, got.Data, pkt.Data)
+	assert.Equal(t, got.ID.Stream, pkt.ID.Stream)
+	assert.Equal(t, got.ID.Message, pkt.ID.Message)
+	assert.Equal(t, got.Kind, pkt.Kind)
+	assert.Equal(t, got.Done, true)
+}
+
+func TestSharedWriteBuf_CloseIdempotent(t *testing.T) {
+	sw := newSharedWriteBuf()
+	sw.Close()
+	sw.Close() // must not panic
+}
+
+func TestSharedWriteBuf_AppendAfterClose(t *testing.T) {
+	sw := newSharedWriteBuf()
+	sw.Close()
+
+	err := sw.Append(drpcwire.Packet{})
+	assert.Error(t, err)
+}
+
+func TestSharedWriteBuf_WaitAndDrainBlocks(t *testing.T) {
+	sw := newSharedWriteBuf()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		data, ok := sw.WaitAndDrain(nil)
+		assert.That(t, ok)
+		assert.That(t, len(data) > 0)
+	}()
+
+	// Append should wake the blocked WaitAndDrain.
+	assert.NoError(t, sw.Append(drpcwire.Packet{Data: []byte("a")}))
+	<-done
+}
+
+func TestSharedWriteBuf_WaitAndDrainCloseEmpty(t *testing.T) {
+	sw := newSharedWriteBuf()
+
+	done := make(chan struct{})
+	go func() {
+		defer close(done)
+		_, ok := sw.WaitAndDrain(nil)
+		assert.That(t, !ok)
+	}()
+
+	// Close on empty buffer should return ok=false.
+	sw.Close()
+	<-done
+}

drpcmanager/manager.go

@@ -60,6 +60,11 @@ type Options struct {
 	// handling. When enabled, the server stream will decode incoming metadata
 	// into grpc metadata in the context.
 	GRPCMetadataCompatMode bool
+
+	// Mux enables stream multiplexing on the transport, allowing multiple
+	// concurrent streams. When false (default), the manager uses the
+	// original single-stream-at-a-time behavior.
+	Mux bool
 }
 
 // Manager handles the logic of managing a transport for a drpc client or
@@ -306,6 +311,7 @@ func (m *Manager) newStream(ctx context.Context, sid uint64, kind drpc.StreamKin
 		drpcopts.SetStreamStats(&opts.Internal, cb(rpc))
 	}
 
+	m.wr.Reset()
 	stream := drpcstream.NewWithOptions(ctx, sid, m.wr, opts)
 	select {
 	case m.streams <- streamInfo{ctx: ctx, stream: stream}: