Fix error API and implement SetReadDeadline

errors.go

@@ -17,14 +17,30 @@ const (
 )
 
 var (
-	errClosed = errors.New("session closed")
+	errClosed        = errors.New("session closed")
+	errNilConnection = errors.New("nil connection")
 )
 
+// SessionError is the error returned by operations on a closed session. It
+// carries the RoQ error code and reason the session was closed with, which
+// callers can inspect with errors.As.
 type SessionError struct {
 	code   uint64
 	reason string
 }
 
+// Code returns the RoQ error code the session was closed with, e.g.
+// ErrRoQPacketError.
+func (e SessionError) Code() uint64 {
+	return e.code
+}
+
+// Reason returns the reason phrase the session was closed with. It may be
+// empty.
+func (e SessionError) Reason() string {
+	return e.reason
+}
+
 func (e SessionError) Error() string {
 	return fmt.Sprintf("roq session error %v: %v", e.code, e.reason)
 }

receive_flow.go

@@ -6,6 +6,7 @@ import (
 	"errors"
 	"io"
 	"log"
+	"os"
 	"sync"
 	"time"
 
@@ -31,6 +32,13 @@ type ReceiveFlow struct {
 	streams map[int64]ReceiveStream
 	closed  bool
 
+	// deadlineCh is closed when the current read deadline expires or is
+	// replaced, which wakes blocked readers so that they pick up the new
+	// deadline. deadlineExpired records whether it was closed by expiry.
+	deadlineCh      chan struct{}
+	deadlineExpired bool
+	deadlineTimer   *time.Timer
+
 	qlog *qlog.Logger
 }
 
@@ -44,12 +52,13 @@ func newReceiveFlow(id uint64, receiveBufferSize int, qlog *qlog.Logger) *Receiv
 				return bytes.NewBuffer(make([]byte, 0, maxPacketBufferSize))
 			},
 		},
-		ctx:       ctx,
-		cancelCtx: cancel,
-		lock:      sync.Mutex{},
-		streams:   map[int64]ReceiveStream{},
-		closed:    false,
-		qlog:      qlog,
+		ctx:        ctx,
+		cancelCtx:  cancel,
+		lock:       sync.Mutex{},
+		streams:    map[int64]ReceiveStream{},
+		closed:     false,
+		deadlineCh: make(chan struct{}),
+		qlog:       qlog,
 	}
 }
 
@@ -145,12 +154,27 @@ func (f *ReceiveFlow) readStream(rs ReceiveStream) {
 // across calls: if buf is too small to hold the complete packet, the packet is
 // dropped and Read returns 0 and io.ErrShortBuffer. Callers should provide a
 // buffer large enough for the largest packet they expect to receive.
+//
+// Read blocks until a packet arrives, the flow is closed, or the deadline set
+// with SetReadDeadline passes.
 func (f *ReceiveFlow) Read(buf []byte) (int, error) {
-	packet, ok := <-f.buffer
-	if !ok {
-		return 0, f.ctx.Err()
+	for {
+		f.lock.Lock()
+		expired, deadlineCh := f.deadlineExpired, f.deadlineCh
+		f.lock.Unlock()
+		if expired {
+			return 0, os.ErrDeadlineExceeded
+		}
+		select {
+		case packet, ok := <-f.buffer:
+			if !ok {
+				return 0, f.ctx.Err()
+			}
+			return f.copyPacket(buf, packet)
+		case <-deadlineCh:
+			// The deadline expired or was replaced: re-read it and wait again.
+		}
 	}
-	return f.copyPacket(buf, packet)
 }
 
 func (f *ReceiveFlow) copyPacket(buf []byte, packet *bytes.Buffer) (int, error) {
@@ -163,11 +187,53 @@ func (f *ReceiveFlow) copyPacket(buf []byte, packet *bytes.Buffer) (int, error)
 	return n, nil
 }
 
+// SetReadDeadline sets the deadline for future Read calls and for any Read
+// that is currently blocked. Once the deadline has passed, Read returns
+// os.ErrDeadlineExceeded without waiting for a packet, even if packets are
+// still buffered, until the deadline is extended or cleared. A zero value for
+// t clears the deadline.
 func (f *ReceiveFlow) SetReadDeadline(t time.Time) error {
-	// TODO
+	f.lock.Lock()
+	defer f.lock.Unlock()
+	if f.deadlineTimer != nil {
+		f.deadlineTimer.Stop()
+		f.deadlineTimer = nil
+	}
+	// Wake the readers waiting on the previous deadline and give them a fresh
+	// channel to wait on, so that they observe the new one.
+	if !f.deadlineExpired {
+		close(f.deadlineCh)
+	}
+	f.deadlineCh = make(chan struct{})
+	f.deadlineExpired = false
+	if t.IsZero() {
+		return nil
+	}
+	d := time.Until(t)
+	if d <= 0 {
+		f.expireDeadline(f.deadlineCh)
+		return nil
+	}
+	ch := f.deadlineCh
+	f.deadlineTimer = time.AfterFunc(d, func() {
+		f.lock.Lock()
+		defer f.lock.Unlock()
+		f.expireDeadline(ch)
+	})
 	return nil
 }
 
+// expireDeadline reports the deadline ch belongs to as expired. It does
+// nothing if that deadline has meanwhile been replaced. Must be called with
+// f.lock held.
+func (f *ReceiveFlow) expireDeadline(ch chan struct{}) {
+	if f.deadlineCh != ch || f.deadlineExpired {
+		return
+	}
+	f.deadlineExpired = true
+	close(ch)
+}
+
 func (f *ReceiveFlow) ID() uint64 {
 	return f.id
 }
@@ -199,5 +265,9 @@ func (f *ReceiveFlow) closeBuffer() {
 		return
 	}
 	f.closed = true
+	if f.deadlineTimer != nil {
+		f.deadlineTimer.Stop()
+		f.deadlineTimer = nil
+	}
 	close(f.buffer)
 }

receive_flow_test.go

@@ -375,3 +375,153 @@ func TestReadStreamWrappedErrors(t *testing.T) {
 		})
 	}
 }
