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..f0f370b3a 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 @@ -135,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 { @@ -149,7 +178,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 +198,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 @@ -216,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. @@ -233,24 +266,21 @@ 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 - clientConn *http3.ClientConn + 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 + transportOnce sync.Once + clientConn *http3.ClientConn useCount atomic.Int64 removeFromCache func() @@ -259,10 +289,32 @@ 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. + 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() + }) } // 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..955db4a01 --- /dev/null +++ b/src/libraries/go/worker/proxy/host_transport_test.go @@ -0,0 +1,205 @@ +/* +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 times 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", + } +} + +// 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() + 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 +} + +// 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) + + cache := createH3RoundTripper() + t.Cleanup(func() { _ = cache.Close() }) + + 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") + + firstPort := <-first + require.NotEmpty(t, firstPort, "the first attempt should have opened a socket") + + // 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 + + 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 +// 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() + t.Cleanup(func() { _ = cache.Close() }) + + 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) + + 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 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) }() + + port := <-watch + require.NotEmpty(t, port, "a socket should have been opened for the host") + + require.NoError(t, cache.Close()) + + assert.Eventually(t, func() bool { return canBindUDP(t, port) }, 3*time.Second, 10*time.Millisecond, + "Close left source port %s bound", port) +} + +// 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() + t.Cleanup(func() { _ = cache.Close() }) + + 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.Zero(t, n, "failed dials should leave no entries behind, found %d", n) +}