Add functional options for session setup

errors.go

@@ -19,6 +19,7 @@ const (
 var (
 	errClosed        = errors.New("session closed")
 	errNilConnection = errors.New("nil connection")
+	errInvalidOption = errors.New("invalid option")
 )
 
 // SessionError is the error returned by operations on a closed session. It

options.go

@@ -0,0 +1,63 @@
+package roq
+
+import "fmt"
+
+const (
+	// defaultReceiveBufferSize is the number of packets buffered per receive
+	// flow.
+	defaultReceiveBufferSize = 1024
+	// defaultUnknownFlowBufferSize is the number of receive flows buffered for
+	// flow IDs the application has not created a ReceiveFlow for yet.
+	defaultUnknownFlowBufferSize = 16
+)
+
+// Option configures a Session.
+type Option func(*sessionConfig) error
+
+type sessionConfig struct {
+	receiveBufferSize     int
+	unknownFlowBufferSize int
+}
+
+// newSessionConfig applies opts on top of the defaults.
+func newSessionConfig(opts []Option) (*sessionConfig, error) {
+	c := &sessionConfig{
+		receiveBufferSize:     defaultReceiveBufferSize,
+		unknownFlowBufferSize: defaultUnknownFlowBufferSize,
+	}
+	for _, opt := range opts {
+		if err := opt(c); err != nil {
+			return nil, err
+		}
+	}
+	return c, nil
+}
+
+// WithReceiveBufferSize sets the number of packets each ReceiveFlow of the
+// session buffers for the application to read. Packets that arrive for a flow
+// whose buffer is full are dropped. It defaults to 1024 packets and must be
+// greater than zero.
+func WithReceiveBufferSize(packets int) Option {
+	return func(c *sessionConfig) error {
+		if packets <= 0 {
+			return fmt.Errorf("%w: receive buffer size must be greater than zero, got %v", errInvalidOption, packets)
+		}
+		c.receiveBufferSize = packets
+		return nil
+	}
+}
+
+// WithUnknownFlowBufferSize sets the number of receive flows the session
+// buffers for flow IDs the application has not called NewReceiveFlow for yet.
+// Once the buffer is full, the oldest of those flows is evicted and its
+// streams are cancelled with ErrRoQUnknownFlowID. It defaults to 16 flows and
+// must be greater than zero.
+func WithUnknownFlowBufferSize(flows int) Option {
+	return func(c *sessionConfig) error {
+		if flows <= 0 {
+			return fmt.Errorf("%w: unknown flow buffer size must be greater than zero, got %v", errInvalidOption, flows)
+		}
+		c.unknownFlowBufferSize = flows
+		return nil
+	}
+}

options_test.go

@@ -0,0 +1,123 @@
+package roq
+
+import (
+	"errors"
+	"testing"
+	"time"
+
+	"github.com/quic-go/quic-go/quicvarint"
+)
+
+// Without options a session uses the documented defaults.
+func TestSessionDefaultKnobs(t *testing.T) {
+	c := newStubConn()
+	s, err := NewSession(c, true, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer closeWithin(t, s, 5*time.Second)
+
+	f, err := s.NewReceiveFlow(1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := cap(f.buffer); got != defaultReceiveBufferSize {
+		t.Errorf("receive flow buffers %d packets, want %d", got, defaultReceiveBufferSize)
+	}
+	if got := s.receiveFlowBuffer.maxLen; got != defaultUnknownFlowBufferSize {
+		t.Errorf("unknown flow buffer holds %d flows, want %d", got, defaultUnknownFlowBufferSize)
+	}
+}
+
+// WithReceiveBufferSize must reach the flows the application creates and the
+// flows the session creates for unknown flow IDs.
+func TestWithReceiveBufferSize(t *testing.T) {
+	c := newStubConn()
+	s, err := NewSession(c, true, nil, WithReceiveBufferSize(4))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer closeWithin(t, s, 5*time.Second)
+
+	f, err := s.NewReceiveFlow(1)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := cap(f.buffer); got != 4 {
+		t.Errorf("receive flow buffers %d packets, want 4", got)
+	}
+
+	// A datagram for a flow ID the application has not claimed yet creates a
+	// buffered flow, which must be sized the same way.
+	c.datagrams <- append(quicvarint.Append(nil, 2), 0x42)
+	unknown := awaitBufferedFlow(t, s, 2)
+	if got := cap(unknown.buffer); got != 4 {
+		t.Errorf("buffered flow holds %d packets, want 4", got)
+	}
+}
+
+// WithUnknownFlowBufferSize must bound how many unclaimed flows are buffered.
+func TestWithUnknownFlowBufferSize(t *testing.T) {
+	c := newStubConn()
+	s, err := NewSessionWithAppHandledConn(c, true, nil, WithUnknownFlowBufferSize(2))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer closeWithin(t, s, 5*time.Second)
+
+	for id := range uint64(4) {
+		s.HandleDatagram(append(quicvarint.Append(nil, id), 0x42))
+	}
+
+	s.receiveFlowBuffer.mutex.Lock()
+	defer s.receiveFlowBuffer.mutex.Unlock()
+	if got := len(s.receiveFlowBuffer.buffer); got != 2 {
+		t.Errorf("session buffered %d unknown flows, want 2", got)
+	}
+	// The two oldest flow IDs must have been evicted, not the newest.
+	for _, id := range []uint64{2, 3} {
+		if _, ok := s.receiveFlowBuffer.buffer[id]; !ok {
+			t.Errorf("flow %d was evicted, want the oldest flows evicted first", id)
+		}
+	}
+}
+
+// Invalid option values are rejected by both constructors rather than silently
+// producing a session with a useless buffer.
+func TestInvalidOptions(t *testing.T) {
+	for _, tc := range []struct {
+		name string
+		opt  Option
+	}{
+		{"zero receive buffer", WithReceiveBufferSize(0)},
+		{"negative receive buffer", WithReceiveBufferSize(-1)},
+		{"zero unknown flow buffer", WithUnknownFlowBufferSize(0)},
+		{"negative unknown flow buffer", WithUnknownFlowBufferSize(-1)},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			if _, err := NewSession(newStubConn(), true, nil, tc.opt); !errors.Is(err, errInvalidOption) {
+				t.Errorf("NewSession error = %v, want errInvalidOption", err)
+			}
+			if _, err := NewSessionWithAppHandledConn(newStubConn(), true, nil, tc.opt); !errors.Is(err, errInvalidOption) {
+				t.Errorf("NewSessionWithAppHandledConn error = %v, want errInvalidOption", err)
+			}
+		})
+	}
+}
+
+// awaitBufferedFlow waits for the session to buffer a flow for an unclaimed
+// flow ID and returns it.
+func awaitBufferedFlow(t *testing.T, s *Session, id uint64) *ReceiveFlow {
+	t.Helper()
+	for deadline := time.Now().Add(5 * time.Second); time.Now().Before(deadline); {
+		s.receiveFlowBuffer.mutex.Lock()
+		f, ok := s.receiveFlowBuffer.buffer[id]
+		s.receiveFlowBuffer.mutex.Unlock()
+		if ok {
+			return f
+		}
+		time.Sleep(time.Millisecond)
+	}
+	t.Fatalf("session never buffered a flow for ID %d", id)
+	return nil
+}

receive_flow_buffer.go

@@ -14,9 +14,9 @@ type receiveFlowBuffer struct {
 	queue  []uint64
 }
 
-func newReceiveFlowBuffer(maxLen uint) *receiveFlowBuffer {
+func newReceiveFlowBuffer(maxLen int) *receiveFlowBuffer {
 	b := &receiveFlowBuffer{
-		maxLen: int(maxLen),
+		maxLen: maxLen,
 		mutex:  sync.Mutex{},
 		buffer: map[uint64]*ReceiveFlow{},
 		queue:  []uint64{},

session.go

@@ -12,8 +12,6 @@ import (
 	"github.com/quic-go/quic-go/quicvarint"
 )
 
-const defaultReceiveBufferSize = 1024 // Number of packets to buffer
-
 type SendStream interface {
 	io.Writer
 	io.Closer
@@ -62,12 +60,16 @@ type Session struct {
 }
 
 // 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) {
+// It returns an error if conn is nil or if an option is invalid.
+func NewSession(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger, opts ...Option) (*Session, error) {
 	if conn == nil {
 		return nil, errNilConnection
 	}
-	s := newSession(conn, acceptDatagrams, qlogger)
+	config, err := newSessionConfig(opts)
+	if err != nil {
+		return nil, err
+	}
+	s := newSession(conn, acceptDatagrams, qlogger, config)
 	s.start()
 
 	return s, nil
@@ -76,25 +78,29 @@ 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. It returns an error if conn is
-// nil.
-func NewSessionWithAppHandledConn(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger) (*Session, error) {
+// nil or if an option is invalid.
+func NewSessionWithAppHandledConn(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger, opts ...Option) (*Session, error) {
 	if conn == nil {
 		return nil, errNilConnection
 	}
-	s := newSession(conn, acceptDatagrams, qlogger)
+	config, err := newSessionConfig(opts)
+	if err != nil {
+		return nil, err
+	}
+	s := newSession(conn, acceptDatagrams, qlogger, config)
 
 	return s, nil
 }
 
-func newSession(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger) *Session {
+func newSession(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger, config *sessionConfig) *Session {
 	ctx, cancel := context.WithCancel(context.Background())
 	return &Session{
-		receiveBufferSize: defaultReceiveBufferSize,
+		receiveBufferSize: config.receiveBufferSize,
 		acceptDatagrams:   acceptDatagrams,
 		conn:              conn,
 		sendFlows:         newSyncMap[uint64, *SendFlow](),
 		receiveFlows:      newSyncMap[uint64, *ReceiveFlow](),
-		receiveFlowBuffer: newReceiveFlowBuffer(16),
+		receiveFlowBuffer: newReceiveFlowBuffer(config.unknownFlowBufferSize),
 		mutex:             sync.Mutex{},
 		closedErr:         nil,
 		wg:                sync.WaitGroup{},