fixup! enable stream multiplexing

Author: shubhamdhama

drpcconn/conn_test.go

@@ -41,7 +41,7 @@ func TestConn_InvokeFlushesSendClose(t *testing.T) {
 
 	ctx.Run(func(ctx context.Context) {
 		wr := drpcwire.NewMuxWriter(ps, nil)
-		defer wr.StopWait()
+		defer func() { wr.Stop(nil); <-wr.Done() }()
 		rd := drpcwire.NewReader(ps)
 
 		_, _ = rd.ReadFrame()    // Invoke
@@ -97,7 +97,7 @@ func TestConn_InvokeSendsGrpcAndDrpcMetadata(t *testing.T) {
 
 	ctx.Run(func(ctx context.Context) {
 		wr := drpcwire.NewMuxWriter(ps, nil)
-		defer wr.StopWait()
+		defer func() { wr.Stop(nil); <-wr.Done() }()
 		rd := drpcwire.NewReader(ps)
 
 		md, err := rd.ReadFrame() // Metadata

drpcmanager/active_streams.go

@@ -12,9 +12,10 @@ import (
 // activeStreams 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 activeStreams struct {
-	mu      sync.RWMutex
-	streams map[uint64]*drpcstream.Stream
-	closed  bool
+	mu       sync.RWMutex
+	streams  map[uint64]*drpcstream.Stream
+	closed   bool
+	closeErr error
 }
 
 func newActiveStreams() *activeStreams {
@@ -34,7 +35,7 @@ func (r *activeStreams) Add(id uint64, stream *drpcstream.Stream) error {
 	defer r.mu.Unlock()
 
 	if r.closed {
-		return managerClosed.New("add to closed collection")
+		return r.closeErr
 	}
 	if _, ok := r.streams[id]; ok {
 		return managerClosed.New("duplicate stream id")
@@ -73,6 +74,7 @@ func (r *activeStreams) Close(err error) {
 	defer r.mu.Unlock()
 
 	r.closed = true
+	r.closeErr = err
 	for id, s := range r.streams {
 		s.Cancel(err)
 		delete(r.streams, id)

drpcmanager/active_streams_test.go

@@ -17,7 +17,7 @@ import (
 
 func testMuxWriter(t *testing.T) *drpcwire.MuxWriter {
 	mw := drpcwire.NewMuxWriter(io.Discard, func(error) {})
-	t.Cleanup(func() { mw.Stop(); <-mw.Done() })
+	t.Cleanup(func() { mw.Stop(nil); <-mw.Done() })
 	return mw
 }
 

drpcmanager/manager.go

@@ -152,14 +152,14 @@ func (m *Manager) log(what string, cb func() string) {
 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
 			if m.kind == Client {
 				err = drpc.ClosedError.New("connection closed")
 			}
 		}
+		m.wr.Stop(err)
+		m.sigs.tport.Set(m.tr.Close())
 		m.streams.Close(err)
 	}
 }
@@ -191,7 +191,7 @@ func (m *Manager) manageReader() {
 
 		switch {
 		// if the packet is for an active stream, deliver it.
-		case ok && stream != nil:
+		case ok:
 			if err := stream.HandleFrame(incomingFrame); err != nil {
 				m.terminate(managerClosed.Wrap(err))
 				return

drpcmanager/random_test.go

@@ -201,7 +201,7 @@ func runRandomized(t *testing.T, prog []byte, r runner) {
 	defer func() { _ = ps.Close() }()
 
 	wr := drpcwire.NewMuxWriter(pc, func(error) {})
-	defer func() { wr.Stop(); <-wr.Done() }()
+	defer func() { wr.Stop(nil); <-wr.Done() }()
 
 	man := New(ps, Server)
 	defer func() { _ = man.Close() }()

drpcstream/stream_test.go

@@ -23,7 +23,7 @@ import (
 func testMuxWriter(t *testing.T) *drpcwire.MuxWriter {
 	t.Helper()
 	mw := drpcwire.NewMuxWriter(io.Discard, func(error) {})
-	t.Cleanup(mw.StopWait)
+	t.Cleanup(func() { mw.Stop(nil); <-mw.Done() })
 	return mw
 }
 

drpcwire/mux_writer.go

@@ -6,16 +6,15 @@ package drpcwire
 import (
 	"io"
 	"sync"
-
-	"storj.io/drpc"
 )
 
 type MuxWriter struct {
 	w       io.Writer
 	buf     []byte
 	mu      sync.Mutex
 	cond    *sync.Cond
-	closed  bool
+	closed   bool
+	closeErr error
 	onError func(error)
 	done    chan struct{}
 }
@@ -59,6 +58,7 @@ func (mw *MuxWriter) run() {
 				return
 			}
 			mw.closed = true
+			mw.closeErr = err
 			mw.mu.Unlock()
 			if mw.onError != nil {
 				mw.onError(err)
@@ -74,18 +74,18 @@ func (mw *MuxWriter) WriteFrame(fr Frame) (err error) {
 	mw.mu.Lock()
 	defer mw.mu.Unlock()
 	if mw.closed {
-		return drpc.ClosedError.New("mux writer closed")
+		return mw.closeErr
 	}
 	mw.buf = AppendFrame(mw.buf, fr)
 	mw.cond.Signal()
 	return nil
 }
 
-func (mw *MuxWriter) Stop() {
+func (mw *MuxWriter) Stop(err error) {
 	mw.mu.Lock()
-	if mw.closed {
-	} else {
+	if !mw.closed {
 		mw.closed = true
+		mw.closeErr = err
 		mw.cond.Broadcast()
 	}
 	mw.mu.Unlock()
@@ -94,8 +94,3 @@ func (mw *MuxWriter) Stop() {
 func (mw *MuxWriter) Done() <-chan struct{} {
 	return mw.done
 }
-
-func (mw *MuxWriter) StopWait() {
-	mw.Stop()
-	<-mw.Done()
-}

drpcwire/mux_writer_test.go

@@ -12,8 +12,6 @@ import (
 	"time"
 
 	"github.com/zeebo/assert"
-
-	"storj.io/drpc"
 )
 
 // blockingWriter blocks in Write until unblock is closed, then returns err.
@@ -80,7 +78,7 @@ func TestMuxWriter(t *testing.T) {
 	assert.NoError(t, err)
 
 	// Now stop the writer and close the pipe.
-	mw.Stop()
+	mw.Stop(errors.New("stopped"))
 	<-mw.Done()
 	pw.Close()
 	pr.Close()
@@ -90,12 +88,12 @@ func TestMuxWriter(t *testing.T) {
 
 func TestMuxWriter_WriteFrameAfterStop(t *testing.T) {
 	mw := NewMuxWriter(io.Discard, func(error) {})
-	mw.Stop()
+	mw.Stop(errors.New("stopped"))
 	<-mw.Done()
 
 	err := mw.WriteFrame(RandFrame())
 	assert.Error(t, err)
-	assert.That(t, drpc.ClosedError.Has(err))
+	assert.Equal(t, err.Error(), "stopped")
 }
 
 func TestMuxWriter_ConcurrentWriteFrame(t *testing.T) {
@@ -141,7 +139,7 @@ func TestMuxWriter_ConcurrentWriteFrame(t *testing.T) {
 	got := make([]byte, expSize)
 	_, err := io.ReadFull(pr, got)
 	assert.NoError(t, err)
-	mw.Stop()
+	mw.Stop(errors.New("stopped"))
 	<-mw.Done()
 	pw.Close()
 	pr.Close()
@@ -190,7 +188,7 @@ func TestMuxWriter_OnErrorCallingStopDoesNotDeadlock(t *testing.T) {
 	var mw *MuxWriter
 	mw = NewMuxWriter(fw, func(err error) {
 		// Simulate manager.terminate calling Stop.
-		mw.Stop()
+		mw.Stop(errors.New("stopped"))
 	})
 
 	assert.NoError(t, mw.WriteFrame(RandFrame()))
@@ -218,7 +216,7 @@ func TestMuxWriter_BlockedWriteUnblockedByClose(t *testing.T) {
 	}
 
 	// Simulate terminate: Stop, then unblock the writer (like tr.Close()).
-	mw.Stop()
+	mw.Stop(errors.New("stopped"))
 	bw.err = errors.New("closed")
 	close(bw.unblock)
 
@@ -241,7 +239,7 @@ func TestMuxWriter_ConcurrentStop(t *testing.T) {
 	for range n {
 		go func() {
 			defer wg.Done()
-			mw.Stop()
+			mw.Stop(errors.New("stopped"))
 		}()
 	}
 	wg.Wait()
@@ -276,7 +274,7 @@ func TestMuxWriter_StopDiscardsBufferedData(t *testing.T) {
 	}
 
 	// Stop without letting the blocked Write complete.
-	mw.Stop()
+	mw.Stop(errors.New("stopped"))
 	bw.err = errors.New("closed")
 	close(bw.unblock)
 
@@ -322,7 +320,7 @@ func TestMuxWriter_WriteFrameDuringActiveDrain(t *testing.T) {
 	close(g2.ch)
 
 	// Both batches were written. Stop and verify.
-	mw.Stop()
+	mw.Stop(errors.New("stopped"))
 	<-mw.Done()
 }