fix: join removed UDP session workers on close

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

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)
 }

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")
+}