multiplex: add (*Multiplex).CloseChan

multiplex.go

@@ -165,6 +165,11 @@ func (mp *Multiplex) IsClosed() bool {
 	}
 }
 
+// CloseChan returns a read-only channel which will be closed when the session is closed
+func (mp *Multiplex) CloseChan() <-chan struct{} {
+	return mp.closed
+}
+
 func (mp *Multiplex) sendMsg(timeout, cancel <-chan struct{}, header uint64, data []byte) error {
 	buf := pool.Get(len(data) + 20)
 

multiplex_test.go

@@ -394,6 +394,42 @@ func TestClosing(t *testing.T) {
 	}
 }
 
+func TestCloseChan(t *testing.T) {
+	a, b := net.Pipe()
+
+	mpa := NewMultiplex(a, false)
+	mpb := NewMultiplex(b, true)
+
+	ctx, cancel := context.WithTimeout(context.Background(), time.Second*1)
+	defer cancel()
+
+	_, err := mpb.NewStream(ctx)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	_, err = mpa.Accept()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	go func() {
+		mpa.Close()
+	}()
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("did not receive from CloseChan for mpa within timeout")
+	case <-mpa.CloseChan():
+	}
+
+	select {
+	case <-ctx.Done():
+		t.Fatal("did not receive from CloseChan for mpb within timeout")
+	case <-mpb.CloseChan():
+	}
+}
+
 func TestReset(t *testing.T) {
 	a, b := net.Pipe()