package main import ( "bytes" "context" "crypto/tls" "database/sql" "encoding/base64" "encoding/binary" "errors" "flag" "fmt" "io" "net" "net/http" "os" "os/signal" "runtime" "slices" "strconv" "strings" "sync" "syscall" "time" entsql "entgo.io/ent/dialect/sql" "github.com/grafana/pyroscope-go" // Replacement for /net/pprof. "github.com/lightsparkdev/spark/common/btcnetwork" "go.uber.org/zap" "go.uber.org/zap/zapcore" "golang.org/x/sync/errgroup" "github.com/XSAM/otelsql" "github.com/go-co-op/gocron/v2" grpcmiddleware "github.com/grpc-ecosystem/go-grpc-middleware" "github.com/improbable-eng/grpc-web/go/grpcweb" "github.com/jackc/pgx/v5/stdlib" _ "github.com/lib/pq" "github.com/lightsparkdev/spark/common" "github.com/lightsparkdev/spark/common/keys" "github.com/lightsparkdev/spark/common/logging" "github.com/lightsparkdev/spark/so" "github.com/lightsparkdev/spark/so/authn" "github.com/lightsparkdev/spark/so/authninternal" "github.com/lightsparkdev/spark/so/authz" "github.com/lightsparkdev/spark/so/chain" "github.com/lightsparkdev/spark/so/consensus" "github.com/lightsparkdev/spark/so/db" "github.com/lightsparkdev/spark/so/ent" _ "github.com/lightsparkdev/spark/so/ent/runtime" "github.com/lightsparkdev/spark/so/entephemeral" sparkerrors "github.com/lightsparkdev/spark/so/errors" "github.com/lightsparkdev/spark/so/partner" "github.com/lightsparkdev/spark/so/rpcauth/brontide" "github.com/lightsparkdev/spark/so/rpcpolicy" sparkgrpc "github.com/lightsparkdev/spark/so/grpc" "github.com/lightsparkdev/spark/so/grpc/grpcutil" "github.com/lightsparkdev/spark/so/handler" "github.com/lightsparkdev/spark/so/knobs" "github.com/lightsparkdev/spark/so/middleware" events "github.com/lightsparkdev/spark/so/stream" "github.com/lightsparkdev/spark/so/task" _ "github.com/mattn/go-sqlite3" "github.com/prometheus/client_golang/prometheus/promhttp" "go.opentelemetry.io/contrib/instrumentation/google.golang.org/grpc/otelgrpc" "go.opentelemetry.io/contrib/instrumentation/net/http/otelhttp" "go.opentelemetry.io/otel" "go.opentelemetry.io/otel/attribute" otelprom "go.opentelemetry.io/otel/exporters/prometheus" "go.opentelemetry.io/otel/propagation" "go.opentelemetry.io/otel/sdk/metric" semconv "go.opentelemetry.io/otel/semconv/v1.17.0" "go.opentelemetry.io/otel/trace/noop" "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials" "google.golang.org/grpc/health/grpc_health_v1" "google.golang.org/grpc/keepalive" ) type args struct { LogLevel string LogJSON bool LogRequestStats bool ConfigFilePath string Index uint64 IdentityPrivateKeyFilePath string OperatorsFilePath string Threshold uint64 SignerAddress string Port uint64 HttpPort uint64 GrpcPort uint64 InternalGrpcPort uint64 DatabasePath string EphemeralDatabasePath string RunningLocally bool ChallengeTimeout time.Duration SessionDuration time.Duration AuthzEnforced bool DisableDKG bool DisableChainwatcher bool SupportedNetworks string AWS bool EphemeralAWS bool ServerCertPath string ServerKeyPath string RunDirectory string RateLimiterEnabled bool RateLimiterMemcachedAddrs string EntDebug bool PyroscopeServer string RisingWaveDatabasePath string } const operatorPoolKnobRefreshInterval = time.Minute // requestBodyReadTimeout bounds how long the server spends reading a single request body on the // ServeHTTP multiplex path. It is sized to let a large (~MaxRequestSize) upload complete over a // slow mobile link while still cutting off a client that dribbles a body indefinitely. See the // usage site for why neither ReadHeaderTimeout, IdleTimeout, ReadTimeout, nor grpc-go's // MaxConnectionAge covers this case. const requestBodyReadTimeout = 60 * time.Second func (a *args) SupportedNetworksList() []btcnetwork.Network { var networks []btcnetwork.Network if strings.Contains(a.SupportedNetworks, "mainnet") || a.SupportedNetworks == "" { networks = append(networks, btcnetwork.Mainnet) } if strings.Contains(a.SupportedNetworks, "testnet") || a.SupportedNetworks == "" { networks = append(networks, btcnetwork.Testnet) } if strings.Contains(a.SupportedNetworks, "regtest") || a.SupportedNetworks == "" { networks = append(networks, btcnetwork.Regtest) } if strings.Contains(a.SupportedNetworks, "signet") || a.SupportedNetworks == "" { networks = append(networks, btcnetwork.Signet) } return networks } func loadArgs() (*args, error) { args := &args{} // Define flags flag.StringVar(&args.LogLevel, "log-level", "debug", "Logging level: debug|info|warn|error") flag.BoolVar(&args.LogJSON, "log-json", false, "Output logs in JSON format") flag.BoolVar(&args.LogRequestStats, "log-request-stats", false, "Log request stats (requires log-json)") flag.StringVar(&args.ConfigFilePath, "config", "so_config.yaml", "Path to config file") flag.Uint64Var(&args.Index, "index", 0, "Index value") flag.StringVar(&args.IdentityPrivateKeyFilePath, "key", "", "Identity private key") flag.StringVar(&args.OperatorsFilePath, "operators", "", "Path to operators file") flag.Uint64Var(&args.Threshold, "threshold", 0, "Threshold value") flag.StringVar(&args.SignerAddress, "signer", "", "Signer address") flag.Uint64Var(&args.Port, "port", 0, "DEPRECATED: Use --http-port instead. HTTP port (grpc-web + metrics)") flag.Uint64Var(&args.HttpPort, "http-port", 0, "HTTP port (grpc-web + metrics)") flag.Uint64Var(&args.GrpcPort, "grpc-port", 0, "Native gRPC port (if 0 or same as http-port, uses ServeHTTP multiplexing)") flag.Uint64Var(&args.InternalGrpcPort, "internal-grpc-port", 0, "If non-zero, start a second gRPC listener on this port that requires TLS + brontide mutual auth. Hosts the SO-to-SO services (SparkInternal, SparkTokenInternal, Gossip, DKG). Internal services remain registered on the public listener too; a future change will introduce a flag to take them off the public listener once peer clients are brontide-aware.") flag.StringVar(&args.DatabasePath, "database", "", "Path to database file") flag.StringVar(&args.EphemeralDatabasePath, "ephemeral-database", "", "Path to ephemeral database file") flag.BoolVar(&args.RunningLocally, "local", false, "Running locally") flag.DurationVar(&args.ChallengeTimeout, "challenge-timeout", time.Minute, "Challenge timeout") flag.DurationVar(&args.SessionDuration, "session-duration", time.Minute*15, "Session duration") flag.BoolVar(&args.AuthzEnforced, "authz-enforced", true, "Enforce authorization checks") flag.BoolVar(&args.DisableDKG, "disable-dkg", false, "Disable DKG") flag.BoolVar(&args.DisableChainwatcher, "disable-chainwatcher", false, "Disable Chainwatcher") flag.StringVar(&args.SupportedNetworks, "supported-networks", "", "Supported networks") flag.BoolVar(&args.AWS, "aws", false, "Use AWS RDS") flag.BoolVar(&args.EphemeralAWS, "ephemeral-aws", false, "Use AWS RDS for the ephemeral database (defaults to the value of --aws if not explicitly set)") flag.StringVar(&args.ServerCertPath, "server-cert", "", "Path to server certificate") flag.StringVar(&args.ServerKeyPath, "server-key", "", "Path to server key") flag.StringVar(&args.RunDirectory, "run-dir", "", "Run directory for resolving relative paths") flag.StringVar(&args.RateLimiterMemcachedAddrs, "rate-limiter-memcached-addrs", "", "Comma-separated list of Memcached addresses") flag.BoolVar(&args.EntDebug, "ent-debug", false, "Log all the SQL queries") flag.StringVar(&args.PyroscopeServer, "pyroscope-server", "", "The address of the Pyroscope server to connect to. Leave blank to skip Pyroscope monitoring.") flag.StringVar(&args.RisingWaveDatabasePath, "risingwave-database", "", "RisingWave Postgres-compatible DSN for partner analytics queries") flag.Parse() var ephemeralAWSFlagSet bool flag.Visit(func(f *flag.Flag) { if f.Name == "ephemeral-aws" { ephemeralAWSFlagSet = true } }) if !ephemeralAWSFlagSet { args.EphemeralAWS = args.AWS } if args.IdentityPrivateKeyFilePath == "" { return nil, errors.New("identity private key file path is required") } if args.OperatorsFilePath == "" { return nil, errors.New("operators file is required") } if args.SignerAddress == "" { return nil, errors.New("signer address is required") } if args.HttpPort == 0 && args.Port == 0 { return nil, errors.New("http-port (or deprecated --port) is required") } if args.HttpPort == 0 && args.Port != 0 { args.HttpPort = args.Port _, _ = fmt.Fprintf(os.Stderr, "WARNING: --port is deprecated, use --http-port instead\n") } if args.HttpPort != 0 && args.Port != 0 && args.HttpPort != args.Port { _, _ = fmt.Fprintf(os.Stderr, "WARNING: Both --port (%d) and --http-port (%d) specified; using --http-port value\n", args.Port, args.HttpPort) } if args.InternalGrpcPort != 0 { if args.InternalGrpcPort == args.HttpPort { return nil, fmt.Errorf("internal-grpc-port (%d) must differ from http-port", args.InternalGrpcPort) } // grpc-port of 0 means ServeHTTP multiplexing with no separate gRPC listener, so it only conflicts when non-zero. if args.GrpcPort != 0 && args.InternalGrpcPort == args.GrpcPort { return nil, fmt.Errorf("internal-grpc-port (%d) must differ from grpc-port", args.InternalGrpcPort) } } return args, nil } func createRateLimiter(config *so.Config, opts ...middleware.RateLimiterOption) (*middleware.RateLimiter, error) { if !config.RateLimiter.Enabled { return nil, nil } return middleware.NewRateLimiter(config, opts...) } type BufferedBody struct { bodyReader io.ReadCloser body []byte position int readErr error loaded bool } func (body *BufferedBody) Read(p []byte) (n int, err error) { if !body.loaded { body.body, body.readErr = io.ReadAll(body.bodyReader) body.loaded = true } // The read error is sticky: surface it on every call rather than only on the read that triggered buffering. // Otherwise, a caller that retries after a partial read would drain the buffered prefix and see a clean io.EOF, // mistaking a truncated body for a complete one. if body.readErr != nil { return 0, body.readErr } n = copy(p, body.body[body.position:]) body.position += n if body.position == len(body.body) { err = io.EOF } return n, err } func (body *BufferedBody) Close() error { return body.bodyReader.Close() } func NewBufferedBody(bodyReader io.ReadCloser) *BufferedBody { return &BufferedBody{bodyReader: bodyReader} } // writeGrpcResourceExhausted writes a terminal ResourceExhausted status directly to w in the wire format // matching contentType, without invoking the gRPC server. gRPC-web carries the status in a trailer frame in // the response body (base64-encoded for the "-text" variants); native gRPC uses a Trailers-Only response // where the status rides in the header frame. In all cases the HTTP status is 200 — the gRPC status is the // real result. func writeGrpcResourceExhausted(w http.ResponseWriter, contentType, msg string) { code := strconv.Itoa(int(codes.ResourceExhausted)) lowerCT := strings.ToLower(contentType) if strings.HasPrefix(lowerCT, "application/grpc-web") { trailers := fmt.Sprintf("grpc-status:%s\r\ngrpc-message:%s\r\n", code, msg) var frame bytes.Buffer frame.WriteByte(1 << 7) // high bit marks a trailer frame // A trailer block this small cannot overflow uint32. _ = binary.Write(&frame, binary.BigEndian, uint32(len(trailers))) frame.WriteString(trailers) body := frame.Bytes() if strings.HasPrefix(lowerCT, "application/grpc-web-text") { encoded := make([]byte, base64.StdEncoding.EncodedLen(len(body))) base64.StdEncoding.Encode(encoded, body) body = encoded } w.Header().Set("Content-Type", contentType) w.WriteHeader(http.StatusOK) _, _ = w.Write(body) return } w.Header().Set("Content-Type", "application/grpc") w.Header().Set("Grpc-Status", code) w.Header().Set("Grpc-Message", msg) w.WriteHeader(http.StatusOK) } func main() { args, err := loadArgs() // We have to use the package-level logger until we get the real one set up. if err != nil { zap.S().Fatalf("Failed to load args: %v", err) } logConfig := zap.NewProductionConfig() logLevel, err := zap.ParseAtomicLevel(args.LogLevel) if err != nil { zap.S().Fatalf("Failed to parse log level: %v", err) } logConfig.Level = logLevel if args.LogJSON { logConfig.Encoding = "json" logConfig.EncoderConfig = zap.NewProductionEncoderConfig() } else { logConfig.Encoding = "console" logConfig.EncoderConfig = zap.NewDevelopmentEncoderConfig() } // Various settings to make logs more similar to slog (both so they're backwards compatible with // downstream ingestion and just generally similar). logConfig.EncoderConfig.TimeKey = "time" logConfig.EncoderConfig.CallerKey = zapcore.OmitKey logConfig.EncoderConfig.EncodeTime = zapcore.RFC3339NanoTimeEncoder logConfig.EncoderConfig.EncodeLevel = zapcore.CapitalLevelEncoder // Disable sampling to ensure all logs are captured. logConfig.Sampling = nil logger, err := logConfig.Build() if err != nil { zap.S().Fatalf("Failed to build logger: %v", err) } defer func() { // Try to make sure we log any panics that occur. if r := recover(); r != nil { logger.Error("Panic in main", zap.Any("panic", r), zap.Stack("stack"), ) } _ = logger.Sync() }() // Now we can start using the logger itself. logger = logger.WithOptions(zap.WrapCore(func(core zapcore.Core) zapcore.Core { return &logging.SourceCore{Core: core} })) sigCtx, done := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM) defer done() errGrp, errCtx := errgroup.WithContext(sigCtx) errCtx = logging.Inject(errCtx, logger) config, err := so.NewConfig( errCtx, args.ConfigFilePath, args.Index, args.IdentityPrivateKeyFilePath, args.OperatorsFilePath, // TODO: Refactor this into the yaml config args.Threshold, args.SignerAddress, args.DatabasePath, args.EphemeralDatabasePath, args.AWS, args.EphemeralAWS, args.AuthzEnforced, args.SupportedNetworksList(), args.ServerCertPath, args.ServerKeyPath, args.RunDirectory, so.RateLimiterConfig{ Enabled: args.RateLimiterEnabled, }, ) if err != nil { logger.Fatal("Failed to create config", zap.Error(err)) } // The operator map includes our own entry and some flows self-dial it (e.g. DKG), so brontide client transport // without the local brontide listener would fail on the first self-dial once the knob flips on. Fail at boot instead. if config.InternalRPC.Transport == so.InternalRPCTransportBrontide && args.InternalGrpcPort == 0 { logger.Fatal("internal_rpc.transport=brontide requires --internal-grpc-port so the local brontide listener is running") } config.RisingWaveDSN = args.RisingWaveDatabasePath // OBSERVABILITY promExporter, err := otelprom.New() if err != nil { logger.Fatal("Failed to create prometheus exporter", zap.Error(err)) } // Default gRPC OTEL histogram boundaries top out at 10s; some of our calls are going longer than that. // We top out at 180s because that's our current client and server timeouts grpcLatencyBucketsMs := []float64{ 0, 5, 10, 25, 50, 75, 100, 250, 500, 750, 1000, 2500, 5000, 7500, 10000, 15000, 30000, 60000, 90000, 120000, 180000, } meterProvider := metric.NewMeterProvider( metric.WithReader(promExporter), metric.WithView( metric.NewView( metric.Instrument{Name: "rpc.server.duration"}, metric.Stream{Aggregation: metric.AggregationExplicitBucketHistogram{Boundaries: grpcLatencyBucketsMs}}, ), metric.NewView( metric.Instrument{Name: "rpc.client.duration"}, metric.Stream{Aggregation: metric.AggregationExplicitBucketHistogram{Boundaries: grpcLatencyBucketsMs}}, ), ), ) otel.SetMeterProvider(meterProvider) otel.SetTextMapPropagator(propagation.TraceContext{}) if config.Tracing.Enabled { shutdown, err := common.ConfigureTracing(errCtx, config.Tracing) if err != nil { logger.Fatal("Failed to configure tracing", zap.Error(err)) } defer func() { shutdownCtx, shutdownRelease := context.WithTimeout(context.Background(), 10*time.Second) defer shutdownRelease() logger.Info("Shutting down tracer provider") if err := shutdown(shutdownCtx); err != nil { logger.Error("Error shutting down tracer provider", zap.Error(err)) } else { logger.Info("Tracer provider shut down") } }() } var valuesProvider knobs.KnobsValuesProvider if config.Knobs.IsEnabled() { if provider, err := knobs.NewKnobsK8ValuesProvider(errCtx, config.Knobs.Namespace); err != nil { // Knobs has failed to fetch the config, so the controllers will rely on the default values. logger.Error("Failed to create K8 knobs", zap.Error(err)) } else { valuesProvider = provider } } else if args.RunningLocally && len(config.Knobs.StaticValues) > 0 { // Fallback for bare-process deployments (dev-cli, run-everything.sh): // use static values from the YAML config valuesProvider = knobs.NewStaticValuesProvider(config.Knobs.StaticValues) } // Knobs service is always defined, no need to check for nil. // If the provider is nil, the knobs service will use the default values. knobsService := knobs.New(valuesProvider) // Start profiling server if enabled (localhost only) pprofServer := setUpPprof(errGrp, logger) shutDownPyroscope := setUpPyroscope(args, logger) defer shutDownPyroscope() dbDriver := config.DatabaseDriver() connector, err := so.NewDBConnector(context.Background(), config, knobsService) if err != nil { logger.Fatal("Failed to create db connector", zap.Error(err)) } defer connector.Close() ephemeralEnabled := config.EphemeralDatabasePath != "" var ephemeralConnector *so.DBConnector if ephemeralEnabled { ephemeralConnector, err = so.NewEphemeralDBConnector(context.Background(), config, knobsService) if err != nil { logger.Fatal("Failed to create ephemeral db connector", zap.Error(err)) } defer ephemeralConnector.Close() } else { logger.Info("Ephemeral DB disabled because no URI was configured") } for _, op := range config.SigningOperatorMap { op.SetTimeoutProvider(knobs.NewKnobsTimeoutProvider(knobsService, config.GRPC.ClientTimeout)) op.SetConnectionPoolConfig(operatorPoolConfigFromKnobs(knobsService, op.Identifier)) } errGrp.Go(func() error { ticker := time.NewTicker(operatorPoolKnobRefreshInterval) defer ticker.Stop() for { select { case <-errCtx.Done(): return nil case <-ticker.C: for _, op := range config.SigningOperatorMap { op.SetConnectionPoolConfig(operatorPoolConfigFromKnobs(knobsService, op.Identifier)) } } } }) config.FrostGRPCConnectionFactory.SetTimeoutProvider( knobs.NewKnobsTimeoutProvider(knobsService, config.GRPC.ClientTimeout)) var sqlDb entsql.ExecQuerier var ephemeralDb entsql.ExecQuerier if dbDriver == "postgres" { sqlDb = stdlib.OpenDBFromPool(connector.Pool()) if ephemeralEnabled { if ephemeralConnector.Pool() == nil { logger.Fatal( "Ephemeral DB pool is not initialized. Check ephemeral DB configuration", zap.String("ephemeral_database_path", config.EphemeralDatabasePath), ) } ephemeralDb = stdlib.OpenDBFromPool(ephemeralConnector.Pool()) } } else { sqlDb = otelsql.OpenDB(connector, otelsql.WithSpanOptions(so.OtelSQLSpanOptions)) if ephemeralEnabled { ephemeralDb = otelsql.OpenDB(ephemeralConnector, otelsql.WithSpanOptions(so.OtelSQLSpanOptions)) } } dialectDriver := entsql.NewDriver(dbDriver, entsql.Conn{ExecQuerier: sqlDb}) var dbClient *ent.Client var ephemeralDbClient *entephemeral.Client var ephemeralDialectDriver *entsql.Driver if ephemeralEnabled { ephemeralDialectDriver = entsql.NewDriver(dbDriver, entsql.Conn{ExecQuerier: ephemeralDb}) } if args.EntDebug { dbClient = ent.NewClient(ent.Driver(dialectDriver), ent.Debug()) if ephemeralEnabled { ephemeralDbClient = entephemeral.NewClient(entephemeral.Driver(ephemeralDialectDriver), entephemeral.Debug()) } } else { dbClient = ent.NewClient(ent.Driver(dialectDriver)) if ephemeralEnabled { ephemeralDbClient = entephemeral.NewClient(entephemeral.Driver(ephemeralDialectDriver)) } } // Add interceptor for query stats and read operation metrics dbClient.Intercept(ent.DatabaseStatsInterceptor("main")) if ephemeralDbClient != nil { ephemeralDbClient.Intercept(ent.DatabaseStatsInterceptor("ephemeral")) } // Add hook for mutation operation metrics (insert, update, delete) dbClient.Use(ent.DatabaseOperationsHook("main")) if ephemeralDbClient != nil { ephemeralDbClient.Use(ent.DatabaseOperationsHook("ephemeral")) } defer func() { _ = dbClient.Close() if ephemeralDbClient != nil { _ = ephemeralDbClient.Close() } }() if dbDriver == "sqlite3" { applySQLitePragmas := func(dbPath string, dbName string) { sqliteDb, err := sql.Open("sqlite3", dbPath) if err != nil { logger.With(zap.Error(err)).Sugar().Fatalf("Failed to open sqlite database for PRAGMA initialization (%s)", dbName) } if sqliteDb != nil { defer func() { _ = sqliteDb.Close() }() } // journal_mode=WAL is a file-level setting that persists across connections, // so setting it once here is sufficient. if _, err := sqliteDb.ExecContext(errCtx, "PRAGMA journal_mode=WAL;"); err != nil { logger.With(zap.Error(err)).Sugar().Fatalf("Failed to set journal_mode for %s sqlite database", dbName) } // busy_timeout is a per-connection setting; it is applied via _busy_timeout=5000 // in the SQLite DSN inside newDBConnector so it takes effect on every pool connection. } applySQLitePragmas(config.DatabasePath, "main") if ephemeralEnabled { applySQLitePragmas(config.EphemeralDatabasePath, "ephemeral") } } dbEvents, err := db.NewDBEvents(errCtx, dbClient, logger.With(zap.String("component", "dbevents"))) if err != nil { logger.Fatal("Failed to create db events", zap.Error(err)) } if config.Database.DBEventsEnabled != nil && *config.Database.DBEventsEnabled { errGrp.Go(func() error { if err := dbEvents.Start(); err != nil { logger.Error("Error in dbevents", zap.Error(err)) return err } if errCtx.Err() == nil { // This technically isn't an error, but raise it as one because dbevents should never // stop unless we explicitly tell it to when shutting down! return fmt.Errorf("dbevents stopped unexpectedly") } return nil }) } frostConnection, err := config.NewFrostGRPCConnection() if err != nil { logger.Fatal("Failed to create frost client", zap.Error(err)) } if !args.DisableChainwatcher { // Chain watchers for network, bitcoindConfig := range config.BitcoindConfigs { errGrp.Go(func() error { chainCtx, chainCancel := context.WithCancel(errCtx) defer chainCancel() chainLogger := logger.With(zap.String("component", "chainwatcher"), zap.String("network", network)) chainCtx = logging.Inject(chainCtx, chainLogger) chainCtx = knobs.InjectKnobsService(chainCtx, knobsService) if err := chain.WatchChain(chainCtx, config, dbClient, ephemeralDbClient, bitcoindConfig); err != nil { logger.Error("Error in chain watcher", zap.Error(err)) return err } if errCtx.Err() == nil { // This technically isn't an error, but raise it as one because our chain watcher should never // stop unless we explicitly tell it to when shutting down! return fmt.Errorf("chain watcher for %s stopped unexpectedly", network) } return nil }) } } if !args.DisableDKG { // Scheduled tasks setup cronCtx, cronCancel := context.WithCancel(errCtx) defer cronCancel() taskLogger := logger.With(zap.String("component", "cron")) cronCtx = logging.Inject(cronCtx, taskLogger) taskLogger.Info("Starting scheduler") taskMonitor, err := task.NewMonitor() if err != nil { taskLogger.Fatal("Failed to create task monitor", zap.Error(err)) } scheduler, err := gocron.NewScheduler( gocron.WithGlobalJobOptions( gocron.WithContext(cronCtx), gocron.WithSingletonMode(gocron.LimitModeReschedule), ), gocron.WithLogger(task.NewZapLoggerAdapter(taskLogger)), gocron.WithMonitorStatus(taskMonitor), ) if err != nil { logger.Fatal("Failed to create scheduler", zap.Error(err)) } for _, scheduled := range task.AllScheduledTasks() { // Don't run the task if the task specifies it should not be run in // test environments and RunningLocally is set (eg. we are in a test environment) if (!args.RunningLocally || scheduled.RunInTestEnv) && !scheduled.Disabled { err := scheduled.Schedule(scheduler, config, dbClient, ephemeralDbClient, knobsService) if err != nil { logger.Fatal("Failed to create job", zap.Error(err)) } } } scheduler.Start() defer func() { _ = scheduler.Shutdown() }() // Run startup tasks startupCtx, startupCancel := context.WithCancel(errCtx) defer startupCancel() errGrp.Go(func() error { // TODO(mhr): Do this properly, have a waitgroup in `RunStartupTasks` that waits until all tasks // are done before returning. startupCtx = logging.Inject(startupCtx, logger.With(zap.String("component", "startup"))) return task.RunStartupTasks(startupCtx, config, dbClient, ephemeralDbClient, args.RunningLocally, knobsService) }) } sessionTokenCreatorVerifier, err := authninternal.NewSessionTokenCreatorVerifier(config.IdentityPrivateKey, nil) if err != nil { logger.Fatal("Failed to create token verifier", zap.Error(err)) } var rateLimiter *middleware.RateLimiter logger.Sugar().Infof( "Rate limiter config: enabled %t", config.RateLimiter.Enabled, ) if config.RateLimiter.Enabled { var err error rlOpts := []middleware.RateLimiterOption{middleware.WithKnobs(knobsService)} memcachedURI := strings.TrimSpace(config.CacheURI) if knobsService.RolloutRandom(knobs.KnobRateLimitMemcacheEnabled, 0) && memcachedURI != "" { baseMaxIdleConns := 32 maxIdleConns := int(knobsService.GetValue( knobs.KnobRateLimitMemcacheMaxIdleConns, float64(baseMaxIdleConns), )) store, sErr := middleware.NewMemcacheStore(maxIdleConns, memcachedURI) if sErr != nil { logger.Warn("Memcached rate limiter store unavailable, falling back to in-memory", zap.Error(sErr)) } else { rlOpts = append(rlOpts, middleware.WithStore(store)) logger.Sugar().Infof("Rate limiter using Memcached store. memcached_addr=%s", memcachedURI) } } rateLimiter, err = createRateLimiter(config, rlOpts...) if err != nil { logger.Fatal("Failed to create rate limiter", zap.Error(err)) } } clientInfoProvider := sparkgrpc.NewGRPCClientInfoProvider(config.XffClientIpPosition) var tableLogger *logging.TableLogger if args.LogRequestStats && args.LogJSON { tableLogger = logging.NewTableLogger(clientInfoProvider) } var ephemeralSessionFactory db.EphemeralSessionFactory if ephemeralDbClient != nil { ephemeralSessionFactory = db.NewDefaultEphemeralSessionFactory(ephemeralDbClient) } serverOpts := []grpc.ServerOption{ grpc.StatsHandler( sparkgrpc.NewInstrumentedStatsHandler(otelgrpc.NewServerHandler()), ), // Align the gRPC transport receive limit with the MaxRequestSize validation cap. gRPC's default is 4MB, so a // message larger than 4MB but within MaxRequestSize would be rejected at the transport layer before the // ValidationInterceptor's MaxRequestSize check ever runs. This lives in the shared serverOpts deliberately: it // applies to both the public server and the brontide-authenticated internal (SO-to-SO) server, which run the // same ValidationInterceptor, so both surfaces enforce one consistent ceiling rather than leaving the internal // server on the 4MB transport default beneath a 10MB validation cap. grpc.MaxRecvMsgSize(sparkgrpc.MaxRequestSize), } // Establish base values from config, then allow runtime knobs to override. // All of these (grpcConnTimeout, grpcKeepaliveTime, grpcKeepaliveTimeout, // and the max connection age values) are read once when the server is // created and cannot be changed at runtime; flipping the knob takes effect // on the next operator restart. The enforcement-policy values // (grpcKeepaliveMinTime, grpcKeepalivePermitWithoutStream) are config-only // for the same reason — there's no live override that could affect an // already-running server. grpcConnTimeout := knobsService.GetDuration(knobs.KnobGrpcServerConnectionTimeout, config.GRPC.ServerConnectionTimeout) grpcKeepaliveTime := knobsService.GetDuration(knobs.KnobGrpcServerKeepaliveTime, config.GRPC.ServerKeepaliveTime) grpcKeepaliveTimeout := knobsService.GetDuration(knobs.KnobGrpcServerKeepaliveTimeout, config.GRPC.ServerKeepaliveTimeout) grpcKeepaliveMinTime := config.GRPC.ServerKeepaliveMinTime grpcKeepalivePermitWithoutStream := config.GRPC.ServerKeepalivePermitWithoutStream != nil && *config.GRPC.ServerKeepalivePermitWithoutStream grpcMaxConnectionAge := knobsService.GetDuration(knobs.KnobGrpcServerMaxConnectionAge, config.GRPC.ServerMaxConnectionAge) grpcMaxConnectionAgeGrace := knobsService.GetDuration(knobs.KnobGrpcServerMaxConnectionAgeGrace, config.GRPC.ServerMaxConnectionAgeGrace) // This uses SetDeadline in net.Conn to set the timeout for the connection // establishment, after which the connection is closed with error // `DeadlineExceeded`. if grpcConnTimeout > 0 { serverOpts = append(serverOpts, grpc.ConnectionTimeout(grpcConnTimeout)) } // Keepalive detects dead connections and closes them. // Time is the interval between keepalive pings. // Timeout is the interval between keepalive pings after which the connection is closed. serverOpts = append(serverOpts, grpc.KeepaliveParams(keepalive.ServerParameters{ Time: grpcKeepaliveTime, Timeout: grpcKeepaliveTimeout, MaxConnectionAge: grpcMaxConnectionAge, MaxConnectionAgeGrace: grpcMaxConnectionAgeGrace, })) // Enforcement policy bounds how aggressively clients are allowed to ping. // The grpc-go default (MinTime=5m, PermitWithoutStream=false) is stricter // than sparkcore Python clients, which ping idle channels every 30s; the // resulting ENHANCE_YOUR_CALM GOAWAYs forced reconnects and added tail // latency to subsequent calls (e.g. Lightning completion). serverOpts = append(serverOpts, grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{ MinTime: grpcKeepaliveMinTime, PermitWithoutStream: grpcKeepalivePermitWithoutStream, })) logger.Sugar().Infof( "gRPC server keepalive: time=%s timeout=%s max_conn_age=%s max_conn_age_grace=%s enforcement_min_time=%s permit_without_stream=%t", grpcKeepaliveTime, grpcKeepaliveTimeout, grpcMaxConnectionAge, grpcMaxConnectionAgeGrace, grpcKeepaliveMinTime, grpcKeepalivePermitWithoutStream, ) concurrencyGuard := sparkgrpc.NewConcurrencyGuard(knobsService, sparkgrpc.KnobTargetName_UnaryGlobalLimit) concurrencyStreamGuard := sparkgrpc.NewConcurrencyGuard(knobsService, sparkgrpc.KnobTargetName_StreamGlobalLimit) var eventsRouter *events.EventRouter if config.Database.DBEventsEnabled != nil && *config.Database.DBEventsEnabled { eventsRouter = events.NewEventRouter(errCtx, dbClient, dbEvents, logger.With(zap.String("component", "events_router")), config) } // Add Interceptors aka gRPC middleware // // Interceptors wrap RPC handlers so we can apply cross‑cutting concerns in one place // and in a defined order. We install separate chains for unary (request/response) // and streaming RPCs. dbSessionMiddleware := sparkgrpc.DatabaseSessionMiddleware( dbClient, db.NewDefaultSessionFactory(dbClient), ephemeralSessionFactory, config.Database.NewTxTimeout, ) // The 2PC engine is stateless across requests (it just carries config, // the gossip sender, and the unwrapped *ent.Client used for engine // bookkeeping writes that must outlive the request transaction), so a // single instance is constructed here and injected into every request // ctx by ConsensusEngineInterceptor below. Handlers fetch it via // consensus.GetEngine(ctx) instead of constructing one per call. consensusEngine := consensus.NewTwoPCEngine(config, handler.NewSendGossipHandler(config), db.NewDefaultSessionFactory(dbClient)) serverOpts = append(serverOpts, grpc.UnaryInterceptor(grpcmiddleware.ChainUnaryServer( sparkgrpc.TimestampHeaderInterceptor(), sparkerrors.ErrorInterceptor(config.ReturnDetailedErrors), sparkgrpc.TracingInterceptor(), sparkgrpc.LogInterceptor(logger.With(zap.String("component", "grpc")), tableLogger), sparkgrpc.SparkTokenMetricsInterceptor(), // Inject knobs into context for unary requests func() grpc.UnaryServerInterceptor { return func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { ctx = knobs.InjectKnobsService(ctx, knobsService) return handler(ctx, req) } }(), sparkgrpc.MethodDisableInterceptor(), sparkgrpc.TimeoutInterceptor(knobsService, config.GRPC.ServerUnaryHandlerTimeout), sparkgrpc.PanicRecoveryInterceptor(config.ReturnDetailedPanicErrors), authn.NewInterceptor(sessionTokenCreatorVerifier).AuthnInterceptor, // SparkPartnerService is AuthPartnerBasic in rpcpolicy (no session token); this // interceptor enforces its HTTP Basic Auth and is a no-op for all other services. partner.NewBasicAuthInterceptor(dbClient).UnaryServerInterceptor, partner.NewInterceptor(dbClient).KnobGatedInterceptor(knobsService, config.Index == 0), // SO0 (coordinator) only // Concurrency and rate limiting after authentication so pubkey is available for rate limiting // but before DB session. sparkgrpc.ConcurrencyInterceptor(concurrencyGuard, clientInfoProvider, knobsService), func() grpc.UnaryServerInterceptor { if rateLimiter != nil { return rateLimiter.UnaryServerInterceptor() } return func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (any, error) { return handler(ctx, req) } }(), dbSessionMiddleware, // Inject the 2PC engine after the DB session so handlers that // need consensus can pull it from ctx without constructing one. sparkgrpc.ConsensusEngineInterceptor(consensusEngine), authz.NewAuthzInterceptor(authz.NewAuthzConfig( authz.WithMode(config.ServiceAuthz.Mode), authz.WithAllowedIPs(config.ServiceAuthz.IPAllowlist), authz.WithIsProtectedMethod(rpcpolicy.IsInternalOnly), authz.WithXffClientIpPosition(config.XffClientIpPosition), )).UnaryServerInterceptor, sparkgrpc.ValidationInterceptor(), // Idempotency must be after the DB session so it can store keys, // and after authz/validation so cached responses cannot bypass // internal-service allowlist checks or request-shape validation. sparkgrpc.IdempotencyInterceptor(), )), grpc.StreamInterceptor(grpcmiddleware.ChainStreamServer( sparkerrors.ErrorStreamingInterceptor(), sparkgrpc.StreamLogInterceptor(logger.With(zap.String("component", "grpc"))), func() grpc.StreamServerInterceptor { return func(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error { ctx := knobs.InjectKnobsService(ss.Context(), knobsService) return handler(srv, &grpcmiddleware.WrappedServerStream{ServerStream: ss, WrappedContext: ctx}) } }(), sparkgrpc.MethodDisableStreamInterceptor(), sparkgrpc.PanicRecoveryStreamInterceptor(), authn.NewInterceptor(sessionTokenCreatorVerifier).StreamAuthnInterceptor, sparkgrpc.ConcurrencyStreamInterceptor(concurrencyStreamGuard, clientInfoProvider, knobsService), func() grpc.StreamServerInterceptor { if rateLimiter != nil { return rateLimiter.StreamServerInterceptor() } return func(srv any, ss grpc.ServerStream, _ *grpc.StreamServerInfo, handler grpc.StreamHandler) error { return handler(srv, ss) } }(), authz.NewAuthzInterceptor(authz.NewAuthzConfig( authz.WithMode(config.ServiceAuthz.Mode), authz.WithAllowedIPs(config.ServiceAuthz.IPAllowlist), authz.WithIsProtectedMethod(rpcpolicy.IsInternalOnly), authz.WithXffClientIpPosition(config.XffClientIpPosition), )).StreamServerInterceptor, sparkgrpc.StreamValidationInterceptor(), )), ) cert, err := tls.LoadX509KeyPair(args.ServerCertPath, args.ServerKeyPath) if err != nil { logger.Fatal("Failed to load server certificate", zap.Error(err)) } tlsConfig := tls.Config{ Certificates: []tls.Certificate{cert}, ClientAuth: tls.NoClientCert, MinVersion: tls.VersionTLS12, } tlsCreds := credentials.NewTLS(&tlsConfig) publicOpts := append(slices.Clone(serverOpts), grpc.Creds(tlsCreds)) grpcServer := grpc.NewServer(publicOpts...) // RisingWave client for partner analytics queries (connects lazily; nil when the // --risingwave-database DSN is empty). Owned here and passed in like dbClient. rwClient := partner.NewRisingWaveClient(config.RisingWaveDSN) defer func() { if rwClient != nil { _ = rwClient.Close() } }() err = RegisterPublicGrpcServers( grpcServer, args, config, logger, dbClient, ephemeralDbClient, sessionTokenCreatorVerifier, eventsRouter, rwClient, ) if err != nil { logger.Fatal("Failed to register all gRPC servers", zap.Error(err)) } // SO-to-SO services are registered on the public listener so existing peer clients continue to reach them. A follow-up // change will make the client factory brontide-aware and introduce a flag to take these services off the public listener entirely. if err := RegisterInternalGrpcServers(grpcServer, args, config, frostConnection, dbClient); err != nil { logger.Fatal("Failed to register internal gRPC servers on public listener", zap.Error(err)) } // Optionally start a second listener that requires TLS + brontide mutual auth and hosts only the SO-to-SO services. // When --internal-grpc-port is non-zero, we build a separate grpc.Server with brontide-wrapped TLS credentials and // serve it on its own listener. var internalGrpcServer *grpc.Server if args.InternalGrpcPort != 0 { peerLookup := brontide.PeerLookupFunc(func(pub keys.Public) *brontide.Peer { id := config.GetOperatorIdentifierFromIdentityPublicKey(pub) if id == "" { return nil } return &brontide.Peer{Identifier: id, IdentityPublicKey: pub} }) brontideCreds, err := brontide.NewServerCredentials(brontide.ServerConfig{ Inner: tlsCreds, LocalPrivateKey: config.IdentityPrivateKey, Peers: peerLookup, }) if err != nil { logger.Fatal("Failed to construct brontide server credentials", zap.Error(err)) } internalOpts := append(slices.Clone(serverOpts), grpc.Creds(brontideCreds)) internalGrpcServer = grpc.NewServer(internalOpts...) if err := RegisterInternalGrpcServers(internalGrpcServer, args, config, frostConnection, dbClient); err != nil { logger.Fatal("Failed to register internal gRPC servers on internal listener", zap.Error(err)) } healthServer := sparkgrpc.NewHealthServer(errCtx, dbClient, ephemeralDbClient) grpc_health_v1.RegisterHealthServer(internalGrpcServer, healthServer) } healthServer := sparkgrpc.NewHealthServer(errCtx, dbClient, ephemeralDbClient) grpc_health_v1.RegisterHealthServer(grpcServer, healthServer) // Fail-closed: every method registered on either server must have an rpcpolicy entry. Catches the case where a new // RPC is added without declaring its auth policy. if missing := registeredMethodsMissingPolicy(grpcServer); len(missing) > 0 { logger.Sugar().Fatalf("gRPC methods registered without an rpcpolicy entry: %v", missing) } if internalGrpcServer != nil { if missing := registeredMethodsMissingPolicy(internalGrpcServer); len(missing) > 0 { logger.Sugar().Fatalf("gRPC methods registered on internal listener without an rpcpolicy entry: %v", missing) } } // Web compatibility layer wrappedGrpcServer := grpcweb.WrapServer(grpcServer, grpcweb.WithOriginFunc(func(_ string) bool { return true }), grpcweb.WithWebsockets(true), grpcweb.WithWebsocketOriginFunc(func(_ *http.Request) bool { return true }), grpcweb.WithCorsForRegisteredEndpointsOnly(false), ) // Add a value to contexts that arrive via gRPC-web so downstream handlers can // distinguish them from native gRPC requests. wrappedGrpc := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { if wrappedGrpcServer.IsGrpcWebRequest(r) || wrappedGrpcServer.IsGrpcWebSocketRequest(r) { r = r.WithContext(grpcutil.WithGrpcWebRequest(r.Context())) } wrappedGrpcServer.ServeHTTP(w, r) }) // Determine if we should serve native gRPC separately or use ServeHTTP multiplexing useNativeGRPC := args.GrpcPort != 0 && args.GrpcPort != args.HttpPort mux := http.NewServeMux() // This health check isn't used by k8s, which uses the gRPC health check. However, it // is used by ALB for checking the health of the HTTP server, so we need to keep it. mux.Handle("/-/ready", http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) { w.WriteHeader(http.StatusOK) })) mux.Handle("/metrics", promhttp.Handler()) otelHTTPOpts := []otelhttp.Option{ otelhttp.WithTracerProvider(noop.TracerProvider{}), // Disable tracing, let gRPC server handle it. otelhttp.WithMetricAttributesFn(func(r *http.Request) []attribute.KeyValue { return []attribute.KeyValue{ attribute.String(string(semconv.HTTPRouteKey), r.URL.Path), } }), } var grpcListener net.Listener if useNativeGRPC { logger.Info("Starting with separate gRPC + HTTP servers", zap.Uint64("grpc_port", args.GrpcPort), zap.Uint64("http_port", args.HttpPort)) var err error grpcListener, err = net.Listen("tcp", fmt.Sprintf(":%d", args.GrpcPort)) if err != nil { logger.Fatal("Failed to create gRPC listener", zap.Error(err)) } errGrp.Go(func() error { logger.Info("Native gRPC server listening", zap.Uint64("port", args.GrpcPort)) if err := grpcServer.Serve(grpcListener); err != nil { logger.Error("Native gRPC server failed", zap.Error(err)) return err } return nil }) mux.Handle("/", otelhttp.NewHandler(wrappedGrpc, "server", otelHTTPOpts...)) } else { logger.Info("Starting with ServeHTTP multiplexing", zap.Uint64("port", args.HttpPort)) mux.Handle("/", otelhttp.NewHandler( http.HandlerFunc( func(w http.ResponseWriter, r *http.Request) { // Reject an over-large body up front when the client declares its size. If we instead let it // reach the gRPC server, MaxBytesReader (below) trips mid-body-read and grpc-go maps that to a // transport.ConnectionError -> codes.Unknown (or, if the truncation surfaces as a connection // reset, a client-side codes.Unavailable that retry-configured SDKs re-send indefinitely), // masking what is really a permanent size violation. Returning ResourceExhausted here gives the // client a terminal code. Content-Length is absent (-1) for native gRPC streaming, so this only // fires for unary gRPC-web requests (the SDK's transport), which is exactly the retry-loop case. contentType := r.Header.Get("Content-Type") if r.ContentLength > sparkgrpc.MaxHTTPBodySize && strings.HasPrefix(strings.ToLower(contentType), "application/grpc") { writeGrpcResourceExhausted(w, contentType, fmt.Sprintf( "request body of %d bytes exceeds max of %d", r.ContentLength, sparkgrpc.MaxHTTPBodySize)) return } // Bound the time spent reading the request body. ReadHeaderTimeout only covers headers and // IdleTimeout only fires on connections with no open streams, so neither caps body-read time. // ReadTimeout is intentionally unset because it's a whole-request deadline that would cut off // long-lived server-streaming responses (e.g. subscribe_to_events). On this ServeHTTP multiplex // path the HTTP/2 layer is net/http's, so grpc-go's MaxConnectionAge doesn't govern these // connections either. Without a bound, a client can send headers quickly then dribble a body // just under MaxHTTPBodySize indefinitely, pinning a goroutine and buffer per stream. A read // deadline closes this for native gRPC and gRPC-web-over-HTTP: it bounds only request reads, // not response writes, and is safe because Spark RPCs send the whole request body upfront // (no client-streaming). This does not cut off long-lived server-streams (e.g. // subscribe_to_events), including over an HTTP/1.1 hop: although h1 read deadlines are // connection-level, net/http clears the deadline once the request body reaches EOF (see // connReader.startBackgroundRead), so it only ever governs the body-read window. // // Caveat: this does NOT cover the grpc-web websocket transport (WithWebsockets, above). Once // a request is upgraded, the library hijacks the connection and replaces r.Body with its own // frame reader, discarding both this deadline and the MaxBytesReader/BufferedBody wrap below. // Websocket frame reads are bounded only by nhooyr's default 32KB per-message read limit — // there is no total-byte or time bound on that path. rc := http.NewResponseController(w) if err := rc.SetReadDeadline(time.Now().Add(requestBodyReadTimeout)); err != nil { logger.Error("could not set request body read deadline", zap.Error(err)) } // The gRPC server doesn't read the request body until EOF before processing the request. This // can result in the HTTP server receiving a DATA(END_FRAME) frame after sending the response, // which elicits a RST_STREAM(STREAM_CLOSED) frame. ALB and nginx then respond to the client // with RST_STREAM(INTERNAL_ERROR) which causes the request to fail. The workaround is to buffer // the entire request body before passing to the gRPC server. // // Since we're now reading the whole body into memory on first read, we have to cap it before // buffering to avoid giant requests. The gRPC MaxRequestSize check only runs after message // parsing, meaning the message will already have been buffered in memory here. Spark RPCs don't // use client-streaming, so a single per-request cap is safe. r.Body = NewBufferedBody(http.MaxBytesReader(w, r.Body, sparkgrpc.MaxHTTPBodySize)) if strings.ToLower(r.Header.Get("Content-Type")) == "application/grpc" { grpcServer.ServeHTTP(w, r) return } wrappedGrpc.ServeHTTP(w, r) }, ), "server", otelHTTPOpts..., ), ) } server := &http.Server{ Addr: fmt.Sprintf(":%d", args.HttpPort), Handler: mux, TLSConfig: &tlsConfig, // Prevent a slowloris attack where a peer opens a connection and dribbles headers to pin server resources. // // We deliberately don't set ReadTimeout or WriteTimeout, which are whole-request/response deadlines applied to // every connection, since this server multiplexes long-lived HTTP/2 gRPC server streams, like subscribe_to_events, // whose response stays open indefinitely. ReadHeaderTimeout: 10 * time.Second, IdleTimeout: 120 * time.Second, } errGrp.Go(func() error { logger.Info("HTTP server listening", zap.Uint64("port", args.HttpPort)) if err := server.ListenAndServeTLS("", ""); !errors.Is(err, http.ErrServerClosed) { logger.Error("HTTP server failed", zap.Error(err)) return err } return nil }) if internalGrpcServer != nil { internalListener, err := net.Listen("tcp", fmt.Sprintf(":%d", args.InternalGrpcPort)) if err != nil { logger.Fatal("Failed to create internal gRPC listener", zap.Error(err)) } errGrp.Go(func() error { logger.Sugar().Infof("Internal gRPC server listening (TLS + brontide) on port %d", args.InternalGrpcPort) if err := internalGrpcServer.Serve(internalListener); err != nil { logger.Error("Internal gRPC server failed", zap.Error(err)) return err } return nil }) } // Now we wait... for something to fail. <-errCtx.Done() if sigCtx.Err() != nil { logger.Info("Received shutdown signal, shutting down gracefully...") } else { logger.Error("Shutting down due to error...") } shutdownCtx, shutdownRelease := context.WithTimeout(context.Background(), 30*time.Second) defer shutdownRelease() // We wrap the gRPC server for grpc-web, which proxies HTTP/WebSocket connections into gRPC calls. // If grpcServer.GracefulStop() is called while the HTTP server is still accepting // connections, grpc-web can receive new requests against a draining gRPC server and // panic, crashing the process. Shutting down the HTTP server first drains and closes // all HTTP and WebSocket connections (including those proxied by grpc-web), so that by // the time we call GracefulStop the gRPC server only needs to wait for in-flight RPCs // that arrived through the native gRPC port. // // See: https://github.com/grpc/grpc-go/issues/1384 logger.Info("Stopping HTTP server...") if err := server.Shutdown(shutdownCtx); err != nil { logger.Error("HTTP server failed to shutdown gracefully", zap.Error(err)) } else { logger.Info("HTTP server stopped") } // GracefulStop blocks until all in-flight RPCs drain and ignores shutdownCtx, so a single hung or long-lived RPC // would otherwise block shutdown indefinitely, past the pod's termination grace period and into a SIGKILL. // This method bounds each stop by shutdownCtx, falling back to a hard Stop(). // The public and internal servers are stopped concurrently so total shutdown time is max(public, internal) rather // than their sum. The HTTP server is already drained above, which is required before stopping the grpc-web-wrapped // public server; the internal server has no such dependency. stopGraceful := func(name string, s *grpc.Server) { logger.Sugar().Infof("Stopping %s gRPC server...", name) done := make(chan struct{}) go func() { s.GracefulStop() close(done) }() select { case <-done: logger.Sugar().Infof("%s gRPC server stopped", name) case <-shutdownCtx.Done(): logger.Sugar().Warnf("%s gRPC server graceful stop timed out, forcing stop", name) s.Stop() } } var wg sync.WaitGroup wg.Go(func() { stopGraceful("public", grpcServer) }) if internalGrpcServer != nil { wg.Go(func() { stopGraceful("internal", internalGrpcServer) }) } wg.Wait() shutDownPprof(shutdownCtx, pprofServer, logger) if err := errGrp.Wait(); err != nil { logger.Error("Shutdown due to error", zap.Error(err)) } } func setUpPprof(errGrp *errgroup.Group, logger *zap.Logger) *http.Server { if os.Getenv("SPARK_PROFILING_ENABLED") != "true" { return nil } profilingPort := "6060" // default port if port := os.Getenv("SPARK_PROFILING_PORT"); port != "" { profilingPort = port } pprofServer := &http.Server{ Addr: net.JoinHostPort("127.0.0.1", profilingPort), // localhost only Handler: http.DefaultServeMux, // pprof endpoints are registered here } errGrp.Go(func() error { profilingLogger := logger.With(zap.String("component", "profiling")) profilingLogger.Info("Starting profiling server", zap.String("port", profilingPort)) if err := pprofServer.ListenAndServe(); !errors.Is(err, http.ErrServerClosed) { profilingLogger.Error("Profiling server failed", zap.Error(err)) return err } return nil }) return pprofServer } func shutDownPprof(shutdownCtx context.Context, pprofServer *http.Server, logger *zap.Logger) { if pprofServer != nil { logger.Info("Stopping profiling server...") if err := pprofServer.Shutdown(shutdownCtx); err != nil { logger.Error("Profiling server failed to shutdown gracefully", zap.Error(err)) return } logger.Info("Profiling server stopped") } } func setUpPyroscope(args *args, logger *zap.Logger) (shutDown func()) { pyroLogger := logger.With(zap.String("component", "profiling")) if len(args.PyroscopeServer) == 0 { pyroLogger.Info("No Pyroscope server specified; skipping") return func() {} } pyroLogger.Info("Connecting to Pyroscope server", zap.String("server", args.PyroscopeServer)) runtime.SetMutexProfileFraction(1000) runtime.SetBlockProfileRate(int(10 * time.Microsecond)) pyroLogger.With(zap.String("server", args.PyroscopeServer)) profiler, err := pyroscope.Start(pyroscope.Config{ ApplicationName: "spark", ServerAddress: args.PyroscopeServer, // e.g. "http://pyroscope.opentelemetry.svc.cluster.local:4040" Logger: pyroLogger.Sugar(), Tags: map[string]string{ "index": strconv.FormatUint(args.Index, 10), "supportedNetworks": args.SupportedNetworks, }, ProfileTypes: []pyroscope.ProfileType{ pyroscope.ProfileCPU, pyroscope.ProfileAllocObjects, pyroscope.ProfileAllocSpace, pyroscope.ProfileInuseObjects, pyroscope.ProfileInuseSpace, pyroscope.ProfileGoroutines, pyroscope.ProfileMutexCount, pyroscope.ProfileMutexDuration, pyroscope.ProfileBlockCount, pyroscope.ProfileBlockDuration, }, }) // This is only possible if our configuration is bad. if err != nil { pyroLogger.Error("Failed to connect to Pyroscope. Profiling data will not be stored.", zap.Error(err)) return func() {} } pyroLogger.Info("Connected to Pyroscope server", zap.String("server", args.PyroscopeServer)) return func() { if err := profiler.Stop(); err != nil { pyroLogger.Error("Failed to stop profiling server", zap.Error(err)) } } } func operatorPoolConfigFromKnobs(knobsService knobs.Knobs, operatorID string) so.OperatorConnectionPoolConfig { target := operatorID defaults := so.DefaultOperatorConnPoolConfig() // Helper to get int value from knob getInt := func(knob string, defaultVal int) int { return int(knobsService.GetValueTarget(knob, &target, float64(defaultVal))) } return so.OperatorConnectionPoolConfig{ MinConnections: getInt(knobs.KnobGrpcClientPoolMinConnections, defaults.MinConnections), MaxConnections: getInt(knobs.KnobGrpcClientPoolMaxConnections, defaults.MaxConnections), IdleTimeout: knobsService.GetDurationTarget(knobs.KnobGrpcClientPoolIdleTimeoutSeconds, &target, defaults.IdleTimeout), MaxLifetime: knobsService.GetDurationTarget(knobs.KnobGrpcClientPoolMaxLifetimeSeconds, &target, defaults.MaxLifetime), UsersPerConnectionCap: getInt(knobs.KnobGrpcClientPoolUsersPerConnectionCap, defaults.UsersPerConnectionCap), ScaleConcurrency: getInt(knobs.KnobGrpcClientPoolScaleConcurrency, defaults.ScaleConcurrency), } } // registeredMethodsMissingPolicy returns the full names of any gRPC methods registered on the server that lack an // rpcpolicy entry. Used as a fail-closed startup guard so a new RPC can't be exposed without declaring its auth policy. func registeredMethodsMissingPolicy(srv *grpc.Server) []string { var missing []string for serviceName, info := range srv.GetServiceInfo() { for _, m := range info.Methods { full := "/" + serviceName + "/" + m.Name if _, ok := rpcpolicy.LookUp(full); !ok { missing = append(missing, full) } } } slices.Sort(missing) return missing }