diff --git a/docs/ARCHITECTURE.md b/docs/ARCHITECTURE.md index aafc0ef..ee4cc6f 100644 --- a/docs/ARCHITECTURE.md +++ b/docs/ARCHITECTURE.md @@ -202,6 +202,13 @@ owner until its client stream, physical-session reservation and UDP admission slot have been released. Only then is opening completion published to Close and other waiters; cleanup runs outside the packet mutex. +Session worker ownership is tracked separately from the target index. A target +can release admission and be removed before its canceled workers finish; +PacketConn Close still joins those workers, including concurrent Close callers. +The client also owns each packet's cancellation watcher, so client Close joins +logical packet shutdown. Worker registration uses the same closed-state mutex +as publication, and all completion waits run outside those mutexes. + The SOCKS frontend also isolates lazy H3 target opening, but only for transports that explicitly implement `transport.PacketConcurrentSender`. Canonical target aliases share one FIFO lane; the WebClient and close-once wrappers preserve diff --git a/internal/tunnel/web_udp.go b/internal/tunnel/web_udp.go index c3119e1..20a273b 100644 --- a/internal/tunnel/web_udp.go +++ b/internal/tunnel/web_udp.go @@ -438,6 +438,9 @@ type webUDPPacketConn struct { closed bool sessions map[string]*webUDPClientSession pending map[string]*webUDPPendingSession + // Session workers remain owned after terminate removes their target index. + // Add is protected by mu and its closed gate; shutdown waits outside mu. + workers sync.WaitGroup // Reports only a newly verified CONNECT-UDP response, never local enqueue. // Protected by mu; invoked after releasing mu to avoid a Close lock cycle. onAuthenticated func(bool) @@ -456,7 +459,12 @@ func newWebUDPPacketConn(client *WebH3Client) *webUDPPacketConn { sessions: make(map[string]*webUDPClientSession), pending: make(map[string]*webUDPPendingSession), } + // DialPacket holds client.mu and has checked client.closed. Register before + // publication so client Close also joins this packet's shutdown. This watcher + // must not belong to packet.workers, which its own shutdown waits for. + client.workers.Add(1) go func() { + defer client.workers.Done() <-ctx.Done() packet.shutdown() }() @@ -674,6 +682,9 @@ func (p *webUDPPacketConn) shutdown() { for _, done := range pending { <-done } + // Removed sessions are no longer in the snapshot but may still be + // finishing their sender, drain, receive or shutdown-watch workers. + p.workers.Wait() }) } @@ -707,11 +718,20 @@ type webUDPClientSession struct { } func (s *webUDPClientSession) start() { + // Called under packet.mu before its closed gate can transition. No worker + // may be registered after shutdown begins waiting for packet ownership. s.wg.Add(4) - go s.receiveDatagrams() - go s.drainCapsules() - go s.sendDatagrams() - go s.watchDatagramSender() + s.packet.workers.Add(4) + go s.runWorker(s.receiveDatagrams) + go s.runWorker(s.drainCapsules) + go s.runWorker(s.sendDatagrams) + go s.runWorker(s.watchDatagramSender) +} + +func (s *webUDPClientSession) runWorker(worker func()) { + defer s.packet.workers.Done() + defer s.wg.Done() + worker() } func (s *webUDPClientSession) send(payload []byte) error { @@ -732,7 +752,6 @@ func (s *webUDPClientSession) send(payload []byte) error { } func (s *webUDPClientSession) sendDatagrams() { - defer s.wg.Done() defer close(s.senderDone) defer s.terminate() for { @@ -753,7 +772,6 @@ func (s *webUDPClientSession) sendDatagrams() { } func (s *webUDPClientSession) watchDatagramSender() { - defer s.wg.Done() <-s.ctx.Done() select { case <-s.senderDone: @@ -767,7 +785,6 @@ func (s *webUDPClientSession) watchDatagramSender() { } func (s *webUDPClientSession) receiveDatagrams() { - defer s.wg.Done() defer s.terminate() for { frame, err := s.stream.ReceiveDatagram(s.ctx) @@ -783,7 +800,6 @@ func (s *webUDPClientSession) receiveDatagrams() { } func (s *webUDPClientSession) drainCapsules() { - defer s.wg.Done() defer s.terminate() _, _ = io.Copy(io.Discard, s.stream) } diff --git a/internal/tunnel/web_udp_removed_worker_join_test.go b/internal/tunnel/web_udp_removed_worker_join_test.go new file mode 100644 index 0000000..09829ff --- /dev/null +++ b/internal/tunnel/web_udp_removed_worker_join_test.go @@ -0,0 +1,136 @@ +package tunnel + +import ( + "context" + "errors" + "net" + "sync" + "sync/atomic" + "testing" + "time" + + "github.com/apernet/quic-go" + "github.com/apernet/quic-go/http3" +) + +// This gate pauses the three context-watching session workers before their +// first cancellation select. It returns the real context's unchanged Done +// channel, and never fabricates an error, packet, stream, or authentication +// result. Holding a worker here models a scheduling pause, not blocked I/O. +type webUDPRemovedWorkerContext struct { + context.Context + budget context.Context + entered chan struct{} + release <-chan struct{} + calls atomic.Int32 +} + +func (c *webUDPRemovedWorkerContext) Done() <-chan struct{} { + if c.calls.Add(1) == 3 { + close(c.entered) + } + select { + case <-c.release: + case <-c.budget.Done(): + } + return c.Context.Done() +} + +func webUDPRemovedWorkerClose(closeOwner func() error, started chan<- struct{}, result chan<- error, done chan struct{}) { + defer close(done) + started <- struct{}{} + result <- closeOwner() +} + +func TestWebUDPCloseJoinsRemovedSessionWorkers(t *testing.T) { + t.Run("packet_close", func(t *testing.T) { testWebUDPCloseJoinsRemovedSessionWorkers(t, false) }) + t.Run("client_close", func(t *testing.T) { testWebUDPCloseJoinsRemovedSessionWorkers(t, true) }) +} + +func testWebUDPCloseJoinsRemovedSessionWorkers(t *testing.T, closeClient bool) { + x := newWebUDPLateCloseJoinFixture(t, false) + p := x.packet() + // Open a real authenticated request, then install the scheduling context + // before starting its workers. Registration/start is the same successful + // handoff tail as session(), under p.mu; no concurrent context mutation. + s, err := p.openSession(x.target.String()) + if err != nil { + t.Fatal(err) + } + release := make(chan struct{}) + ungate := sync.OnceFunc(func() { close(release) }) + ctx := &webUDPRemovedWorkerContext{Context: s.ctx, budget: x.f.ctx, entered: make(chan struct{}), release: release} + s.ctx = ctx + p.mu.Lock() + p.sessions[s.target] = s + s.start() + p.mu.Unlock() + workersDone := make(chan struct{}) + x.f.register(workersDone) + go func() { defer close(workersDone); s.wg.Wait() }() + // Always release workers before packet/client cleanup, including failure + // against the original implementation. Their independent join is not Close. + t.Cleanup(ungate) + x.f.wait(ctx.entered, "three actual session workers at scheduling gate") + // The fourth actual worker is reading Capsules on the real H3 stream. + // Resetting that direction makes it call terminate and remove its target. + s.stream.CancelRead(quic.StreamErrorCode(http3.ErrCodeRequestCanceled)) + x.f.wait(s.done, "actual Capsule worker terminated/removed session") + p.mu.Lock() + remaining := len(p.sessions) + p.mu.Unlock() + if remaining != 0 || len(x.c.udpSlots) != 0 { + t.Fatalf("removed target/UDP slots=%d/%d, want 0/0", remaining, len(x.c.udpSlots)) + } + x.requestJoin() + x.zeroUsers() + select { + case <-workersDone: + t.Fatal("scheduling gate did not retain actual session workers") + default: + } + started := make(chan struct{}, 2) + results := make(chan error, 2) + first, second := make(chan struct{}), make(chan struct{}) + closeOwner := p.Close + if closeClient { + closeOwner = x.c.Close + } + x.f.register(first) + x.f.register(second) + go webUDPRemovedWorkerClose(closeOwner, started, results, first) + x.f.wait(p.done, "Packet.Close cancellation/closed transition") + go webUDPRemovedWorkerClose(closeOwner, started, results, second) + for i := 0; i < 2; i++ { + select { + case <-started: + case <-x.f.ctx.Done(): + t.Fatal("Close callers did not start") + } + } + for _, closeDone := range []<-chan struct{}{first, second} { + select { + case <-closeDone: + t.Error("owner Close returned while removed session workers were still owned") + case <-time.After(100 * time.Millisecond): + } + } + select { + case <-workersDone: + t.Error("workers exited before their scheduling gate was released") + default: + } + ungate() + x.f.wait(workersDone, "independent removed session worker join") + x.f.wait(first, "first Close after worker completion") + x.f.wait(second, "concurrent Close after worker completion") + for i := 0; i < 2; i++ { + if err := <-results; err != nil { + t.Errorf("Close=%v", err) + } + } + if err := p.Send([]byte("after-close"), x.target.String()); !errors.Is(err, net.ErrClosed) { + t.Errorf("Send after Close=%v, want net.ErrClosed", err) + } + t.Log("real TLS13/H3 request removed and admission released; all four session workers independently joined after gate release") +}