flags.go

package featureflags

import (
	"context"
	"fmt"
	"strings"

	"github.com/launchdarkly/go-sdk-common/v3/ldcontext"
	"github.com/launchdarkly/go-sdk-common/v3/ldvalue"

	"github.com/e2b-dev/infra/packages/shared/pkg/env"
)

// kinds
const (
	SandboxKind                        ldcontext.Kind = "sandbox"
	SandboxTemplateAttribute           string         = "template-id"
	SandboxKernelVersionAttribute      string         = "kernel-version"
	SandboxFirecrackerVersionAttribute string         = "firecracker-version"

	TeamKind       ldcontext.Kind = "team"
	UserKind       ldcontext.Kind = "user"
	ClusterKind    ldcontext.Kind = "cluster"
	deploymentKind ldcontext.Kind = "deployment"
	TierKind       ldcontext.Kind = "tier"
	ServiceKind    ldcontext.Kind = "service"
	TemplateKind   ldcontext.Kind = "template"
	VolumeKind     ldcontext.Kind = "volume"

	OrchestratorKind            ldcontext.Kind = "orchestrator"
	OrchestratorCommitAttribute string         = "commit"
)

// All flags must be defined here: https://app.launchdarkly.com/projects/default/flags/

type JSONFlag struct {
	name     string
	fallback ldvalue.Value
}

func (f JSONFlag) Key() string {
	return f.name
}

func (f JSONFlag) String() string {
	return f.name
}

func (f JSONFlag) Fallback() ldvalue.Value {
	return f.fallback
}

func NewJSONFlag(name string, fallback ldvalue.Value) JSONFlag {
	flag := JSONFlag{name: name, fallback: fallback}
	builder := launchDarklyOfflineStore.Flag(flag.name).ValueForAll(fallback)
	launchDarklyOfflineStore.Update(builder)

	return flag
}

var CleanNFSCache = NewJSONFlag("clean-nfs-cache", ldvalue.Null())

// RateLimitConfigFlag provides per-team rate limit overrides.
// JSON format:
//
//	{
//	  "/sandboxes/": {"rate": 50, "burst": 100},
//	  "/sandboxes/:sandboxID/pause": {"rate": 10, "burst": 20}
//	}
//
// When non-null, values override the code defaults. Target specific teams in LaunchDarkly.
var RateLimitConfigFlag = NewJSONFlag("rate-limit-config", ldvalue.Null())

type BoolFlag struct {
	name     string
	fallback bool
}

func (f BoolFlag) Key() string {
	return f.name
}

func (f BoolFlag) String() string {
	return f.name
}

func (f BoolFlag) Fallback() bool {
	return f.fallback
}

func NewBoolFlag(name string, fallback bool) BoolFlag {
	flag := BoolFlag{name: name, fallback: fallback}
	builder := launchDarklyOfflineStore.Flag(flag.name).VariationForAll(fallback)
	launchDarklyOfflineStore.Update(builder)

	return flag
}

