router: reuse signed sandbox JWTs for a session

pkg/router/handlers.go

@@ -287,18 +287,7 @@ func (s *Server) generateSandboxJWT(c *gin.Context, sandbox *types.SandboxInfo)
 		return "", true
 	}
 
-	claims := map[string]interface{}{
-		"session_id": sandbox.SessionID,
-	}
-
-	if oidcClaims := extractClaims(c); oidcClaims != nil {
-		claims["user_sub"] = oidcClaims.Subject
-		if oidcClaims.Email != "" {
-			claims["user_email"] = oidcClaims.Email
-		}
-	}
-
-	token, err := s.jwtManager.GenerateToken(claims)
+	token, err := s.jwtManager.GenerateSandboxToken(sandbox.SessionID, extractClaims(c))
 	if err != nil {
 		klog.Errorf("Failed to generate JWT token (session: %s): %v", sandbox.SessionID, err)
 		c.JSON(http.StatusInternalServerError, gin.H{

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/jwt.go

@@ -24,6 +24,7 @@ import (
 	"encoding/pem"
 	"fmt"
 	"os"
+	"sync"
 	"time"
 
 	"github.com/golang-jwt/jwt/v5"
@@ -33,13 +34,20 @@ import (
 	"k8s.io/client-go/kubernetes"
 	"k8s.io/client-go/rest"
 	"k8s.io/klog/v2"
+	"k8s.io/utils/lru"
 )
 
 const (
 	// RSA key size for JWT signing
 	rsaKeySize = 2048
 	// JWT token expiration time
 	jwtExpiration = 5 * time.Minute
+	// jwtReuseWindow is how long a cached token may be handed out again. It is
+	// half of jwtExpiration, so a reused token always has at least 2.5 minutes
+	// of validity left when the sandbox receives it.
+	jwtReuseWindow = jwtExpiration / 2
+	// jwtCacheSize bounds the number of cached tokens (one per session and caller).
+	jwtCacheSize = 4096
 	// IdentitySecretName is the name of the secret storing Router's keys
 	IdentitySecretName = "picod-router-identity" //nolint:gosec // This is a name reference, not a credential
 	// PrivateKeyDataKey is the key in the secret data map for private key
@@ -62,6 +70,19 @@ type JWTManager struct {
 	privateKey *rsa.PrivateKey
 	publicKey  *rsa.PublicKey
 	clientset  kubernetes.Interface
+
+	// tokenCache holds recently signed sandbox tokens keyed by session and
+	// caller, so that repeated requests on one session do not pay for an RSA
+	// signature each time.
+	tokenMu    sync.Mutex
+	tokenCache *lru.Cache
+}
+
+// cachedToken is a signed token together with the time after which it must
+// be signed again.
+type cachedToken struct {
+	token   string
+	reuseTo time.Time
 }
 
 // NewJWTManager creates a new JWT manager with a fresh RSA key pair
@@ -74,9 +95,64 @@ func NewJWTManager() (*JWTManager, error) {
 	return &JWTManager{
 		privateKey: privateKey,
 		publicKey:  &privateKey.PublicKey,
+		tokenCache: lru.New(jwtCacheSize),
 	}, nil
 }
 
+// sandboxTokenKey identifies the claims of a sandbox token. It is built only
+// by GenerateSandboxToken, together with the claims, so the two cannot drift.
+type sandboxTokenKey struct {
+	sessionID string
+	hasCaller bool
+	userSub   string
+	userEmail string
+}
+
+// GenerateSandboxToken returns the token the router sends to a sandbox for a
+// session and, when OIDC is on, the calling user. A token signed for the same
+// session and caller is reused while it is younger than jwtReuseWindow, because
+// signing with RSA-2048 costs roughly half a millisecond, several times the
+// rest of the router request path.
+func (jm *JWTManager) GenerateSandboxToken(sessionID string, caller *Claims) (string, error) {
+	key := sandboxTokenKey{sessionID: sessionID}
+	claims := map[string]interface{}{
+		"session_id": sessionID,
+	}
+	if caller != nil {
+		key.hasCaller = true
+		key.userSub = caller.Subject
+		claims["user_sub"] = caller.Subject
+		if caller.Email != "" {
+			key.userEmail = caller.Email
+			claims["user_email"] = caller.Email
+		}
+	}
+
+	now := time.Now()
+
+	jm.tokenMu.Lock()
+	if jm.tokenCache == nil {
+		jm.tokenCache = lru.New(jwtCacheSize)
+	}
+	if v, ok := jm.tokenCache.Get(key); ok {
+		if ct, isToken := v.(cachedToken); isToken && now.Before(ct.reuseTo) {
+			jm.tokenMu.Unlock()
+			return ct.token, nil
+		}
+	}
+	jm.tokenMu.Unlock()
+
+	token, err := jm.GenerateToken(claims)
+	if err != nil {
+		return "", err
+	}
+
+	jm.tokenMu.Lock()
+	jm.tokenCache.Add(key, cachedToken{token: token, reuseTo: now.Add(jwtReuseWindow)})
+	jm.tokenMu.Unlock()
+	return token, nil
+}
+
 // GenerateToken generates a JWT token with the given claims
 func (jm *JWTManager) GenerateToken(claims map[string]interface{}) (string, error) {
 	// Create JWT claims
@@ -218,5 +294,9 @@ func (jm *JWTManager) loadPrivateKeyPEM(privateKeyPEM []byte) error {
 
 	jm.privateKey = privateKey
 	jm.publicKey = &privateKey.PublicKey
+	// Tokens signed with the previous key would no longer verify.
+	jm.tokenMu.Lock()
+	jm.tokenCache = lru.New(jwtCacheSize)
+	jm.tokenMu.Unlock()
 	return nil
 }