drpcwire: introduce MuxWriter for stream multiplexing

Replace per-stream drpcwire.Writer with a shared MuxWriter that uses a
dedicated drain goroutine. This decouples stream writes from transport I/O,
enabling true concurrent stream multiplexing.

Core Architecture:

MuxWriter (drpcwire/mux_writer.go):
- Single instance per Manager, shared across all streams
- Dedicated goroutine continuously drains buffered frames to transport
- Non-blocking WriteFrame: appends to buffer under lock, signals goroutine
- Double-buffer swap (buf/spare) minimizes time spent under lock

Stream changes (drpcstream/stream.go):
- wr field: *drpcwire.Writer -> *drpcwire.MuxWriter
- Removed: ManualFlush option, RawFlush, rawFlushLocked, checkRecvFlush
- sendPacketLocked/rawWriteLocked: WriteFrame only, no Flush

Manager integration (drpcmanager/manager.go):
- Creates MuxWriter with onError: m.terminate
- terminate(): wr.Stop() THEN tr.Close() — Stop makes WriteFrame reject
  immediately; transport close unblocks any in-flight Write in the drain
  goroutine
- Close(): <-wr.Done() to wait for drain goroutine exit

Key Design Decisions:

1. No explicit flush needed.
The drain goroutine continuously pulls from the buffer. Natural batching
occurs because appends accumulate while the goroutine is mid-Write. The
old cork pattern (delay flush until first recv) is unnecessary — appending
is a memcpy under lock, and the goroutine controls when transport I/O
happens.

2. sync.Cond over channels.
Signal coalescing: multiple WriteFrame calls while run() is in Write
produce a single wakeup. No allocation overhead. Stop uses closed bool +
Broadcast. Consistent with packetQueue.

3. Two-phase shutdown (Stop/Done split to avoid deadlock).
Stop() is non-blocking: sets closed, Broadcast, returns immediately.
Done() returns a channel that closes when run() exits. This split is
critical for the onError path: run() -> Write fails -> sets closed ->
onError -> terminate -> Stop (finds closed=true, noop) -> run() returns.
If Stop blocked until run() exited, this path would self-join.

4. run() owns its lifecycle on write failure.
When Write fails, run() sets closed=true itself before calling onError.
The subsequent onError -> terminate -> Stop path finds closed already set.
No coordination needed; the flag is idempotent.

5. No per-stream FrameWriter wrapper.
Initially considered a per-stream FrameWriter wrapping *MuxWriter, but
the only value was a closed check before append. That check lives in
MuxWriter.WriteFrame directly. Streams hold *MuxWriter and call
WriteFrame.

What this unlocks:
- Concurrent multiplexing: streams no longer serialize on writes
- Simplified stream: all flush/cork complexity removed
- Natural batching from continuous drain
- Direct error propagation: transport write failures fire manager
  termination via onError callback

Breaking changes:
- drpcstream.Options.ManualFlush removed
- Stream.RawFlush(), SetManualFlush() removed
- Stream constructor: *drpcwire.Writer -> *drpcwire.MuxWriter

Test coverage: 8 concurrency tests for MuxWriter covering concurrent
WriteFrame, write errors, onError->Stop deadlock path, blocked Write
unblocked by Close, concurrent Stop, abort semantics (Stop discards
buffered data), and write-during-active-drain. A data race in the
initial implementation (reading buf capacity without lock) was caught
by these tests and fixed.

Author: shubhamdhama

drpcconn/conn_test.go