+
+// A deadline that passes while Read is blocked must unblock it.
+func TestReadDeadlineExpires(t *testing.T) {
+	f := newReceiveFlow(1, 10, nil)
+	if err := f.SetReadDeadline(time.Now().Add(20 * time.Millisecond)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	done := make(chan error, 1)
+	go func() {
+		_, err := f.Read(make([]byte, 100))
+		done <- err
+	}()
+	select {
+	case err := <-done:
+		if !errors.Is(err, os.ErrDeadlineExceeded) {
+			t.Errorf("Read error = %v, want os.ErrDeadlineExceeded", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Read did not return after the deadline passed")
+	}
+	// The deadline stays expired until it is extended or cleared.
+	if _, err := f.Read(make([]byte, 100)); !errors.Is(err, os.ErrDeadlineExceeded) {
+		t.Errorf("second Read error = %v, want os.ErrDeadlineExceeded", err)
+	}
+}
+
+// A deadline in the past times out immediately, even with a packet buffered.
+func TestReadDeadlineInThePast(t *testing.T) {
+	f := newReceiveFlow(1, 10, nil)
+	f.push(bytes.NewBuffer([]byte{0x42}))
+	if err := f.SetReadDeadline(time.Now().Add(-time.Second)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	if _, err := f.Read(make([]byte, 100)); !errors.Is(err, os.ErrDeadlineExceeded) {
+		t.Errorf("Read error = %v, want os.ErrDeadlineExceeded", err)
+	}
+	// Clearing the deadline makes the buffered packet readable again.
+	if err := f.SetReadDeadline(time.Time{}); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	n, err := f.Read(make([]byte, 100))
+	if err != nil || n != 1 {
+		t.Errorf("Read = %v, %v, want 1, nil", n, err)
+	}
+}
+
+// Setting a deadline must also apply to a Read that is already blocked.
+func TestSetReadDeadlineWhileReadBlocked(t *testing.T) {
+	f := newReceiveFlow(1, 10, nil)
+	done := make(chan error, 1)
+	go func() {
+		_, err := f.Read(make([]byte, 100))
+		done <- err
+	}()
+	// Give the reader a chance to block on the empty buffer first.
+	time.Sleep(10 * time.Millisecond)
+	if err := f.SetReadDeadline(time.Now().Add(20 * time.Millisecond)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	select {
+	case err := <-done:
+		if !errors.Is(err, os.ErrDeadlineExceeded) {
+			t.Errorf("Read error = %v, want os.ErrDeadlineExceeded", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("blocked Read did not pick up the new deadline")
+	}
+}
+
+// Extending the deadline of a blocked reader must keep it blocked, and the
+// replaced deadline must not fire on the new one.
+func TestExtendReadDeadlineWhileReadBlocked(t *testing.T) {
+	f := newReceiveFlow(1, 10, nil)
+	if err := f.SetReadDeadline(time.Now().Add(20 * time.Millisecond)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	done := make(chan error, 1)
+	go func() {
+		_, err := f.Read(make([]byte, 100))
+		done <- err
+	}()
+	if err := f.SetReadDeadline(time.Now().Add(time.Hour)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	select {
+	case err := <-done:
+		t.Fatalf("Read returned %v, want it to keep waiting for the extended deadline", err)
+	case <-time.After(100 * time.Millisecond):
+	}
+	f.push(bytes.NewBuffer([]byte{0x42}))
+	select {
+	case err := <-done:
+		if err != nil {
+			t.Errorf("Read error = %v, want nil", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Read did not return the pushed packet")
+	}
+}
+
+// A deadline must not keep Read from reporting the closed flow.
+func TestReadDeadlineAndClose(t *testing.T) {
+	f := newReceiveFlow(1, 10, nil)
+	if err := f.SetReadDeadline(time.Now().Add(time.Hour)); err != nil {
+		t.Fatalf("SetReadDeadline: %v", err)
+	}
+	done := make(chan error, 1)
+	go func() {
+		_, err := f.Read(make([]byte, 100))
+		done <- err
+	}()
+	if err := f.Close(); err != nil {
+		t.Fatalf("Close: %v", err)
+	}
+	select {
+	case err := <-done:
+		if !errors.Is(err, context.Canceled) {
+			t.Errorf("Read error = %v, want context.Canceled", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatal("Read did not return after the flow was closed")
+	}
+}
+
+// Racing SetReadDeadline against readers and expiring timers must not panic on
+// a double close of the deadline channel.
+func TestConcurrentSetReadDeadline(t *testing.T) {
+	f := newReceiveFlow(1, 10, nil)
+	var wg sync.WaitGroup
+	for range 8 {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for range 100 {
+				if err := f.SetReadDeadline(time.Now().Add(time.Millisecond)); err != nil {
+					t.Error(err)
+					return
+				}
+			}
+		}()
+	}
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for range 100 {
+			_, _ = f.Read(make([]byte, 100))
+		}
+	}()
+	wg.Wait()
+}

session.go

@@ -49,17 +49,24 @@ type Session struct {
 	receiveFlows      *syncMap[uint64, *ReceiveFlow]
 	receiveFlowBuffer *receiveFlowBuffer
 
-	mutex     sync.Mutex
-	closedErr error
-	wg        sync.WaitGroup
-	ctx       context.Context
-	cancelCtx context.CancelFunc
+	mutex sync.Mutex
+	// closedErr is the SessionError reported by operations on the closed
+	// session; connCloseErr is what closing the QUIC connection returned.
+	closedErr    error
+	connCloseErr error
+	wg           sync.WaitGroup
+	ctx          context.Context
+	cancelCtx    context.CancelFunc
 
 	qlog *qlog.Logger
 }
 
 // NewSession creates a new roq session. QUIC connection is handled by roq.
+// It returns an error if conn is nil.
 func NewSession(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger) (*Session, error) {
+	if conn == nil {
+		return nil, errNilConnection
+	}
 	s := newSession(conn, acceptDatagrams, qlogger)
 	s.start()
 
@@ -68,8 +75,12 @@ func NewSession(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger) (*S
 
 // NewSessionWithAppHandledConn creates a new roq session. QUIC connection is
 // handled by the application. HandleDatagram and HandleUniStreamWithFlowID have
-// to be called for each datagram / new stream.
+// to be called for each datagram / new stream. It returns an error if conn is
+// nil.
 func NewSessionWithAppHandledConn(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger) (*Session, error) {
+	if conn == nil {
+		return nil, errNilConnection
+	}
 	s := newSession(conn, acceptDatagrams, qlogger)
 
 	return s, nil
@@ -165,7 +176,7 @@ func (s *Session) close(code uint64, reason string) {
 	}
 	s.sendFlows.rangeFn(func(_ uint64, v *SendFlow) { _ = v.Close() })
 	s.receiveFlows.rangeFn(func(_ uint64, v *ReceiveFlow) { _ = v.Close() })
-	_ = s.conn.CloseWithError(code, reason)
+	s.connCloseErr = s.conn.CloseWithError(code, reason)
 	s.cancelCtx()
 }
 
@@ -181,10 +192,16 @@ func (s *Session) closeWithError(code uint64, reason string) {
 
 // Close closes the session and waits for its receive loops to return. It must
 // not be called from a callback running on one of those loops.
+//
+// Close returns the error from closing the underlying QUIC connection. It is
+// idempotent: later calls close nothing and report the same error as the call
+// that closed the session.
 func (s *Session) Close() error {
 	s.close(ErrRoQNoError, "")
 	s.wg.Wait()
-	return nil
+	s.mutex.Lock()
+	defer s.mutex.Unlock()
+	return s.connCloseErr
 }
 
 func (s *Session) receiveUniStreams() error {