var (
	MetricsWriteFlag                    = NewBoolFlag("sandbox-metrics-write", true)
	MetricsReadFlag                     = NewBoolFlag("sandbox-metrics-read", true)
	SnapshotFeatureFlag                 = NewBoolFlag("use-nfs-for-snapshots", env.IsDevelopment())
	TemplateFeatureFlag                 = NewBoolFlag("use-nfs-for-templates", env.IsDevelopment())
	EnableWriteThroughCacheFlag         = NewBoolFlag("write-to-cache-on-writes", false)
	UseNFSCacheForBuildingTemplatesFlag = NewBoolFlag("use-nfs-for-building-templates", env.IsDevelopment())
	BestOfKCanFitFlag                   = NewBoolFlag("best-of-k-can-fit", true)
	BestOfKTooManyStartingFlag          = NewBoolFlag("best-of-k-too-many-starting", false)
	EdgeProvidedSandboxMetricsFlag      = NewBoolFlag("edge-provided-sandbox-metrics", false)
	CreateStorageCacheSpansFlag         = NewBoolFlag("create-storage-cache-spans", env.IsDevelopment())
	SandboxAutoResumeFlag               = NewBoolFlag("sandbox-auto-resume", env.IsDevelopment())
	OrchAcceptsCombinedHostFlag         = NewBoolFlag("orch-accepts-combined-host", false)

	// PeerToPeerChunkTransferFlag enables peer-to-peer chunk routing.
	PeerToPeerChunkTransferFlag = NewBoolFlag("peer-to-peer-chunk-transfer", false)
	// PeerToPeerAsyncCheckpointFlag makes Checkpoint upload fire-and-forget instead
	// of synchronous. Only safe to enable after PeerToPeerChunkTransferFlag is ON.
	PeerToPeerAsyncCheckpointFlag = NewBoolFlag("peer-to-peer-async-checkpoint", false)

	PersistentVolumesFlag            = NewBoolFlag("can-use-persistent-volumes", env.IsDevelopment())
	ExecutionMetricsOnWebhooksFlag   = NewBoolFlag("execution-metrics-on-webhooks", false) // TODO: Remove NLT 20250315
	SandboxLabelBasedSchedulingFlag  = NewBoolFlag("sandbox-label-based-scheduling", false)
	OptimisticResourceAccountingFlag = NewBoolFlag("sandbox-placement-optimistic-resource-accounting", false)
	FreePageReportingFlag            = NewBoolFlag("free-page-reporting", false)
)

type IntFlag struct {
	name     string
	fallback int
}

func (f IntFlag) Key() string {
	return f.name
}

func (f IntFlag) String() string {
	return f.name
}

func (f IntFlag) Fallback() int {
	return f.fallback
}

func NewIntFlag(name string, fallback int) IntFlag {
	flag := IntFlag{name: name, fallback: fallback}
	builder := launchDarklyOfflineStore.Flag(flag.name).ValueForAll(ldvalue.Int(fallback))
	launchDarklyOfflineStore.Update(builder)

	return flag
}

