Implements "CloseChan" functionality for "Stream"

session_test.go

@@ -1562,3 +1562,146 @@ func TestCancelAccept(t *testing.T) {
 
 	wg.Wait()
 }
+
+func TestStreamCloseChan_SessionClose(t *testing.T) {
+	client, server := testClientServer()
+	defer client.Close()
+	var wg sync.WaitGroup
+	// server runs in the background and closes the stream as soon as received.
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		_, err := server.AcceptStream()
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		time.Sleep(time.Second)
+		server.Close()
+	}()
+	go func() {
+		defer wg.Done()
+		stream, err := client.OpenStream()
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		defer stream.Close()
+		// wait for close
+		select {
+		case <-stream.CloseChan():
+		case <-time.After(time.Second * 2):
+			t.Fatalf("err: closeChan is open, state=%d", stream.state)
+		}
+	}()
+	wg.Wait()
+}
+
+func TestStreamCloseChan_StreamClose(t *testing.T) {
+	client, server := testClientServer()
+	defer client.Close()
+	defer server.Close()
+	var wg sync.WaitGroup
+	// server runs in the background and closes the stream as soon as received.
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		stream, err := client.AcceptStream()
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		b := []byte("close")
+		_, err = stream.Read(b)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		t.Logf("C: <- %s", string(b))
+		_, err = stream.Write(b)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		t.Logf("C: -> %s", string(b))
+		stream.Close()
+		t.Log("C: closed")
+	}()
+
+	go func() {
+		defer wg.Done()
+		stream, err := server.OpenStream()
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		b := []byte("close")
+		stream.Write(b)
+		t.Logf("S: -> %s", string(b))
+		<-time.After(time.Second)
+		n, err := stream.Read(b)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		if n != len(b) {
+			t.Fatal("err: length mismatch")
+		}
+		t.Logf("S: <- %s", string(b))
+		stream.Read(b)
+
+		// check for close
+		select {
+		case <-stream.CloseChan():
+		default:
+			t.Fatalf("err: closeChan is open, state=%d", stream.state)
+		}
+	}()
+	wg.Wait()
+}
+
+func TestStreamCloseChan_StreamClose_WithFullBuffer(t *testing.T) {
+	client, server := testClientServer()
+	defer client.Close()
+	defer server.Close()
+	var wg sync.WaitGroup
+	// server runs in the background and closes the stream as soon as received.
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		stream, err := client.AcceptStream()
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		b := []byte("close")
+		_, err = stream.Read(b)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		_, err = stream.Write(b)
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		stream.Close()
+	}()
+	go func() {
+		defer wg.Done()
+		stream, err := server.OpenStream()
+		if err != nil {
+			t.Fatalf("err: %v", err)
+		}
+		b := []byte("close")
+		stream.Write(b)
+		<-time.After(time.Second)
+		for i := 0; i < 3; i++ {
+			n, err := stream.Read(b)
+			if err != nil && err != io.EOF {
+				t.Fatalf("err: %v", err)
+			}
+			if err != io.EOF && n != len(b) {
+				t.Fatal("err: length mismatch")
+			}
+		}
+
+		// closeChan should be closed when io.EOF is encountered.
+		select {
+		case <-stream.CloseChan():
+		default:
+			t.Fatalf("err: closeChan is open, state=%d", stream.state)
+		}
+	}()
+	wg.Wait()
+}

stream.go

@@ -57,6 +57,9 @@ type Stream struct {
 	// closeTimer is set with stateLock held to honor the StreamCloseTimeout
 	// setting on Session.
 	closeTimer *time.Timer
+
+	// closeChan is closed in case of stream being shutdown.
+	closeChan chan struct{}
 }
 
 // newStream is used to construct a new stream within
@@ -75,6 +78,7 @@ func newStream(session *Session, id uint32, state streamState) *Stream {
 		recvNotifyCh: make(chan struct{}, 1),
 		sendNotifyCh: make(chan struct{}, 1),
 		establishCh:  make(chan struct{}, 1),
+		closeChan:    make(chan struct{}, 1),
 	}
 	s.readDeadline.Store(time.Time{})
 	s.writeDeadline.Store(time.Time{})
@@ -106,11 +110,13 @@ START:
 		if s.recvBuf == nil || s.recvBuf.Len() == 0 {
 			s.recvLock.Unlock()
 			s.stateLock.Unlock()
+			s.closeCloseChan()
 			return 0, io.EOF
 		}
 		s.recvLock.Unlock()
 	case streamReset:
 		s.stateLock.Unlock()
+		s.closeCloseChan()
 		return 0, ErrConnectionReset
 	}
 	s.stateLock.Unlock()
@@ -181,9 +187,11 @@ START:
 		fallthrough
 	case streamClosed:
 		s.stateLock.Unlock()
+		s.closeCloseChan()
 		return 0, ErrStreamClosed
 	case streamReset:
 		s.stateLock.Unlock()
+		s.closeCloseChan()
 		return 0, ErrConnectionReset
 	}
 	s.stateLock.Unlock()
@@ -388,9 +396,29 @@ func (s *Stream) forceClose() {
 	s.stateLock.Lock()
 	s.state = streamClosed
 	s.stateLock.Unlock()
+	s.closeCloseChan()
 	s.notifyWaiting()
 }
 
+// CloseChan returns a read-only channel which is closed as
+// soon as the stream is closed. Note that when it is closed,
+// doesn't imply that the buffers are empty too.
+func (s *Stream) CloseChan() <-chan struct{} {
+	return s.closeChan
+}
+
+// IsClosed returns true in case of the stream being closed.
+// Note that when it is closed, doesn't imply that the buffers
+// are empty too.
+func (s *Stream) IsClosed() bool {
+	select {
+	case <-s.closeChan:
+		return true
+	default:
+		return false
+	}
+}
+
 // processFlags is used to update the state of the stream
 // based on set flags, if any. Lock must be held
 func (s *Stream) processFlags(flags uint16) error {
@@ -399,6 +427,7 @@ func (s *Stream) processFlags(flags uint16) error {
 
 	// Close the stream without holding the state lock
 	closeStream := false
+	closeCloseChan := false
 	defer func() {
 		if closeStream {
 			if s.closeTimer != nil {
@@ -408,6 +437,9 @@ func (s *Stream) processFlags(flags uint16) error {
 
 			s.session.closeStream(s.id)
 		}
+		if closeCloseChan {
+			s.closeCloseChan()
+		}
 	}()
 
 	if flags&flagACK == flagACK {
@@ -438,6 +470,7 @@ func (s *Stream) processFlags(flags uint16) error {
 	if flags&flagRST == flagRST {
 		s.state = streamReset
 		closeStream = true
+		closeCloseChan = true
 		s.notifyWaiting()
 	}
 	return nil
@@ -542,3 +575,13 @@ func (s *Stream) Shrink() {
 	}
 	s.recvLock.Unlock()
 }
+
+// closeCloseChan closes the closeChan, to mark the end of
+// stream lifetime.
+func (s *Stream) closeCloseChan() {
+	select {
+	case <-s.closeChan:
+	default:
+		close(s.closeChan)
+	}
+}