Fix new flow registration races

errors.go

@@ -17,9 +17,10 @@ const (
 )
 
 var (
-	errClosed        = errors.New("session closed")
-	errNilConnection = errors.New("nil connection")
-	errInvalidOption = errors.New("invalid option")
+	errClosed          = errors.New("session closed")
+	errNilConnection   = errors.New("nil connection")
+	errInvalidOption   = errors.New("invalid option")
+	errDuplicateFlowID = errors.New("duplicate flow ID")
 )
 
 // SessionError is the error returned by operations on a closed session. It

receive_flow_buffer_test.go

@@ -61,30 +61,28 @@ func TestReceiveFlowBufferPoppedFlowIsNotEvicted(t *testing.T) {
 // and everything pushed into it is lost.
 func TestReceiveFlowBufferGetOrCreateIsAtomic(t *testing.T) {
 	const routines = 8
-	for range 100 {
-		b := newReceiveFlowBuffer(16)
-		flows := make([]*ReceiveFlow, routines)
-		var start sync.WaitGroup
-		var done sync.WaitGroup
-		start.Add(1)
-		for i := range routines {
-			done.Add(1)
-			go func() {
-				defer done.Done()
-				start.Wait()
-				flows[i] = b.getOrCreate(1, 1, nil)
-			}()
-		}
-		start.Done()
-		done.Wait()
+	b := newReceiveFlowBuffer(16)
+	flows := make([]*ReceiveFlow, routines)
+	var start sync.WaitGroup
+	var done sync.WaitGroup
+	start.Add(1)
+	for i := range routines {
+		done.Add(1)
+		go func() {
+			defer done.Done()
+			start.Wait()
+			flows[i] = b.getOrCreate(1, 1, nil)
+		}()
+	}
+	start.Done()
+	done.Wait()
 
-		for i, f := range flows {
-			if f != flows[0] {
-				t.Fatalf("goroutine %d got flow %p, goroutine 0 got %p, want the same flow", i, f, flows[0])
-			}
-		}
-		if len(b.queue) != 1 {
-			t.Fatalf("queue has %d entries, want 1", len(b.queue))
+	for i, f := range flows {
+		if f != flows[0] {
+			t.Fatalf("goroutine %d got flow %p, goroutine 0 got %p, want the same flow", i, f, flows[0])
 		}
 	}
+	if len(b.queue) != 1 {
+		t.Fatalf("queue has %d entries, want 1", len(b.queue))
+	}
 }

receive_flow_test.go

@@ -278,40 +278,38 @@ func TestCloseFlowTwice(t *testing.T) {
 // Closing a flow while its producers are still pushing must not lose a packet
 // that was queued, and must not send on the closed buffer.
 func TestConcurrentPushAndClose(t *testing.T) {
-	for range 100 {
-		f := newReceiveFlow(1, 10, nil)
-		var wg sync.WaitGroup
-		wg.Add(2)
-		go func() {
-			defer wg.Done()
-			f.push(bytes.NewBuffer([]byte{0x01}))
-		}()
-		go func() {
-			defer wg.Done()
-			_ = f.Close()
-		}()
-		wg.Wait()
+	f := newReceiveFlow(1, 10, nil)
+	var wg sync.WaitGroup
+	wg.Add(2)
+	go func() {
+		defer wg.Done()
+		f.push(bytes.NewBuffer([]byte{0x01}))
+	}()
+	go func() {
+		defer wg.Done()
+		_ = f.Close()
+	}()
+	wg.Wait()
 
-		// The push either won the race, in which case Read must hand the
-		// packet out, or it lost it and Read reports the closure. Both are
-		// correct; a queued packet that Read never returns is not.
-		buf := make([]byte, 1)
-		n, err := f.Read(buf)
-		if err == nil {
-			if n != 1 || buf[0] != 0x01 {
-				t.Fatalf("Read = %d bytes %v, want the pushed packet", n, buf[:n])
-			}
-			if _, err := f.Read(buf); !errors.Is(err, context.Canceled) {
-				t.Fatalf("Read on drained flow: %v, want %v", err, context.Canceled)
-			}
-			continue
-		}
-		if !errors.Is(err, context.Canceled) {
-			t.Fatalf("Read: %v, want %v", err, context.Canceled)
+	// The push either won the race, in which case Read must hand the
+	// packet out, or it lost it and Read reports the closure. Both are
+	// correct; a queued packet that Read never returns is not.
+	buf := make([]byte, 1)
+	n, err := f.Read(buf)
+	if err == nil {
+		if n != 1 || buf[0] != 0x01 {
+			t.Fatalf("Read = %d bytes %v, want the pushed packet", n, buf[:n])
 		}
-		if len(f.buffer) != 0 {
-			t.Fatalf("closed flow still holds %d packets no Read can return", len(f.buffer))
+		if _, err := f.Read(buf); !errors.Is(err, context.Canceled) {
+			t.Fatalf("Read on drained flow: %v, want %v", err, context.Canceled)
 		}
+		return
+	}
+	if !errors.Is(err, context.Canceled) {
+		t.Fatalf("Read: %v, want %v", err, context.Canceled)
+	}
+	if len(f.buffer) != 0 {
+		t.Fatalf("closed flow still holds %d packets no Read can return", len(f.buffer))
 	}
 }
 

session.go

@@ -3,7 +3,6 @@ package roq
 import (
 	"bytes"
 	"context"
-	"errors"
 	"io"
 	"sync"
 
@@ -46,6 +45,7 @@ type Session struct {
 	sendFlows         *syncMap[uint64, *SendFlow]
 	receiveFlows      *syncMap[uint64, *ReceiveFlow]
 	receiveFlowBuffer *receiveFlowBuffer
+	receiveFlowsMutex sync.Mutex
 
 	mutex sync.Mutex
 	// closedErr is the SessionError reported by operations on the closed
@@ -101,6 +101,7 @@ func newSession(conn Connection, acceptDatagrams bool, qlogger *qlog.Logger, con
 		sendFlows:         newSyncMap[uint64, *SendFlow](),
 		receiveFlows:      newSyncMap[uint64, *ReceiveFlow](),
 		receiveFlowBuffer: newReceiveFlowBuffer(config.unknownFlowBufferSize),
+		receiveFlowsMutex: sync.Mutex{},
 		mutex:             sync.Mutex{},
 		closedErr:         nil,
 		wg:                sync.WaitGroup{},
@@ -134,14 +135,11 @@ func (s *Session) NewSendFlow(id uint64) (*SendFlow, error) {
 	if err := s.isClosed(); err != nil {
 		return nil, err
 	}
-	if _, ok := s.sendFlows.get(id); ok {
-		return nil, errors.New("duplicate flow ID")
-	}
 	f := newFlow(s.conn, id, func() {
 		s.sendFlows.delete(id)
 	}, s.qlog)
-	if err := s.sendFlows.add(id, f); err != nil {
-		return nil, err
+	if _, ok := s.sendFlows.getOrInsert(id, f); !ok {
+		return nil, errDuplicateFlowID
 	}
 	return f, nil
 }
@@ -150,17 +148,16 @@ func (s *Session) NewReceiveFlow(id uint64) (*ReceiveFlow, error) {
 	if err := s.isClosed(); err != nil {
 		return nil, err
 	}
+	s.receiveFlowsMutex.Lock()
+	defer s.receiveFlowsMutex.Unlock()
 	if _, ok := s.receiveFlows.get(id); ok {
-		return nil, errors.New("duplicate flow ID")
+		return nil, errDuplicateFlowID
 	}
-	var f *ReceiveFlow
-	f = s.receiveFlowBuffer.pop(id)
+	f := s.receiveFlowBuffer.pop(id)
 	if f == nil {
 		f = newReceiveFlow(id, s.receiveBufferSize, s.qlog)
 	}
-	if err := s.receiveFlows.add(id, f); err != nil {
-		return nil, err
-	}
+	s.receiveFlows.set(id, f)
 	return f, nil
 }
 
@@ -234,6 +231,15 @@ func (s *Session) receiveDatagrams() error {
 	}
 }
 
+func (s *Session) receiveFlow(flowID uint64) *ReceiveFlow {
+	s.receiveFlowsMutex.Lock()
+	defer s.receiveFlowsMutex.Unlock()
+	if f, ok := s.receiveFlows.get(flowID); ok {
+		return f
+	}
+	return s.receiveFlowBuffer.getOrCreate(flowID, s.receiveBufferSize, s.qlog)
+}
+
 // HandleDatagram handles a datagram. If QUIC connection is handled by the
 // application, this function has to be called by the application for each
 // datagram that belongs to the roq connection.
@@ -259,14 +265,7 @@ func (s *Session) HandleDatagram(datagram []byte) {
 			},
 		})
 	}
-	if f, ok := s.receiveFlows.get(flowID); ok {
-		b := f.bufferPool.Get().(*bytes.Buffer)
-		b.Reset()
-		b.Write(datagram[quicvarint.Len(flowID):])
-		f.push(b)
-		return
-	}
-	f := s.receiveFlowBuffer.getOrCreate(flowID, s.receiveBufferSize, s.qlog)
+	f := s.receiveFlow(flowID)
 	b := f.bufferPool.Get().(*bytes.Buffer)
 	b.Reset()
 	b.Write(datagram[quicvarint.Len(flowID):])
@@ -284,12 +283,7 @@ func (s *Session) HandleUniStreamWithFlowID(flowID uint64, rs ReceiveStream) {
 			StreamID: uint64(rs.ID()),
 		})
 	}
-	if f, ok := s.receiveFlows.get(flowID); ok {
-		f.readStream(rs)
-		return
-	}
-	f := s.receiveFlowBuffer.getOrCreate(flowID, s.receiveBufferSize, s.qlog)
-	f.readStream(rs)
+	s.receiveFlow(flowID).readStream(rs)
 }
 
 func (s *Session) handleUniStream(rs ReceiveStream) {

session_test.go

@@ -7,6 +7,8 @@ import (
 	"sync"
 	"testing"
 	"time"
+
+	"github.com/quic-go/quic-go/quicvarint"
 )
 
 // stubSendStream is a SendStream that records writes and closes. writeFn, when
@@ -421,3 +423,69 @@ func TestNewSessionNilConnection(t *testing.T) {
 		t.Errorf("NewSessionWithAppHandledConn(nil) error = %v, want errNilConnection", err)
 	}
 }
+
+// A caller that loses the race to register a flow ID must not carry the
+// buffered flow off with it: the winner has to be the flow that already holds
+// the packets that arrived for the ID before it was registered.
+func TestNewReceiveFlowDoesNotStrandBufferedFlow(t *testing.T) {
+	s, err := NewSessionWithAppHandledConn(newStubConn(), true, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	// A packet for a flow ID nobody has registered yet: the session buffers
+	// it, and whoever registers the ID inherits that buffer.
+	s.HandleDatagram(append(quicvarint.Append(nil, 1), 'x'))
+
+	// Race registrations of that ID against each other.
+	type registration struct {
+		flow *ReceiveFlow
+		err  error
+	}
+	const routines = 8
+	registrations := make(chan registration, routines)
+	start := make(chan struct{})
+	var wg sync.WaitGroup
+	wg.Add(routines)
+	for range routines {
+		go func() {
+			defer wg.Done()
+			<-start
+			f, err := s.NewReceiveFlow(1)
+			registrations <- registration{flow: f, err: err}
+		}()
+	}
+	close(start)
+	wg.Wait()
+	close(registrations)
+
+	// Exactly one caller may win; the rest must be told the ID is taken.
+	var winners []*ReceiveFlow
+	for r := range registrations {
+		if r.err != nil {
+			if !errors.Is(r.err, errDuplicateFlowID) {
+				t.Fatalf("losing caller: error = %v, want errDuplicateFlowID", r.err)
+			}
+			continue
+		}
+		winners = append(winners, r.flow)
+	}
+	if len(winners) != 1 {
+		t.Fatalf("%d callers registered flow ID 1, want 1", len(winners))
+	}
+
+	// The winner must be able to read the buffered packet. The deadline makes
+	// a packet that was dropped with a stranded flow a failure, not a hang.
+	flow := winners[0]
+	if err := flow.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
+		t.Fatal(err)
+	}
+	buf := make([]byte, 8)
+	n, err := flow.Read(buf)
+	if err != nil {
+		t.Fatalf("Read: %v, want the packet buffered before the flow was registered", err)
+	}
+	if string(buf[:n]) != "x" {
+		t.Fatalf("Read returned %q, want %q", buf[:n], "x")
+	}
+	closeWithin(t, s, 5*time.Second)
+}

sync_map.go

@@ -1,7 +1,6 @@
 package roq
 
 import (
-	"errors"
 	"maps"
 	"sync"
 )
@@ -18,14 +17,23 @@ func newSyncMap[K comparable, V any]() *syncMap[K, V] {
 	}
 }
 
-func (m *syncMap[K, V]) add(k K, v V) error {
+// getOrInsert returns the value stored under k, inserting v if the map holds
+// no entry for k yet. The second return value reports whether v was inserted.
+func (m *syncMap[K, V]) getOrInsert(k K, v V) (V, bool) {
 	m.mutex.Lock()
 	defer m.mutex.Unlock()
-	if _, ok := m.elements[k]; ok {
-		return errors.New("duplicate entry")
+	if old, ok := m.elements[k]; ok {
+		return old, false
 	}
 	m.elements[k] = v
-	return nil
+	return v, true
+}
+
+// set stores v under k, replacing any entry the map holds for k.
+func (m *syncMap[K, V]) set(k K, v V) {
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
+	m.elements[k] = v
 }
 
 func (m *syncMap[K, V]) get(k K) (V, bool) {