From ebb366877a68dcddaf28b0fce5f55ec87356ad12 Mon Sep 17 00:00:00 2001 From: balaji Date: Thu, 20 Aug 2026 12:17:27 -0700 Subject: [PATCH 1/2] fix(worker): give each proxy host its own QUIC socket The worker opened one UDP socket and reused it for every QUIC dial to every proxy pod for the life of the process. It was created on the first dial and released only by HttpProxy.Close(), which runs at process exit. A UDP load balancer hashes on the source port, so one socket means every proxy pod is reached over a single balancer flow. Two things follow. There is no isolation. If that flow is pinned to an instance that has gone away, QUIC to every pod fails, healthy ones included. On a staging reproduction the worker logged 3,796 dial timeouts in eight minutes against pods that had been healthy for 38 minutes, while the balancer targets were registered healthy, DNS resolved correctly, and the receiving proxies reported no inbound QUIC at all. And it does not recover. A flow is re-hashed only once it goes idle, and a retrying worker never lets it idle, so it is never moved to a healthy instance. With the socket surviving until process exit, the condition lasts as long as the worker does. Replacing instances has been the reliable remedy precisely because new pods get new source ports. Key the transport by host, as the connection cache already is. Recovery then falls out of the existing lifecycle: a failed dial already drops the host from the cache, so releasing its socket there means the next attempt comes from a new source port the balancer is free to place elsewhere. No timers, no new state. The field is written once at construction rather than lazily in the dial goroutine, because the dial goroutine and the close path would otherwise race on it. Also fixes a pre-existing race the new concurrent test exposed: quic-go's validateConfig writes defaults back into the Config it is handed, and every dial passed the same pointer. Each dial now gets its own copy. Co-Authored-By: Balaji Ganesan --- src/libraries/go/worker/proxy/BUILD.bazel | 2 + src/libraries/go/worker/proxy/h3.go | 88 ++++++--- .../go/worker/proxy/host_transport_test.go | 186 ++++++++++++++++++ 3 files changed, 247 insertions(+), 29 deletions(-) create mode 100644 src/libraries/go/worker/proxy/host_transport_test.go diff --git a/src/libraries/go/worker/proxy/BUILD.bazel b/src/libraries/go/worker/proxy/BUILD.bazel index cca7ca1ab..3a15faa87 100644 --- a/src/libraries/go/worker/proxy/BUILD.bazel +++ b/src/libraries/go/worker/proxy/BUILD.bazel @@ -60,6 +60,7 @@ go_test( name = "proxy_test", srcs = [ "h3_addrlist_test.go", + "host_transport_test.go", "proxy_e2e_test.go", "proxy_extra_test.go", "proxy_test.go", @@ -82,6 +83,7 @@ go_test( "@com_github_nats_io_nats_go//jetstream", "@com_github_quic_go_quic_go//http3", "@com_github_quic_go_quic_go//integrationtests/tools", + "@com_github_stretchr_testify//assert", "@com_github_stretchr_testify//require", "@org_golang_google_protobuf//proto", "@org_golang_x_net//http2", diff --git a/src/libraries/go/worker/proxy/h3.go b/src/libraries/go/worker/proxy/h3.go index 8a41058b2..36e3eba59 100644 --- a/src/libraries/go/worker/proxy/h3.go +++ b/src/libraries/go/worker/proxy/h3.go @@ -66,12 +66,32 @@ func createH3RoundTripper() *h3ConnectionCache { return &h3ConnectionCache{wrappedTransport: h3, clients: make(map[string]*roundTripperWithCount)} } +// newHostTransport opens a UDP socket for one proxy host. +// +// A socket per host rather than one shared by all of them, because the source +// port is what a UDP load balancer hashes on. Sharing one socket means every +// proxy pod is reached over a single balancer flow, so if that flow is pinned +// to an instance that has gone away, QUIC to every pod fails at once and keeps +// failing for as long as traffic continues: the flow never idles out and +// therefore never re-hashes. Restarting the whole worker was the only way back, +// which is why replacing function instances has been the standard remedy. +// +// Per host, a failed host is discarded together with its socket, and the next +// attempt arrives from a new source port that the balancer is free to place on +// a healthy instance. It also isolates hosts from one another. +func newHostTransport() (*quic.Transport, error) { + udpConn, err := net.ListenUDP("udp", nil) + if err != nil { + return nil, err + } + return &quic.Transport{Conn: udpConn}, nil +} + // mostly copied from http3.Transport because we need to hijack the http3 client stream // and when doing that we can't use the built in http3.Transport.RoundTrip function which caches // connections. type h3ConnectionCache struct { wrappedTransport *http3.Transport - quicTransport *quic.Transport mutex sync.Mutex clients map[string]*roundTripperWithCount } @@ -103,9 +123,17 @@ func (t *h3ConnectionCache) getClient(ctx context.Context, hostname string) (rtc if !ok { ctx, cancel := context.WithCancel(ctx) removeOnce := sync.Once{} + var hostTransport *quic.Transport + if t.wrappedTransport.Dial == nil { + hostTransport, err = newHostTransport() + if err != nil { + return nil, false, err + } + } cl = &roundTripperWithCount{ - dialing: make(chan struct{}), - cancel: cancel, + dialing: make(chan struct{}), + cancel: cancel, + transport: hostTransport, // may be called multiple times if many callers detect failure removeFromCache: func() { removeOnce.Do(func() { @@ -116,7 +144,7 @@ func (t *h3ConnectionCache) getClient(ctx context.Context, hostname string) (rtc go func() { defer close(cl.dialing) defer cancel() - conn, rt, err := t.dial(ctx, hostname) + conn, rt, err := t.dial(ctx, hostname, cl) if err != nil { cl.dialErr = err return @@ -149,7 +177,7 @@ func (t *h3ConnectionCache) getClient(ctx context.Context, hostname string) (rtc return cl, isReused, nil } -func (t *h3ConnectionCache) dial(ctx context.Context, hostname string) (*quic.Conn, *http3.ClientConn, error) { +func (t *h3ConnectionCache) dial(ctx context.Context, hostname string, cl *roundTripperWithCount) (*quic.Conn, *http3.ClientConn, error) { var tlsConf *tls.Config if t.wrappedTransport.TLSClientConfig == nil { tlsConf = &tls.Config{} @@ -169,24 +197,20 @@ func (t *h3ConnectionCache) dial(ctx context.Context, hostname string) (*quic.Co dial := t.wrappedTransport.Dial if dial == nil { - if t.quicTransport == nil { - udpConn, err := net.ListenUDP("udp", nil) - if err != nil { - return nil, nil, err - } - t.quicTransport = &quic.Transport{Conn: udpConn} - } dial = func(ctx context.Context, addr string, tlsCfg *tls.Config, cfg *quic.Config) (*quic.Conn, error) { network := "udp" udpAddr, err := t.resolveUDPAddr(ctx, network, addr) if err != nil { return nil, err } - conn, err := t.quicTransport.DialEarly(ctx, udpAddr, tlsCfg, cfg) + conn, err := cl.transport.DialEarly(ctx, udpAddr, tlsCfg, cfg) return conn, err } } - conn, err := dial(ctx, hostname, tlsConf, t.wrappedTransport.QUICConfig) + // Per-dial copy: quic-go's validateConfig writes defaults back into the + // Config it is handed, so concurrent dials sharing one pointer race on it. + quicConf := *t.wrappedTransport.QUICConfig + conn, err := dial(ctx, hostname, tlsConf, &quicConf) if err != nil { zap.L().Warn("failed to dial quic connection", zap.Error(err), zap.String("hostname", hostname)) return nil, nil, err @@ -233,23 +257,19 @@ func (t *h3ConnectionCache) Close() error { } } t.clients = nil - if t.quicTransport != nil { - if err := t.quicTransport.Close(); err != nil { - return err - } - if err := t.quicTransport.Conn.Close(); err != nil { - return err - } - t.quicTransport = nil - } return nil } type roundTripperWithCount struct { - cancel context.CancelFunc - dialing chan struct{} // closed as soon as quic.Dial(Early) returned - dialErr error - conn *quic.Conn + cancel context.CancelFunc + dialing chan struct{} // closed as soon as quic.Dial(Early) returned + dialErr error + conn *quic.Conn + // transport owns this host's UDP socket. Held per host rather than shared + // so that one unreachable proxy pod cannot take QUIC to every other pod + // down with it, and so that discarding a failed host also discards the + // source port it was using. See the note on newHostTransport. + transport *quic.Transport clientConn *http3.ClientConn useCount atomic.Int64 @@ -259,10 +279,20 @@ type roundTripperWithCount struct { func (r *roundTripperWithCount) Close() error { r.cancel() <-r.dialing + var connErr error if r.conn != nil { - return r.conn.CloseWithError(0, "") + connErr = r.conn.CloseWithError(0, "") } - return nil + // Closing the socket is the point, not just tidiness: a new one is opened + // from a different source port, which is what lets a load balancer that + // pinned this flow to a dead target route the next attempt somewhere else. + // Deliberately not set to nil: the field is written once at construction so + // that it can be read without synchronisation from the dial goroutine. + if r.transport != nil { + _ = r.transport.Close() + _ = r.transport.Conn.Close() + } + return connErr } // An addrList represents a list of network endpoint addresses. diff --git a/src/libraries/go/worker/proxy/host_transport_test.go b/src/libraries/go/worker/proxy/host_transport_test.go new file mode 100644 index 000000000..9363669c6 --- /dev/null +++ b/src/libraries/go/worker/proxy/host_transport_test.go @@ -0,0 +1,186 @@ +/* +SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +SPDX-License-Identifier: Apache-2.0 + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +*/ + +package proxy + +import ( + "context" + "net" + "testing" + "time" + + "github.com/google/uuid" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + + pb "github.com/NVIDIA/nvcf/src/libraries/go/worker/proto/nvcf" +) + +// deadUDPHost returns the address of a socket that reads packets and never +// answers, which is how a proxy pod that has gone away behaves from the dialler's +// point of view: the handshake has to time out rather than being refused. +func deadUDPHost(t *testing.T) string { + t.Helper() + pc, err := net.ListenPacket("udp", "127.0.0.1:0") + require.NoError(t, err) + t.Cleanup(func() { _ = pc.Close() }) + go func() { + buf := make([]byte, 1500) + for { + if _, _, err := pc.ReadFrom(buf); err != nil { + return + } + } + }() + return pc.LocalAddr().String() +} + +func h3ConnConfig(addr string) *pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig_HTTP3ConnectionConfig { + return &pb.WorkerInvokeFunctionRequest_StatefulConfig_ConnectionConfig_HTTP3ConnectionConfig{ + ProxyURI: "https://" + addr + "/v1/proxy", + ProxyAuthorizationToken: "dummy-token", + } +} + +// localPortFor reports the source port the cache is currently using for a host, +// or "" when it holds no socket for it. +func localPortFor(t *testing.T, cache *h3ConnectionCache, hostname string) string { + t.Helper() + cache.mutex.Lock() + defer cache.mutex.Unlock() + cl, ok := cache.clients[hostname] + if !ok || cl.transport == nil { + return "" + } + _, port, err := net.SplitHostPort(cl.transport.Conn.LocalAddr().String()) + require.NoError(t, err) + return port +} + +// The reason this change exists. A UDP load balancer hashes on the source port, +// so a dead flow only moves to a healthy instance when the source port changes. +// Discarding a failed host has to discard its socket, otherwise every retry +// leaves from the same port, lands on the same dead instance, and the worker +// never recovers while traffic continues. +func TestFailedHostGetsANewSourcePortOnRetry(t *testing.T) { + setupLogger() + allowInsecure(t) + addr := deadUDPHost(t) + hostname := addr + + cache := createH3RoundTripper() + requestId := uuid.New().String() + + _, err := quicConnect(context.Background(), requestId, h3ConnConfig(addr), cache) + require.Error(t, err, "dial to a host that never answers must fail") + + // The failed host must not still be holding its socket. + assert.Empty(t, localPortFor(t, cache, hostname), + "a failed host should have been discarded together with its socket") + + // Dial again and capture the port actually in use during the attempt. + seen := make(chan string, 1) + go func() { + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if p := localPortFor(t, cache, hostname); p != "" { + seen <- p + return + } + time.Sleep(5 * time.Millisecond) + } + seen <- "" + }() + _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(addr), cache) + + second := <-seen + require.NotEmpty(t, second, "second attempt should have opened a socket") + t.Logf("second attempt used source port %s", second) +} + +// Hosts must not share a socket, otherwise one unreachable pod takes QUIC to +// every other pod down with it. +func TestDistinctHostsUseDistinctSourcePorts(t *testing.T) { + setupLogger() + allowInsecure(t) + hostA := deadUDPHost(t) + hostB := deadUDPHost(t) + require.NotEqual(t, hostA, hostB) + + cache := createH3RoundTripper() + + ports := make(chan string, 2) + for _, h := range []string{hostA, hostB} { + go func(h string) { + deadline := time.Now().Add(3 * time.Second) + for time.Now().Before(deadline) { + if p := localPortFor(t, cache, h); p != "" { + ports <- p + return + } + time.Sleep(5 * time.Millisecond) + } + ports <- "" + }(h) + } + go func() { _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(hostA), cache) }() + _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(hostB), cache) + + p1, p2 := <-ports, <-ports + require.NotEmpty(t, p1) + require.NotEmpty(t, p2) + assert.NotEqual(t, p1, p2, "each host must have its own source port, got %s and %s", p1, p2) +} + +// Closing the cache must not leak the per-host sockets. +func TestCloseReleasesHostSockets(t *testing.T) { + setupLogger() + allowInsecure(t) + addr := deadUDPHost(t) + + cache := createH3RoundTripper() + go func() { _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(addr), cache) }() + + require.Eventually(t, func() bool { return localPortFor(t, cache, addr) != "" }, + 3*time.Second, 5*time.Millisecond, "a socket should have been opened for the host") + + require.NoError(t, cache.Close()) + + cache.mutex.Lock() + clients := cache.clients + cache.mutex.Unlock() + assert.Nil(t, clients, "Close should have released the client map") +} + +// A worker holds one socket per proxy pod it talks to, so the count has to track +// distinct hosts rather than growing per dial attempt. +func TestSocketCountTracksDistinctHosts(t *testing.T) { + setupLogger() + allowInsecure(t) + + cache := createH3RoundTripper() + addr := deadUDPHost(t) + + for i := 0; i < 3; i++ { + _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(addr), cache) + } + + cache.mutex.Lock() + n := len(cache.clients) + cache.mutex.Unlock() + assert.LessOrEqual(t, n, 1, "repeated attempts to one host must not accumulate entries, got %d", n) +} From 46dda2ec4886a538bc992aa2879fe2c658fa8cbd Mon Sep 17 00:00:00 2001 From: balaji Date: Thu, 20 Aug 2026 12:32:29 -0700 Subject: [PATCH 2/2] fix(worker): release the host socket on every path that drops the host Review caught that giving each host its own socket leaked one on every failed dial. None of the removal paths closed it: removeClient and the dialErr branch in getClient both dropped the entry from the map and nothing else, and the connection-context callback goes through the same remove-only path. That is worse than the shared socket it replaced. A dial storm leaked a file descriptor per failure, and the staging reproduction produced 3,796 failures in eight minutes. Release the socket from every path that drops a host, exactly once via a sync.Once, since several paths can race to drop the same entry. The release deliberately does not wait on the dial to finish: callers hold the cache lock, and a dial still in flight belongs to a host being discarded anyway. The tests asserted on cache-map state, which could not tell a released socket from a leaked one, and could have passed while a retry reused the old source port. They now bind the old port to prove it was released, hold it so a retry cannot reuse it by chance, and assert the retry leaves from a different port. Verified they fail with the leak reintroduced. Co-Authored-By: Balaji Ganesan --- src/libraries/go/worker/proxy/h3.go | 38 ++++- .../go/worker/proxy/host_transport_test.go | 147 ++++++++++-------- 2 files changed, 113 insertions(+), 72 deletions(-) diff --git a/src/libraries/go/worker/proxy/h3.go b/src/libraries/go/worker/proxy/h3.go index 36e3eba59..f0f370b3a 100644 --- a/src/libraries/go/worker/proxy/h3.go +++ b/src/libraries/go/worker/proxy/h3.go @@ -163,6 +163,7 @@ func (t *h3ConnectionCache) getClient(ctx context.Context, hostname string) (rtc case <-cl.dialing: if cl.dialErr != nil { delete(t.clients, hostname) + cl.closeTransport() return nil, false, cl.dialErr } select { @@ -240,11 +241,19 @@ func (t *h3ConnectionCache) resolveUDPAddr(ctx context.Context, network, addr st func (t *h3ConnectionCache) removeClient(hostname string) { t.mutex.Lock() - defer t.mutex.Unlock() if t.clients == nil { + t.mutex.Unlock() return } + cl := t.clients[hostname] delete(t.clients, hostname) + t.mutex.Unlock() + + // Dropping the host has to drop its socket too, otherwise the source port + // is never released and every retry leaks one. + if cl != nil { + cl.closeTransport() + } } // Close closes the QUIC connections that this Transport has used. @@ -269,8 +278,9 @@ type roundTripperWithCount struct { // so that one unreachable proxy pod cannot take QUIC to every other pod // down with it, and so that discarding a failed host also discards the // source port it was using. See the note on newHostTransport. - transport *quic.Transport - clientConn *http3.ClientConn + transport *quic.Transport + transportOnce sync.Once + clientConn *http3.ClientConn useCount atomic.Int64 removeFromCache func() @@ -286,13 +296,25 @@ func (r *roundTripperWithCount) Close() error { // Closing the socket is the point, not just tidiness: a new one is opened // from a different source port, which is what lets a load balancer that // pinned this flow to a dead target route the next attempt somewhere else. - // Deliberately not set to nil: the field is written once at construction so - // that it can be read without synchronisation from the dial goroutine. - if r.transport != nil { + r.closeTransport() + return connErr +} + +// closeTransport releases this host's UDP socket. Idempotent, because a host is +// dropped from the cache by several paths and the socket must be released +// exactly once by whichever gets there first. Leaking it would be worse than +// the shared socket this replaced: a dial storm would leak one per failure. +// +// It deliberately does not wait on dialing. Callers hold the cache lock, and a +// dial still in flight is one whose host is being discarded anyway. +func (r *roundTripperWithCount) closeTransport() { + r.transportOnce.Do(func() { + if r.transport == nil { + return + } _ = r.transport.Close() _ = r.transport.Conn.Close() - } - return connErr + }) } // An addrList represents a list of network endpoint addresses. diff --git a/src/libraries/go/worker/proxy/host_transport_test.go b/src/libraries/go/worker/proxy/host_transport_test.go index 9363669c6..955db4a01 100644 --- a/src/libraries/go/worker/proxy/host_transport_test.go +++ b/src/libraries/go/worker/proxy/host_transport_test.go @@ -31,8 +31,8 @@ import ( ) // deadUDPHost returns the address of a socket that reads packets and never -// answers, which is how a proxy pod that has gone away behaves from the dialler's -// point of view: the handshake has to time out rather than being refused. +// answers, which is how a proxy pod that has gone away behaves from the +// dialler's point of view: the handshake times out rather than being refused. func deadUDPHost(t *testing.T) string { t.Helper() pc, err := net.ListenPacket("udp", "127.0.0.1:0") @@ -56,9 +56,9 @@ func h3ConnConfig(addr string) *pb.WorkerInvokeFunctionRequest_StatefulConfig_Co } } -// localPortFor reports the source port the cache is currently using for a host, -// or "" when it holds no socket for it. -func localPortFor(t *testing.T, cache *h3ConnectionCache, hostname string) string { +// portInUse reports the source port the cache currently holds for a host, or "" +// if it holds no socket for it. +func portInUse(t *testing.T, cache *h3ConnectionCache, hostname string) string { t.Helper() cache.mutex.Lock() defer cache.mutex.Unlock() @@ -71,45 +71,75 @@ func localPortFor(t *testing.T, cache *h3ConnectionCache, hostname string) strin return port } -// The reason this change exists. A UDP load balancer hashes on the source port, -// so a dead flow only moves to a healthy instance when the source port changes. -// Discarding a failed host has to discard its socket, otherwise every retry -// leaves from the same port, lands on the same dead instance, and the worker -// never recovers while traffic continues. -func TestFailedHostGetsANewSourcePortOnRetry(t *testing.T) { +// capturePort watches for the socket a dial opens and reports its port. +func capturePort(t *testing.T, cache *h3ConnectionCache, hostname string) <-chan string { + t.Helper() + out := make(chan string, 1) + go func() { + deadline := time.Now().Add(5 * time.Second) + for time.Now().Before(deadline) { + if p := portInUse(t, cache, hostname); p != "" { + out <- p + return + } + time.Sleep(2 * time.Millisecond) + } + out <- "" + }() + return out +} + +// canBindUDP reports whether a UDP socket can be bound to a port, which is only +// true once whoever held it has actually released it. This is what proves the +// socket was closed rather than merely dropped from a map. +func canBindUDP(t *testing.T, port string) bool { + t.Helper() + pc, err := net.ListenPacket("udp", "127.0.0.1:"+port) + if err != nil { + return false + } + _ = pc.Close() + return true +} + +// The property the change exists for. A UDP load balancer hashes on the source +// port, so a dead flow only moves to a healthy instance when the port changes. +// Asserting on cache state alone would not prove that, so this checks the +// socket is genuinely released and that the retry leaves from a different port. +func TestFailedDialReleasesSocketAndRetryUsesANewPort(t *testing.T) { setupLogger() allowInsecure(t) addr := deadUDPHost(t) - hostname := addr cache := createH3RoundTripper() - requestId := uuid.New().String() + t.Cleanup(func() { _ = cache.Close() }) - _, err := quicConnect(context.Background(), requestId, h3ConnConfig(addr), cache) - require.Error(t, err, "dial to a host that never answers must fail") + first := capturePort(t, cache, addr) + _, err := quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(addr), cache) + require.Error(t, err, "a dial to a host that never answers must fail") - // The failed host must not still be holding its socket. - assert.Empty(t, localPortFor(t, cache, hostname), - "a failed host should have been discarded together with its socket") + firstPort := <-first + require.NotEmpty(t, firstPort, "the first attempt should have opened a socket") - // Dial again and capture the port actually in use during the attempt. - seen := make(chan string, 1) - go func() { - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - if p := localPortFor(t, cache, hostname); p != "" { - seen <- p - return - } - time.Sleep(5 * time.Millisecond) - } - seen <- "" - }() + // The socket must actually be released, not just forgotten. If it were + // leaked, a dial storm would leak one per failure. + require.Eventually(t, func() bool { return canBindUDP(t, firstPort) }, 3*time.Second, 10*time.Millisecond, + "source port %s was never released after the dial failed", firstPort) + + // Hold the old port so the retry cannot reuse it by chance, which would + // make the assertion below pass for the wrong reason. + guard, err := net.ListenPacket("udp", "127.0.0.1:"+firstPort) + require.NoError(t, err) + defer func() { _ = guard.Close() }() + + second := capturePort(t, cache, addr) _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(addr), cache) + secondPort := <-second - second := <-seen - require.NotEmpty(t, second, "second attempt should have opened a socket") - t.Logf("second attempt used source port %s", second) + require.NotEmpty(t, secondPort, "the retry should have opened a socket") + assert.NotEqual(t, firstPort, secondPort, + "the retry must leave from a new source port so the balancer can re-place the flow") + t.Logf("first attempt port=%s, retry port=%s", firstPort, secondPort) } // Hosts must not share a socket, otherwise one unreachable pod takes QUIC to @@ -122,58 +152,47 @@ func TestDistinctHostsUseDistinctSourcePorts(t *testing.T) { require.NotEqual(t, hostA, hostB) cache := createH3RoundTripper() + t.Cleanup(func() { _ = cache.Close() }) - ports := make(chan string, 2) - for _, h := range []string{hostA, hostB} { - go func(h string) { - deadline := time.Now().Add(3 * time.Second) - for time.Now().Before(deadline) { - if p := localPortFor(t, cache, h); p != "" { - ports <- p - return - } - time.Sleep(5 * time.Millisecond) - } - ports <- "" - }(h) - } + watchA := capturePort(t, cache, hostA) + watchB := capturePort(t, cache, hostB) go func() { _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(hostA), cache) }() _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(hostB), cache) - p1, p2 := <-ports, <-ports - require.NotEmpty(t, p1) - require.NotEmpty(t, p2) - assert.NotEqual(t, p1, p2, "each host must have its own source port, got %s and %s", p1, p2) + portA, portB := <-watchA, <-watchB + require.NotEmpty(t, portA) + require.NotEmpty(t, portB) + assert.NotEqual(t, portA, portB, "each host must have its own source port, got %s and %s", portA, portB) } -// Closing the cache must not leak the per-host sockets. +// Closing the cache must release the sockets, not just drop the map. func TestCloseReleasesHostSockets(t *testing.T) { setupLogger() allowInsecure(t) addr := deadUDPHost(t) cache := createH3RoundTripper() + watch := capturePort(t, cache, addr) go func() { _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(addr), cache) }() - require.Eventually(t, func() bool { return localPortFor(t, cache, addr) != "" }, - 3*time.Second, 5*time.Millisecond, "a socket should have been opened for the host") + port := <-watch + require.NotEmpty(t, port, "a socket should have been opened for the host") require.NoError(t, cache.Close()) - cache.mutex.Lock() - clients := cache.clients - cache.mutex.Unlock() - assert.Nil(t, clients, "Close should have released the client map") + assert.Eventually(t, func() bool { return canBindUDP(t, port) }, 3*time.Second, 10*time.Millisecond, + "Close left source port %s bound", port) } -// A worker holds one socket per proxy pod it talks to, so the count has to track -// distinct hosts rather than growing per dial attempt. -func TestSocketCountTracksDistinctHosts(t *testing.T) { +// Failed dials must not accumulate cache entries, which is the other half of +// not accumulating sockets. +func TestFailedDialsLeaveNoCacheEntries(t *testing.T) { setupLogger() allowInsecure(t) + addr := deadUDPHost(t) cache := createH3RoundTripper() - addr := deadUDPHost(t) + t.Cleanup(func() { _ = cache.Close() }) for i := 0; i < 3; i++ { _, _ = quicConnect(context.Background(), uuid.New().String(), h3ConnConfig(addr), cache) @@ -182,5 +201,5 @@ func TestSocketCountTracksDistinctHosts(t *testing.T) { cache.mutex.Lock() n := len(cache.clients) cache.mutex.Unlock() - assert.LessOrEqual(t, n, 1, "repeated attempts to one host must not accumulate entries, got %d", n) + assert.Zero(t, n, "failed dials should leave no entries behind, found %d", n) }