From d37179d86ded9a5e253b12201164b756cb432483 Mon Sep 17 00:00:00 2001 From: Mathis Engelbart Date: Fri, 28 Aug 2026 17:01:49 +0200 Subject: [PATCH] Fix new flow registration races --- errors.go | 7 ++-- receive_flow_buffer_test.go | 44 ++++++++++++------------ receive_flow_test.go | 60 ++++++++++++++++---------------- session.go | 46 +++++++++++-------------- session_test.go | 68 +++++++++++++++++++++++++++++++++++++ sync_map.go | 18 +++++++--- 6 files changed, 155 insertions(+), 88 deletions(-) diff --git a/errors.go b/errors.go index 432c6e2..eff381a 100644 --- a/errors.go +++ b/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 diff --git a/receive_flow_buffer_test.go b/receive_flow_buffer_test.go index c88a06e..f2d4b21 100644 --- a/receive_flow_buffer_test.go +++ b/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)) + } } diff --git a/receive_flow_test.go b/receive_flow_test.go index 38fbe2b..6e57851 100644 --- a/receive_flow_test.go +++ b/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)) } } diff --git a/session.go b/session.go index 43b7d65..d1b9572 100644 --- a/session.go +++ b/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) { diff --git a/session_test.go b/session_test.go index 6d357e9..4711e0f 100644 --- a/session_test.go +++ b/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) +} diff --git a/sync_map.go b/sync_map.go index 251b65b..bb801db 100644 --- a/sync_map.go +++ b/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) {