router: skip preflight dial when no retries are configured

pkg/router/handlers.go

@@ -238,6 +238,13 @@ func (s *Server) waitForUpstreamReachable(ctx context.Context, targetURL *url.UR
 		}
 	}
 
+	// Without retries the preflight is one more connect to the sandbox that
+	// the proxy repeats right after. The proxy transport has the same dial
+	// timeout (see NewServer), so it reports the same errors.
+	if retryCount == 0 {
+		return nil
+	}
+
 	var lastErr error
 	for attempt := 0; attempt <= retryCount; attempt++ {
 		dialer := &net.Dialer{Timeout: retryInterval}

pkg/router/handlers_bench_test.go

@@ -0,0 +1,96 @@
+/*
+Copyright The Volcano Authors.
+
+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 router
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+
+	"github.com/gin-gonic/gin"
+	"k8s.io/klog/v2"
+
+	"github.com/volcano-sh/agentcube/pkg/common/types"
+)
+
+// closeNotifyRecorder adds http.CloseNotifier, which httputil.ReverseProxy
+// asks for and httptest.ResponseRecorder does not implement.
+type closeNotifyRecorder struct {
+	*httptest.ResponseRecorder
+}
+
+func (closeNotifyRecorder) CloseNotify() <-chan bool { return make(chan bool) }
+
+// newInvokeBenchServer builds a Server whose invoke path is the real one
+// (handleInvoke -> forwardToSandbox) with an in-process upstream sandbox.
+func newInvokeBenchServer(b *testing.B, kind string) (*Server, string) {
+	b.Helper()
+	gin.DefaultWriter = io.Discard
+	klog.LogToStderr(false)
+	klog.SetOutput(io.Discard)
+	upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		_, _ = w.Write([]byte("ok"))
+	}))
+	b.Cleanup(upstream.Close)
+
+	jwtMgr, err := NewJWTManager()
+	if err != nil {
+		b.Fatal(err)
+	}
+	transport := &http.Transport{}
+	b.Cleanup(transport.CloseIdleConnections)
+
+	sandbox := &types.SandboxInfo{
+		Kind:      kind,
+		SessionID: "session-1",
+		Name:      "sb",
+		EntryPoints: []types.SandboxEntryPoint{{
+			Path:     "/",
+			Protocol: "HTTP",
+			Endpoint: upstream.Listener.Addr().String(),
+		}},
+	}
+	s := &Server{
+		config:         &Config{MaxConcurrentRequests: 1 << 20},
+		sessionManager: &mockSessionManager{sandbox: sandbox},
+		storeClient:    &fakeStoreClient{},
+		httpTransport:  transport,
+		jwtManager:     jwtMgr,
+	}
+	s.setupRoutes()
+	return s, "/v1/namespaces/default/code-interpreters/ci/invocations/api/ping"
+}
+
+func BenchmarkInvoke(b *testing.B) {
+	for _, kind := range []string{types.SandboxClaimsKind, types.CodeInterpreterKind} {
+		b.Run(kind, func(b *testing.B) {
+			s, path := newInvokeBenchServer(b, kind)
+			b.ReportAllocs()
+			b.ResetTimer()
+			for i := 0; i < b.N; i++ {
+				req := httptest.NewRequest(http.MethodGet, path, nil)
+				req.Header.Set("x-agentcube-session-id", "session-1")
+				w := closeNotifyRecorder{httptest.NewRecorder()}
+				s.engine.ServeHTTP(w, req)
+				if w.Code != http.StatusOK {
+					b.Fatalf("status %d: %s", w.Code, w.Body.String())
+				}
+			}
+		})
+	}
+}

pkg/router/preflight_test.go

@@ -0,0 +1,92 @@
+/*
+Copyright The Volcano Authors.
+
+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 router
+
+import (
+	"context"
+	"net"
+	"net/http"
+	"net/http/httptest"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/volcano-sh/agentcube/pkg/common/types"
+)
+
+func TestWaitForUpstreamReachableWithoutRetriesDoesNotDial(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	defer func() { _ = listener.Close() }()
+
+	var accepted atomic.Int32
+	go func() {
+		for {
+			conn, acceptErr := listener.Accept()
+			if acceptErr != nil {
+				return
+			}
+			accepted.Add(1)
+			_ = conn.Close()
+		}
+	}()
+
+	targetURL, err := buildURL("HTTP", listener.Addr().String())
+	require.NoError(t, err)
+
+	s := &Server{config: &Config{InitialConnectRetryCount: 0, InitialConnectRetryInterval: 100 * time.Millisecond}}
+	require.NoError(t, s.waitForUpstreamReachable(context.Background(), targetURL))
+
+	time.Sleep(50 * time.Millisecond)
+	assert.Zero(t, accepted.Load(), "preflight should not open a connection when no retries are configured")
+}
+
+func TestForwardToSandboxRefusedUpstreamReturnsBadGateway(t *testing.T) {
+	listener, err := net.Listen("tcp", "127.0.0.1:0")
+	require.NoError(t, err)
+	address := listener.Addr().String()
+	require.NoError(t, listener.Close())
+
+	transport := &http.Transport{DialContext: (&net.Dialer{Timeout: 200 * time.Millisecond}).DialContext}
+	defer transport.CloseIdleConnections()
+	s := &Server{
+		config:        &Config{MaxConcurrentRequests: 10, InitialConnectRetryCount: 0, InitialConnectRetryInterval: 200 * time.Millisecond},
+		httpTransport: transport,
+		sessionManager: &mockSessionManager{sandbox: &types.SandboxInfo{
+			Kind:        types.CodeInterpreterKind,
+			SessionID:   "s1",
+			EntryPoints: []types.SandboxEntryPoint{{Path: "/", Protocol: "HTTP", Endpoint: address}},
+		}},
+		storeClient: &fakeStoreClient{},
+	}
+	s.setupRoutes()
+
+	routerServer := httptest.NewServer(s.engine)
+	defer routerServer.Close()
+
+	req, err := http.NewRequest(http.MethodGet, routerServer.URL+"/v1/namespaces/default/code-interpreters/ci/invocations/x", nil)
+	require.NoError(t, err)
+	req.Header.Set("x-agentcube-session-id", "s1")
+	resp, err := (&http.Client{Timeout: 5 * time.Second}).Do(req)
+	require.NoError(t, err)
+	defer resp.Body.Close()
+
+	assert.Equal(t, http.StatusBadGateway, resp.StatusCode)
+}

pkg/router/server.go

@@ -19,6 +19,7 @@ package router
 import (
 	"context"
 	"fmt"
+	"net"
 	"net/http"
 	"time"
 
@@ -86,6 +87,9 @@ func NewServer(config *Config) (*Server, error) {
 
 	// Create a reusable HTTP transport for connection pooling
 	httpTransport := &http.Transport{
+		// Fail fast on an unreachable sandbox instead of waiting for the OS
+		// connect timeout. The same bound applies to the preflight dial.
+		DialContext:        (&net.Dialer{Timeout: config.InitialConnectRetryInterval}).DialContext,
 		IdleConnTimeout:    0,
 		DisableCompression: false,
 	}