var (
	MaxSandboxesPerNode           = NewIntFlag("max-sandboxes-per-node", 200)
	GcloudConcurrentUploadLimit   = NewIntFlag("gcloud-concurrent-upload-limit", 8)
	GcloudMaxTasks                = NewIntFlag("gcloud-max-tasks", 16)
	ClickhouseBatcherMaxBatchSize = NewIntFlag("clickhouse-batcher-max-batch-size", 100)
	ClickhouseBatcherMaxDelay     = NewIntFlag("clickhouse-batcher-max-delay", 1000) // 1s in milliseconds
	ClickhouseBatcherQueueSize    = NewIntFlag("clickhouse-batcher-queue-size", 1000)
	BestOfKSampleSize             = NewIntFlag("best-of-k-sample-size", 3)                   // Default K=3
	BestOfKMaxOvercommit          = NewIntFlag("best-of-k-max-overcommit", 400)              // Default R=4 (stored as percentage, max over-commit ratio)
	BestOfKAlpha                  = NewIntFlag("best-of-k-alpha", 50)                        // Default Alpha=0.5 (stored as percentage for int flag, current usage weight)
	EnvdInitTimeoutMilliseconds   = NewIntFlag("envd-init-request-timeout-milliseconds", 50) // Timeout for envd init request in milliseconds
	HostStatsSamplingInterval     = NewIntFlag("host-stats-sampling-interval", 5000)         // Host stats sampling interval in milliseconds (default 5s)
	MaxCacheWriterConcurrencyFlag = NewIntFlag("max-cache-writer-concurrency", 10)

	// BuildCacheMaxUsagePercentage the maximum percentage of the cache disk storage
	// that can be used before the cache starts evicting items.
	BuildCacheMaxUsagePercentage = NewIntFlag("build-cache-max-usage-percentage", 85)
	BuildProvisionVersion        = NewIntFlag("build-provision-version", 0)

	// NBDConnectionsPerDevice the number of NBD socket connections per device
	NBDConnectionsPerDevice = NewIntFlag("nbd-connections-per-device", 1)

	// MemoryPrefetchMaxFetchWorkers is the maximum number of parallel fetch workers per sandbox for memory prefetching.
	// Fetching is I/O bound so we can have more parallelism.
	MemoryPrefetchMaxFetchWorkers = NewIntFlag("memory-prefetch-max-fetch-workers", 16)

	// MemoryPrefetchMaxCopyWorkers is the maximum number of parallel copy workers per sandbox for memory prefetching.
	// Copy uses uffd syscalls, so we limit parallelism to avoid overwhelming the system.
	MemoryPrefetchMaxCopyWorkers = NewIntFlag("memory-prefetch-max-copy-workers", 8)

	// TCPFirewallMaxConnectionsPerSandbox is the maximum number of concurrent TCP firewall
	// connections allowed per sandbox. Negative means no limit.
	TCPFirewallMaxConnectionsPerSandbox = NewIntFlag("tcpfirewall-max-connections-per-sandbox", -1)

	// SandboxMaxIncomingConnections is the maximum number of concurrent HTTP proxy
	// connections allowed per sandbox. Negative means no limit.
	SandboxMaxIncomingConnections = NewIntFlag("sandbox-max-incoming-connections", -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)

	// MinAutoResumeTimeoutSeconds is the minimum auto-resume timeout in seconds.
	// This prevents thrashing from very short timeouts.
	MinAutoResumeTimeoutSeconds = NewIntFlag("minimum-autoresume-timeout", 300)

	// BuildReservedDiskSpaceMB is the amount of disk space in MB reserved for root on the guest filesystem.
	// Reserved blocks are only usable by root (uid 0), protecting the guest OS from disk-full conditions.
	BuildReservedDiskSpaceMB = NewIntFlag("build-reserved-disk-space-mb", 0)

	// MaxConcurrentSnapshotUpserts limits concurrent UpsertSnapshot calls (pause + snapshot template paths).
	// 0 or negative disables throttling (unlimited concurrency).
	MaxConcurrentSnapshotUpserts = NewIntFlag("max-concurrent-snapshot-upserts", 0)
	// MaxConcurrentSandboxListQueries limits concurrent GetSnapshotsWithCursor calls in the sandbox list path.
	// 0 or negative disables throttling (unlimited concurrency).
	MaxConcurrentSandboxListQueries = NewIntFlag("max-concurrent-sandbox-list-queries", 0)
	// MaxConcurrentSnapshotBuildQueries limits concurrent GetSnapshotBuilds calls (e.g. sandbox delete).
	// 0 or negative disables throttling (unlimited concurrency).
	MaxConcurrentSnapshotBuildQueries = NewIntFlag("max-concurrent-snapshot-build-queries", 0)
)

type StringFlag struct {
	name     string
	fallback string
}

func (f StringFlag) Key() string {
	return f.name
}

func (f StringFlag) String() string {
	return f.name
}

func (f StringFlag) Fallback() string {
	return f.fallback
}

func NewStringFlag(name string, fallback string) StringFlag {
	flag := StringFlag{name: name, fallback: fallback}
	builder := launchDarklyOfflineStore.Flag(flag.name).ValueForAll(ldvalue.String(fallback))
	launchDarklyOfflineStore.Update(builder)

	return flag
}

// This is currently not configurable via feature flags.
const (
	DefaultKernelVersion = "vmlinux-6.1.158"
)

// The Firecracker version the last tag + the short SHA (so we can build our dev previews)
// TODO: The short tag here has only 7 characters — the one from our build pipeline will likely have exactly 8 so this will break.
const (
	DefaultFirecackerV1_10Version = "v1.10.1_30cbb07"
	DefaultFirecackerV1_12Version = "v1.12.1_210cbac"
	DefaultFirecackerV1_14Version = "v1.14.1_458ca91"
	DefaultFirecrackerVersion     = DefaultFirecackerV1_12Version
)

var FirecrackerVersionMap = map[string]string{
	"v1.10": DefaultFirecackerV1_10Version,
	"v1.12": DefaultFirecackerV1_12Version,
	"v1.14": DefaultFirecackerV1_14Version,
}

