diff --git a/packages/client-proxy/internal/proxy/proxy.go b/packages/client-proxy/internal/proxy/proxy.go index 9d36478edb..aa93be83c4 100644 --- a/packages/client-proxy/internal/proxy/proxy.go +++ b/packages/client-proxy/internal/proxy/proxy.go @@ -16,6 +16,7 @@ import ( "go.uber.org/zap" "google.golang.org/grpc/codes" + "github.com/e2b-dev/infra/packages/shared/pkg/connlimit" "github.com/e2b-dev/infra/packages/shared/pkg/consts" "github.com/e2b-dev/infra/packages/shared/pkg/featureflags" proxygrpc "github.com/e2b-dev/infra/packages/shared/pkg/grpc/proxy" @@ -241,6 +242,18 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port })) } + // Per-sandbox ingress limiter. The client-proxy shares one keepalive pool + // (ConnectionKey constant) across sandboxes, so the limiter is keyed by the + // Destination's LimitKey (the sandbox ID) instead, bounding per-sandbox + // in-flight requests without fragmenting the pool. Default flag -1 keeps the + // limiter inert until an operator sets a finite cap. + connLimitConfig := &reverseproxy.ConnectionLimitConfig{ + Limiter: connlimit.NewConnectionLimiter(), + GetMaxLimit: func(ctx context.Context) int { + return featureFlagsClient.IntFlag(ctx, featureflags.ClientProxyMaxConnectionsPerSandbox) + }, + } + getTargetFromRequest := reverseproxy.GetTargetFromRequest() forwardedHostDomains = normalizeForwardedHostDomains(forwardedHostDomains) proxy := reverseproxy.New( @@ -342,7 +355,11 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port RequestLogger: l, SandboxPort: port, ConnectionKey: connectionKey, - Url: url, + // Share the pool across sandboxes via ConnectionKey, but account + // the ingress limiter per sandbox via LimitKey, so one sandbox + // cannot occupy the shared edge without bound. + LimitKey: sandboxId, + Url: url, MaskRequestHost: clientProxyMaskRequestHost( ctx, featureFlagsClient, @@ -352,7 +369,7 @@ func NewClientProxy(meterProvider metric.MeterProvider, serviceName string, port ), }, nil }, - nil, + connLimitConfig, false, proxyOptions..., ) diff --git a/packages/shared/pkg/featureflags/flags.go b/packages/shared/pkg/featureflags/flags.go index 5cac0e21ad..d274f6ded4 100644 --- a/packages/shared/pkg/featureflags/flags.go +++ b/packages/shared/pkg/featureflags/flags.go @@ -887,6 +887,13 @@ var ( // connections allowed per sandbox. Negative means no limit. SandboxMaxIncomingConnections = NewIntFlag("sandbox-max-incoming-connections", -1) + // ClientProxyMaxConnectionsPerSandbox is the maximum number of concurrent + // client-proxy ingress requests accounted per sandbox. The client-proxy + // shares one keepalive connection pool across sandboxes, so without this the + // only bound on per-sandbox in-flight requests is the process-wide pool + // ceiling. Negative means no limit; the default preserves existing behavior. + ClientProxyMaxConnectionsPerSandbox = NewIntFlag("client-proxy-max-connections-per-sandbox", -1) + // BuildBaseRootfsSizeLimitMB is the maximum size of the base rootfs filesystem created from the OCI image, in MB. BuildBaseRootfsSizeLimitMB = NewIntFlag("build-base-rootfs-size-limit-mb", 25000) diff --git a/packages/shared/pkg/proxy/handler.go b/packages/shared/pkg/proxy/handler.go index 5bd1121218..bb34d02ab3 100644 --- a/packages/shared/pkg/proxy/handler.go +++ b/packages/shared/pkg/proxy/handler.go @@ -207,7 +207,7 @@ func handler(p *pool.ProxyPool, getDestination func(r *http.Request) (*pool.Dest // Connection limiting if connLimitConfig != nil { maxLimit := connLimitConfig.GetMaxLimit(ctx) - count, acquired := connLimitConfig.Limiter.TryAcquire(d.ConnectionKey, maxLimit) + count, acquired := connLimitConfig.Limiter.TryAcquire(d.LimiterKey(), maxLimit) if !acquired { l.Warn(ctx, "sandbox too many incoming connections", zap.String("host", r.Host), @@ -239,7 +239,7 @@ func handler(p *pool.ProxyPool, getDestination func(r *http.Request) (*pool.Dest start := time.Now() defer func() { - connLimitConfig.Limiter.Release(d.ConnectionKey) + connLimitConfig.Limiter.Release(d.LimiterKey()) if connLimitConfig.OnConnectionReleased != nil { connLimitConfig.OnConnectionReleased(ctx, time.Since(start).Milliseconds()) } diff --git a/packages/shared/pkg/proxy/pool/destination.go b/packages/shared/pkg/proxy/pool/destination.go index 9664e3ee47..40ffd68014 100644 --- a/packages/shared/pkg/proxy/pool/destination.go +++ b/packages/shared/pkg/proxy/pool/destination.go @@ -23,13 +23,17 @@ type Destination struct { // Should we return the error about closed port if there is a problem with a connection to upstream? DefaultToPortError bool RequestLogger logger.Logger - // ConnectionKey uniquely identifies a single sandbox lifecycle. It is - // used for two purposes: - // 1. keepalive connection pool isolation, so connections to a reused - // IP:port pair are not accidentally shared across sandboxes; - // 2. per-sandbox ingress connection limiter accounting. - // Embedders should pick a value that is unique per lifecycle. - ConnectionKey string + // ConnectionKey identifies the keepalive connection pool entry: connections + // to a reused IP:port pair are not accidentally shared across sandboxes. + // Embedders that isolate per lifecycle pick a value unique per lifecycle; + // embedders that intentionally share one pool (client-proxy) use a constant. + ConnectionKey string + // LimitKey scopes the ingress connection limiter accounting. When empty the + // limiter falls back to ConnectionKey (see LimiterKey), so a lifecycle-keyed + // pool is limited per lifecycle with no extra field. An embedder that shares + // one ConnectionKey across sandboxes (for pool reuse) sets LimitKey to the + // per-sandbox value so limiting stays per sandbox without fragmenting the pool. + LimitKey string IncludeSandboxIdInProxyErrorLogger bool // MaskRequestHost is used to mask the request host. MaskRequestHost *string @@ -41,3 +45,14 @@ type Destination struct { // ConnectionKey entry. Closing it ends them. Owner *Owner } + +// LimiterKey returns the key the ingress connection limiter accounts against: +// LimitKey when set, otherwise ConnectionKey. This lets an embedder share one +// ConnectionKey for pool reuse while still limiting per sandbox via LimitKey. +func (d Destination) LimiterKey() string { + if d.LimitKey != "" { + return d.LimitKey + } + + return d.ConnectionKey +} diff --git a/packages/shared/pkg/proxy/pool/destination_limiterkey_test.go b/packages/shared/pkg/proxy/pool/destination_limiterkey_test.go new file mode 100644 index 0000000000..f07772b78c --- /dev/null +++ b/packages/shared/pkg/proxy/pool/destination_limiterkey_test.go @@ -0,0 +1,28 @@ +package pool + +import ( + "testing" + + "github.com/stretchr/testify/assert" +) + +// TestDestination_LimiterKey verifies the limiter key is decoupled from the +// pool key: LimitKey when set, else ConnectionKey. This is what lets the +// client-proxy share one ConnectionKey (pool reuse) while limiting per sandbox. +func TestDestination_LimiterKey(t *testing.T) { + t.Parallel() + + // No LimitKey: falls back to ConnectionKey (orchestrator behavior — limiter + // scoped per lifecycle, identical to before this field existed). + lifecycle := Destination{ConnectionKey: "lifecycle-abc"} + assert.Equal(t, "lifecycle-abc", lifecycle.LimiterKey()) + + // LimitKey set: limiter is keyed per sandbox while the pool key stays shared + // (client-proxy behavior — one pool, per-sandbox limit). + shared := Destination{ConnectionKey: "client-proxy", LimitKey: "sbx-123"} + assert.Equal(t, "sbx-123", shared.LimiterKey()) + assert.Equal(t, "client-proxy", shared.ConnectionKey, "pool key must stay shared, not fragmented per sandbox") + + // Empty both: empty (admission with an empty key is the caller's concern). + assert.Equal(t, "", Destination{}.LimiterKey()) +}