@@ -40,31 +40,32 @@ func TestConn_InvokeFlushesSendClose(t *testing.T) {
 	invokeDone := make(chan struct{})
 
 	ctx.Run(func(ctx context.Context) {
-		wr := drpcwire.NewWriter(ps, 64)
+		wr := drpcwire.NewMuxWriter(ps, nil)
+		defer wr.StopWait()
 		rd := drpcwire.NewReader(ps)
 
 		_, _ = rd.ReadFrame()    // Invoke
 		_, _ = rd.ReadFrame()    // Message
 		pkt, _ := rd.ReadFrame() // CloseSend
 
-		_ = wr.WritePacket(drpcwire.Packet{
+		_ = wr.WriteFrame(drpcwire.Frame{
 			Data: []byte("qux"),
 			ID:   drpcwire.ID{Stream: pkt.ID.Stream, Message: 1},
 			Kind: drpcwire.KindMessage,
+			Done: true,
 		})
-		_ = wr.Flush()
 
 		_, _ = rd.ReadFrame() // Close
 		<-invokeDone          // wait for invoke to return
 
 		// ensure that any later packets are dropped by writing one
 		// before closing the transport.
-		for i := 0; i < 5; i++ {
-			_ = wr.WritePacket(drpcwire.Packet{
+		for range 5 {
+			_ = wr.WriteFrame(drpcwire.Frame{
 				ID:   drpcwire.ID{Stream: pkt.ID.Stream, Message: 2},
 				Kind: drpcwire.KindCloseSend,
+				Done: true,
 			})
-			_ = wr.Flush()
 		}
 
 		_ = ps.Close()
@@ -78,7 +79,7 @@ func TestConn_InvokeFlushesSendClose(t *testing.T) {
 
 	invokeDone <- struct{}{} // signal invoke has returned
 
-	// we should eventually notice the transport is closed
+	// we should eventually notice the transport is closed due to ps.Close()
 	select {
 	case <-conn.Closed():
 	case <-time.After(1 * time.Second):
@@ -95,7 +96,8 @@ func TestConn_InvokeSendsGrpcAndDrpcMetadata(t *testing.T) {
 	defer func() { assert.NoError(t, ps.Close()) }()
 
 	ctx.Run(func(ctx context.Context) {
-		wr := drpcwire.NewWriter(ps, 64)
+		wr := drpcwire.NewMuxWriter(ps, nil)
+		defer wr.StopWait()
 		rd := drpcwire.NewReader(ps)
 
 		md, err := rd.ReadFrame() // Metadata
@@ -114,12 +116,12 @@ func TestConn_InvokeSendsGrpcAndDrpcMetadata(t *testing.T) {
 		_, _ = rd.ReadFrame()    // Message
 		pkt, _ := rd.ReadFrame() // CloseSend
 
-		_ = wr.WritePacket(drpcwire.Packet{
+		_ = wr.WriteFrame(drpcwire.Frame{
 			Data: []byte("qux"),
 			ID:   drpcwire.ID{Stream: pkt.ID.Stream, Message: 1},
 			Kind: drpcwire.KindMessage,
+			Done: true,
 		})
-		_ = wr.Flush()
 
 		_, _ = rd.ReadFrame() // Close
 	})
@@ -181,6 +183,10 @@ func TestConn_NewStreamSendsGrpcAndDrpcMetadata(t *testing.T) {
 	s, err := conn.NewStream(ctx, "/com.example.Foo/Bar", testEncoding{})
 	assert.NoError(t, err)
 	_ = s.CloseSend()
+
+	// Wait for the server goroutine to read all frames before defers
+	// close the pipe. With MuxWriter, writes are asynchronous.
+	ctx.Wait()
 }
 
 func TestConn_encodeMetadata(t *testing.T) {

drpcmanager/active_streams_test.go

@@ -6,6 +6,7 @@ package drpcmanager
 import (
 	"context"
 	"errors"
+	"io"
 	"testing"
 
 	"github.com/zeebo/assert"
@@ -14,13 +15,19 @@ import (
 	"storj.io/drpc/drpcwire"
 )
 
-func testStream(id uint64) *drpcstream.Stream {
-	return drpcstream.New(context.Background(), id, &drpcwire.Writer{})
+func testMuxWriter(t *testing.T) *drpcwire.MuxWriter {
+	mw := drpcwire.NewMuxWriter(io.Discard, func(error) {})
+	t.Cleanup(func() { mw.Stop(); <-mw.Done() })
+	return mw
+}
+
+func testStream(t *testing.T, id uint64) *drpcstream.Stream {
+	return drpcstream.New(context.Background(), id, testMuxWriter(t))
 }
 
 func TestActiveStreams_AddAndGet(t *testing.T) {
 	streams := newActiveStreams()
-	s := testStream(1)
+	s := testStream(t, 1)
 
 	assert.NoError(t, streams.Add(1, s))
 
@@ -39,7 +46,7 @@ func TestActiveStreams_GetMissing(t *testing.T) {
 
 func TestActiveStreams_Remove(t *testing.T) {
 	streams := newActiveStreams()
-	s := testStream(1)
+	s := testStream(t, 1)
 
 	assert.NoError(t, streams.Add(1, s))
 	assert.Equal(t, streams.Len(), 1)
@@ -60,8 +67,8 @@ func TestActiveStreams_RemoveIdempotent(t *testing.T) {
 
 func TestActiveStreams_DuplicateAdd(t *testing.T) {
 	streams := newActiveStreams()
-	s1 := testStream(1)
-	s2 := testStream(1)
+	s1 := testStream(t, 1)
+	s2 := testStream(t, 1)
 
 	assert.NoError(t, streams.Add(1, s1))
 	assert.Error(t, streams.Add(1, s2))
@@ -76,13 +83,13 @@ func TestActiveStreams_AddAfterClose(t *testing.T) {
 	streams := newActiveStreams()
 	streams.Close(errors.New("closed"))
 
-	err := streams.Add(1, testStream(1))
+	err := streams.Add(1, testStream(t, 1))
 	assert.Error(t, err)
 }
 
 func TestActiveStreams_RemoveAfterClose(t *testing.T) {
 	streams := newActiveStreams()
-	s := testStream(1)
+	s := testStream(t, 1)
 	assert.NoError(t, streams.Add(1, s))
 
 	streams.Close(errors.New("closed"))
@@ -95,13 +102,12 @@ func TestActiveStreams_Len(t *testing.T) {
 	streams := newActiveStreams()
 	assert.Equal(t, streams.Len(), 0)
 
-	assert.NoError(t, streams.Add(1, testStream(1)))
+	assert.NoError(t, streams.Add(1, testStream(t, 1)))
 	assert.Equal(t, streams.Len(), 1)
 
-	assert.NoError(t, streams.Add(2, testStream(2)))
+	assert.NoError(t, streams.Add(2, testStream(t, 2)))
 	assert.Equal(t, streams.Len(), 2)
 
 	streams.Remove(1)
 	assert.Equal(t, streams.Len(), 1)
 }
-

drpcmanager/manager.go

@@ -30,10 +30,6 @@ var managerClosed = errs.Class("manager closed")
 
 // Options controls configuration settings for a manager.
 type Options struct {
-	// WriterBufferSize controls the size of the buffer that we will fill before
-	// flushing. Normal writes to streams typically issue a flush explicitly.
-	WriterBufferSize int
-
 	// Reader are passed to any readers the manager creates.
 	Reader drpcwire.ReaderOptions
 
@@ -55,7 +51,7 @@ type Options struct {
 // to the appropriate stream.
 type Manager struct {
 	tr   drpc.Transport
-	wr   *drpcwire.Writer
+	wr   *drpcwire.MuxWriter
 	rd   *drpcwire.Reader
 	opts Options
 
@@ -114,14 +110,15 @@ func New(tr drpc.Transport, kind ManagerKind) *Manager {
 func NewWithOptions(tr drpc.Transport, kind ManagerKind, opts Options) *Manager {
 	m := &Manager{
 		tr:   tr,
-		wr:   drpcwire.NewWriter(tr, opts.WriterBufferSize),
 		rd:   drpcwire.NewReaderWithOptions(tr, opts.Reader),
 		opts: opts,
 
 		invokes: make(chan invokeInfo),
 		kind:    kind,
 	}
 
+	m.wr = drpcwire.NewMuxWriter(tr, m.terminate)
+
 	// a buffer of size 1 allows NewServerStream to signal it is done creating a
 	// new server stream without having to coordinate with manageReader.
 	m.pdone.Make(1)
@@ -148,10 +145,14 @@ func (m *Manager) log(what string, cb func() string) {
 }
 
 // terminate puts the Manager into a terminal state and closes any resources
-// that need to be closed to signal the state change.
+// that need to be closed to signal the state change. The mux writer is stopped
+// before closing the transport so that WriteFrame immediately rejects new
+// writes; the subsequent transport close unblocks any in-flight Write in the
+// drain goroutine.
 func (m *Manager) terminate(err error) {
 	if m.sigs.term.Set(err) {
 		m.log("TERM", func() string { return fmt.Sprint(err) })
+		m.wr.Stop()
 		m.sigs.tport.Set(m.tr.Close())
 		if errors.Is(err, io.EOF) {
 			err = context.Canceled
@@ -317,8 +318,9 @@ func (m *Manager) Unblocked() <-chan struct{} {
 func (m *Manager) Close() error {
 	m.terminate(managerClosed.New("Close called"))
 
-	m.wg.Wait() // wait for all stream goroutines
-	m.sigs.read.Wait()
+	<-m.wr.Done()      // wait for writer goroutine to exit
+	m.wg.Wait()        // wait for all stream goroutines
+	m.sigs.read.Wait() // wait for reader goroutine to exit
 	m.sigs.tport.Wait()
 
 	return m.sigs.tport.Err()

drpcmanager/manager_test.go

@@ -48,7 +48,6 @@ func TestDrpcMetadata(t *testing.T) {
 		assert.NoError(t, stream.RawWrite(drpcwire.KindInvokeMetadata, buf))
 		assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("invoke")))
 		assert.NoError(t, stream.RawWrite(drpcwire.KindMessage, []byte("message")))
-		assert.NoError(t, stream.RawFlush())
 		assert.NoError(t, stream.Close())
 	})
 
@@ -105,7 +104,6 @@ func TestDrpcMetadataWithGRPCMetadataCompatMode(t *testing.T) {
 		assert.NoError(t, stream.RawWrite(drpcwire.KindInvokeMetadata, buf))
 		assert.NoError(t, stream.RawWrite(drpcwire.KindInvoke, []byte("invoke")))
 		assert.NoError(t, stream.RawWrite(drpcwire.KindMessage, []byte("message")))
-		assert.NoError(t, stream.RawFlush())
 		assert.NoError(t, stream.Close())
 	})
 

drpcmanager/random_test.go

@@ -45,24 +45,26 @@ func (rc *randClient) newSteam(ctx context.Context, man *Manager) (*drpcstream.S
 	return stream, err
 }
 
-func (rc *randClient) execute(t *testing.T, wr *drpcwire.Writer, op byte) {
+func (rc *randClient) execute(t *testing.T, wr *drpcwire.MuxWriter, op byte) {
 	cmd, arg, done := parseOp(op)
 
 	if !rc.active {
-		assert.NoError(t, wr.WritePacket(drpcwire.Packet{
+		assert.NoError(t, wr.WriteFrame(drpcwire.Frame{
 			Data: make([]byte, arg),
 			ID:   rc.id.incMessage(),
 			Kind: drpcwire.KindInvoke,
+			Done: true,
 		}))
 		rc.active = true
 	}
 
 	switch cmd {
 	case 0: // new invoke
 		if rc.active {
-			assert.NoError(t, wr.WritePacket(drpcwire.Packet{
+			assert.NoError(t, wr.WriteFrame(drpcwire.Frame{
 				ID:   rc.id.incMessage(),
 				Kind: drpcwire.KindClose,
+				Done: true,
 			}))
 		}
 
@@ -99,10 +101,11 @@ func (rc *randClient) execute(t *testing.T, wr *drpcwire.Writer, op byte) {
 		}))
 
 	case 2: // cause the remote side to close
-		assert.NoError(t, wr.WritePacket(drpcwire.Packet{
+		assert.NoError(t, wr.WriteFrame(drpcwire.Frame{
 			Data: []byte("remote-close"),
 			ID:   rc.id.incMessage(),
 			Kind: drpcwire.KindMessage,
+			Done: true,
 		}))
 
 	case 3, 4, 5, 6, 7: // send normal message
@@ -130,7 +133,7 @@ func (rs *randServer) newSteam(ctx context.Context, man *Manager) (*drpcstream.S
 	return man.NewClientStream(ctx, "rpc")
 }
 
-func (rs *randServer) execute(t *testing.T, wr *drpcwire.Writer, op byte) {
+func (rs *randServer) execute(t *testing.T, wr *drpcwire.MuxWriter, op byte) {
 	cmd, arg, done := parseOp(op)
 
 	switch cmd {
@@ -160,10 +163,11 @@ func (rs *randServer) execute(t *testing.T, wr *drpcwire.Writer, op byte) {
 		}))
 
 	case 2: // cause the remote side to close
-		assert.NoError(t, wr.WritePacket(drpcwire.Packet{
+		assert.NoError(t, wr.WriteFrame(drpcwire.Frame{
 			Data: []byte("remote-close"),
 			ID:   rs.id.incMessage(),
 			Kind: drpcwire.KindMessage,
+			Done: true,
 		}))
 
 	case 3, 4, 5, 6, 7: // send random message
@@ -185,7 +189,7 @@ func (rs *randServer) execute(t *testing.T, wr *drpcwire.Writer, op byte) {
 
 type runner interface {
 	newSteam(ctx context.Context, man *Manager) (*drpcstream.Stream, error)
-	execute(t *testing.T, wr *drpcwire.Writer, op byte)
+	execute(t *testing.T, wr *drpcwire.MuxWriter, op byte)
 }
 
 func runRandomized(t *testing.T, prog []byte, r runner) {
@@ -196,7 +200,9 @@ func runRandomized(t *testing.T, prog []byte, r runner) {
 	defer func() { _ = pc.Close() }()
 	defer func() { _ = ps.Close() }()
 
-	wr := drpcwire.NewWriter(pc, 0)
+	wr := drpcwire.NewMuxWriter(pc, func(error) {})
+	defer func() { wr.Stop(); <-wr.Done() }()
+
 	man := New(ps, Server)
 	defer func() { _ = man.Close() }()
 
@@ -223,7 +229,6 @@ func runRandomized(t *testing.T, prog []byte, r runner) {
 
 	for _, op := range prog {
 		r.execute(t, wr, op)
-		assert.NoError(t, wr.Flush())
 	}
 
 	assert.NoError(t, man.Close())