// BuildIoEngine Sync is used by default as there seems to be a bad interaction between Async and a lot of io operations.
var (
	BuildFirecrackerVersion     = NewStringFlag("build-firecracker-version", env.GetEnv("DEFAULT_FIRECRACKER_VERSION", DefaultFirecrackerVersion))
	BuildIoEngine               = NewStringFlag("build-io-engine", "Sync")
	DefaultPersistentVolumeType = NewStringFlag("default-persistent-volume-type", "")
	BuildNodeInfo               = NewJSONFlag("preferred-build-node", ldvalue.Null())
	FirecrackerVersions         = NewJSONFlag("firecracker-versions", ldvalue.FromJSONMarshal(FirecrackerVersionMap))
)

// ResolveFirecrackerVersion resolves the firecracker version using the FirecrackerVersions feature flag.
// The buildVersion format is "v1.12.1_210cbac" — we extract "v1.12" as the lookup key.
func ResolveFirecrackerVersion(ctx context.Context, ff *Client, buildVersion string) string {
	parts := strings.Split(buildVersion, "_")
	if len(parts) < 2 {
		return buildVersion
	}

	versionParts := strings.Split(strings.TrimPrefix(parts[0], "v"), ".")
	if len(versionParts) < 2 {
		return buildVersion
	}

	key := fmt.Sprintf("v%s.%s", versionParts[0], versionParts[1])
	versions := ff.JSONFlag(ctx, FirecrackerVersions).AsValueMap()

	if resolved, ok := versions.Get(key).AsOptionalString().Get(); ok {
		return resolved
	}

	return buildVersion
}

// defaultTrackedTemplates is the default map of template aliases tracked for metrics.
// This is used to reduce metric cardinality.
// JSON format: {"base": true, "code-interpreter-v1": true, ...}
var defaultTrackedTemplates = map[string]bool{
	"base":                  true,
	"code-interpreter-v1":   true,
	"code-interpreter-beta": true,
	"desktop":               true,
}

// TrackedTemplatesForMetrics is a JSON flag that defines which template aliases
// should be tracked in sandbox start time metrics. Templates not in this list
// will be grouped under "other" to reduce metric cardinality.
// JSON format: {"base": true, "code-interpreter-v1": true, ...}
var TrackedTemplatesForMetrics = NewJSONFlag("tracked-templates-for-metrics", ldvalue.FromJSONMarshal(defaultTrackedTemplates))

// GetTrackedTemplatesSet fetches the TrackedTemplatesForMetrics flag and returns it as a set for efficient lookup.
// Only keys with a truthy value are included; keys set to false are ignored.
func GetTrackedTemplatesSet(ctx context.Context, ff *Client) map[string]struct{} {
	value := ff.JSONFlag(ctx, TrackedTemplatesForMetrics)
	valueMap := value.AsValueMap()
	keys := valueMap.Keys(nil)
	result := make(map[string]struct{}, len(keys))
	for _, key := range keys {
		if valueMap.Get(key).BoolValue() {
			result[key] = struct{}{}
		}
	}

	return result
}

// ChunkerConfigFlag is a JSON flag controlling the chunker implementation and tuning.
//
// NOTE: Changing useStreaming has no effect on chunkers already created for
// cached templates. A service restart (redeploy) is required for that change
// to take effect. minReadBatchSizeKB is checked just-in-time on each fetch,
// so it takes effect immediately.
//
// JSON format: {"useStreaming": false, "minReadBatchSizeKB": 16}
var ChunkerConfigFlag = NewJSONFlag("chunker-config", ldvalue.FromJSONMarshal(map[string]any{
	"useStreaming":       false,
	"minReadBatchSizeKB": 16,
}))

// TCPFirewallEgressThrottleConfig controls per-sandbox egress throttling via Firecracker's
// VMM-level token bucket rate limiters on the network interface.
// Structure mirrors the Firecracker RateLimiter API: two independent token buckets.
// Set bucketSize to -1 to disable a bucket.
//
// Ops bucket (packets):    effective rate = ops.bucketSize * 1000 / ops.refillTimeMs ops/s.
// Bandwidth bucket (bytes): effective rate = bandwidth.bucketSize * 1000 / bandwidth.refillTimeMs bytes/s.
var TCPFirewallEgressThrottleConfig = NewJSONFlag("tcpfirewall-egress-throttle-config", ldvalue.FromJSONMarshal(map[string]any{
	"ops":       map[string]any{"bucketSize": -1, "oneTimeBurst": 0, "refillTimeMs": 1000},
	"bandwidth": map[string]any{"bucketSize": -1, "oneTimeBurst": 0, "refillTimeMs": 1000},
}))

// TokenBucketConfig holds parameters for a single Firecracker token bucket.
// BucketSize < 0 disables the bucket.
type TokenBucketConfig struct {
	BucketSize   int64
	OneTimeBurst int64
	RefillTimeMs int64
}

// TCPFirewallEgressThrottleConfigValue holds the parsed values of TCPFirewallEgressThrottleConfig.
type TCPFirewallEgressThrottleConfigValue struct {
	Ops       TokenBucketConfig
	Bandwidth TokenBucketConfig
}

// parseThrottleBuckets parses "ops" and "bandwidth" token bucket configs from a JSON flag value.
func parseThrottleBuckets(value ldvalue.Value) (ops, bandwidth TokenBucketConfig) {
	parseBucket := func(key string) TokenBucketConfig {
		b := value.GetByKey(key)
		if b.IsNull() {
			return TokenBucketConfig{BucketSize: -1} // disabled
		}

		// Validate refill time
		refillTimeMs := int64(b.GetByKey("refillTimeMs").IntValue())
		if refillTimeMs <= 0 {
			return TokenBucketConfig{BucketSize: -1} // disabled — invalid refill time
		}

		return TokenBucketConfig{
			BucketSize:   int64(b.GetByKey("bucketSize").IntValue()),
			OneTimeBurst: int64(b.GetByKey("oneTimeBurst").IntValue()),
			RefillTimeMs: refillTimeMs,
		}
	}

	return parseBucket("ops"), parseBucket("bandwidth")
}

// GetTCPFirewallEgressThrottleConfig fetches and parses the TCPFirewallEgressThrottleConfig flag.
func GetTCPFirewallEgressThrottleConfig(ctx context.Context, ff *Client) TCPFirewallEgressThrottleConfigValue {
	value := ff.JSONFlag(ctx, TCPFirewallEgressThrottleConfig)
	ops, bw := parseThrottleBuckets(value)

	return TCPFirewallEgressThrottleConfigValue{
		Ops:       ops,
		Bandwidth: bw,
	}
}

// BlockDriveThrottleConfig controls per-sandbox block device (disk) throttling via Firecracker's
// VMM-level token bucket rate limiters on the rootfs drive.
// Structure mirrors the Firecracker RateLimiter API: two independent token buckets.
// Set bucketSize to -1 to disable a bucket.
//
// Ops bucket (IOPS):       effective rate = ops.bucketSize * 1000 / ops.refillTimeMs ops/s.
// Bandwidth bucket (bytes): effective rate = bandwidth.bucketSize * 1000 / bandwidth.refillTimeMs bytes/s.
var BlockDriveThrottleConfig = NewJSONFlag("block-drive-throttle-config", ldvalue.FromJSONMarshal(map[string]any{
	"ops":       map[string]any{"bucketSize": -1, "oneTimeBurst": 0, "refillTimeMs": 1000},
	"bandwidth": map[string]any{"bucketSize": -1, "oneTimeBurst": 0, "refillTimeMs": 1000},
}))

// BlockDriveThrottleConfigValue holds the parsed values of BlockDriveThrottleConfig.
type BlockDriveThrottleConfigValue struct {
	Ops       TokenBucketConfig
	Bandwidth TokenBucketConfig
}

// GetBlockDriveThrottleConfig fetches and parses the BlockDriveThrottleConfig flag.
func GetBlockDriveThrottleConfig(ctx context.Context, ff *Client) BlockDriveThrottleConfigValue {
	value := ff.JSONFlag(ctx, BlockDriveThrottleConfig)
	ops, bw := parseThrottleBuckets(value)

	return BlockDriveThrottleConfigValue{
		Ops:       ops,
		Bandwidth: bw,
	}
}