diff --git a/pkg/server/runner.go b/pkg/server/runner.go index b539ad44..73ec86e3 100644 --- a/pkg/server/runner.go +++ b/pkg/server/runner.go @@ -40,7 +40,7 @@ func NewRunner(opts *Options) *Runner { return &Runner{opts: opts} } -func (r *Runner) Run(ctx context.Context) error { +func (r *Runner) Run(ctx context.Context) (err error) { opts := r.opts logging.InitLogging(opts.LoggingOptions(), opts.Observability.Verbosity) @@ -49,177 +49,87 @@ func (r *Runner) Run(ctx context.Context) error { setupLog := ctrl.Log.WithName("setup") setupLog.Info("Logger initialized") - tracerShutdown, err := uotel.InitTracer(logr.NewContext(context.Background(), setupLog)) - if err != nil { - setupLog.Error(err, "Failed to initialize OpenTelemetry tracer") - return err - } + baseCtx := logr.NewContext(context.Background(), setupLog) + + var ( + healthServer *health.Server + gateFactory *flowcontrol.GateFactory + tracerShutdown func(ctx context.Context) error + ) defer func() { - shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + if err != nil { + setupLog.Error(err, "Runner failed") + } + shutdownCtx, cancel := context.WithTimeout(baseCtx, 5*time.Second) defer cancel() - if err := tracerShutdown(shutdownCtx); err != nil { - setupLog.Error(err, "Failed to shutdown tracer") + if healthServer != nil { + if err := healthServer.Shutdown(shutdownCtx); err != nil { + setupLog.Error(err, "Health server shutdown error") + } } - }() - - setupLog.Info("Async Processor starting", "version", version.Version, "commit", version.Commit, "buildDate", version.BuildDate) - - printAllFlags(setupLog) - - var workerPools []pipeline.WorkerPoolConfig - if opts.Worker.PoolConfigFile != "" { - workerPools, err = pipeline.LoadWorkerPools(opts.Worker.PoolConfigFile) - if err != nil { - setupLog.Error(err, "Failed to load pool configuration file") - return err + if gateFactory != nil { + if err := gateFactory.Close(); err != nil { + setupLog.Error(err, "Failed to close gate factory") + } } - setupLog.Info("Loaded named pools config", "count", len(workerPools)) - } else { - workerPools = []pipeline.WorkerPoolConfig{{ - ID: "default", - Workers: opts.Worker.Concurrency, - }} - setupLog.Info("No queue/pool configs set. Created default pool", "workers", opts.Worker.Concurrency) - } - - gateFactory := flowcontrol.NewGateFactoryWithCacheTTL(opts.Prometheus.URL, opts.Prometheus.CacheTTL) - defer func() { - if err := gateFactory.Close(); err != nil { - setupLog.Error(err, "Failed to close gate factory") + if tracerShutdown != nil { + if err := tracerShutdown(shutdownCtx); err != nil { + setupLog.Error(err, "Failed to shutdown tracer") + } } }() - var policy pipeline.RequestMergePolicy - switch opts.Queue.MergePolicy { - case "random-robin": - policy = async.NewRandomRobinPolicy() - default: - return fmt.Errorf("unknown request merge policy: %s", opts.Queue.MergePolicy) + tracerShutdown, err = initTracer(baseCtx) + if err != nil { + return err } - var impl pipeline.Flow - switch opts.Queue.Impl { - case "redis-pubsub": - flow, err := redis.NewRedisMQFlow(opts.Redis, opts.RedisConnection, redis.WithRedisTracing(opts.Observability.RedisTracing), redis.WithWorkerPools(workerPools)) - if err != nil { - setupLog.Error(err, "Failed to create Redis pub/sub flow") - return err - } - impl = flow - case "redis-sortedset": - flow, err := redis.NewRedisSortedSetFlow(opts.RedisSortedSet, opts.RedisConnection, redis.WithGateFactory(gateFactory), redis.WithSortedSetRedisTracing(opts.Observability.RedisTracing), redis.WithSortedSetWorkerPools(workerPools)) - if err != nil { - setupLog.Error(err, "Failed to create Redis sorted-set flow") - return err - } - impl = flow - setupLog.Info("Using Redis sorted-set flow with per-queue gating") - case "gcp-pubsub": - flow, err := pubsub.NewGCPPubSubMQFlow(opts.PubSub, pubsub.WithWorkerPools(workerPools)) - if err != nil { - setupLog.Error(err, "Failed to create GCP PubSub flow") - return err - } - impl = flow - case "gcp-pubsub-gated": - flow, err := pubsub.NewGCPPubSubMQFlow(opts.PubSub, pubsub.WithGateFactory(gateFactory), pubsub.WithWorkerPools(workerPools)) - if err != nil { - setupLog.Error(err, "Failed to create GCP PubSub gated flow") - return err - } - impl = flow - setupLog.Info("Using GCP PubSub flow with per-queue gating") - default: - return fmt.Errorf("unknown message queue implementation: %s", opts.Queue.Impl) - } + setupLog.Info("Async Processor starting", "version", version.Version, "commit", version.Commit, "buildDate", version.BuildDate) - metrics.Register(metrics.GetAsyncProcessorCollectors(impl.Characteristics().SupportsMessageLatency)...) + printAllFlags(setupLog) - var checker health.Checker - if hc, ok := impl.(pipeline.HealthChecker); ok { - checker = hc.HealthCheck - } - healthServer := health.NewServer(opts.Server.HealthPort, checker, setupLog.WithName("health")) - healthLn, err := healthServer.ListenAndServe() + poolsMap, totalConcurrency, err := loadWorkerPools(opts.Worker, setupLog) if err != nil { - setupLog.Error(err, "Failed to bind health server") return err } - go func() { - if err := healthServer.Serve(healthLn); err != nil { - setupLog.Error(err, "Health server failed") - } - }() - - signalCtx := ctx - drainCtx, drainCancel := context.WithCancel(logr.NewContext(context.Background(), setupLog)) - defer drainCancel() + gateFactory = flowcontrol.NewGateFactoryWithCacheTTL(opts.Prometheus.URL, opts.Prometheus.CacheTTL) - metricsServerOptions := metricsserver.Options{ - BindAddress: fmt.Sprintf(":%d", opts.Server.MetricsPort), - FilterProvider: func() func(c *rest.Config, httpClient *http.Client) (metricsserver.Filter, error) { - if opts.Server.MetricsEndpointAuth { - return filters.WithAuthenticationAndAuthorization - } - return nil - }(), + policy, err := loadRequestMergePolicy(opts.Queue.MergePolicy) + if err != nil { + return err } - restConfig := ctrl.GetConfigOrDie() - msrv, err := metricsserver.NewServer(metricsServerOptions, restConfig, http.DefaultClient) + flow, err := loadFlow(opts, gateFactory, poolsMap) if err != nil { - setupLog.Error(err, "Failed to create metrics server") return err } - go msrv.Start(signalCtx) //nolint:errcheck - tlsConfig, err := buildTLSConfig(opts.TLS) + metrics.Register(metrics.GetAsyncProcessorCollectors(flow.Characteristics().SupportsMessageLatency)...) + + healthServer, err = initHealthServer(flow, opts.Server, setupLog) if err != nil { - setupLog.Error(err, "Failed to build TLS configuration") return err } - totalConcurrency := 0 - poolsMap := make(map[string]pipeline.WorkerPoolConfig) - for _, p := range workerPools { - if p.Workers <= 0 { - p.Workers = opts.Worker.Concurrency - } - poolsMap[p.ID] = p - totalConcurrency += p.Workers - metrics.SetPoolWorkerLimit(p.ID, float64(p.Workers)) + if err = startMetricsServer(ctx, opts.Server, setupLog); err != nil { + return err } - inferenceTransport := &http.Transport{ - MaxIdleConns: 100, - MaxIdleConnsPerHost: totalConcurrency, - IdleConnTimeout: 90 * time.Second, - TLSClientConfig: tlsConfig, + inferenceClient, err := initInferenceClient(opts.TLS, totalConcurrency) + if err != nil { + return err } - inferenceHTTPClient := &http.Client{Transport: otelhttp.NewTransport(inferenceTransport, - otelhttp.WithSpanNameFormatter(func(_ string, _ *http.Request) string { - return "http-request" - }), - )} - inferenceClient := asyncworker.NewHTTPInferenceClient(inferenceHTTPClient) - var transforms *transform.Chain - if opts.TransformConfigFile != "" { - cfg, err := transform.LoadConfig(opts.TransformConfigFile) - if err != nil { - setupLog.Error(err, "Failed to load transform configuration file") - return err - } - transforms, err = transform.BuildChain(cfg.RequestTransforms, plugins.NewHandle(signalCtx)) - if err != nil { - setupLog.Error(err, "Failed to build request transform chain") - return err - } - setupLog.Info("Loaded request transform plugins", "count", transforms.Len()) + transforms, err := loadTransforms(ctx, opts.TransformConfigFile, setupLog) + if err != nil { + return err } - dispatch := policy.MergeRequestChannels(impl.RequestChannels(), poolsMap) + drainCtx, drainCancel := context.WithCancel(baseCtx) + defer drainCancel() + + dispatch := policy.MergeRequestChannels(flow.RequestChannels(), poolsMap) poolGates := make(map[string]pipeline.Gate) for poolID, pool := range poolsMap { @@ -236,10 +146,7 @@ func (r *Runner) Run(ctx context.Context) error { var wg sync.WaitGroup for poolID, mergedChan := range dispatch.Channels { - pool, ok := poolsMap[poolID] - if !ok { - return fmt.Errorf("pool %s not found", poolID) - } + pool := poolsMap[poolID] workersCount := pool.Workers poolGate := poolGates[poolID] @@ -248,25 +155,25 @@ func (r *Runner) Run(ctx context.Context) error { wg.Add(1) go func(mergedChan chan pipeline.EmbelishedRequestMessage, poolGate pipeline.Gate) { defer wg.Done() - asyncworker.WorkerWithGate(signalCtx, drainCtx, impl.Characteristics(), inferenceClient, mergedChan, impl.RetryChannel(), impl.ResultChannel(), opts.Worker.RequestTimeout, transforms, poolGate) + asyncworker.WorkerWithGate(ctx, drainCtx, flow.Characteristics(), inferenceClient, mergedChan, flow.RetryChannel(), flow.ResultChannel(), opts.Worker.RequestTimeout, transforms, poolGate) }(mergedChan, poolGate) } } - impl.Start(signalCtx) + flow.Start(ctx) healthServer.SetReady() - if reporter, ok := impl.(pipeline.BacklogReporter); ok && opts.Queue.BacklogPollInterval > 0 { - go pollBacklog(signalCtx, reporter, opts.Queue.BacklogPollInterval) + if reporter, ok := flow.(pipeline.BacklogReporter); ok && opts.Queue.BacklogPollInterval > 0 { + go pollBacklog(ctx, reporter, opts.Queue.BacklogPollInterval) } else if !ok { setupLog.Info("Selected flow does not support broker backlog metrics", "message-queue-impl", opts.Queue.Impl) } - <-signalCtx.Done() + <-ctx.Done() healthServer.SetNotReady() setupLog.Info("Signal received, stopping message consumption") - impl.StopConsuming() + flow.StopConsuming() setupLog.Info("Draining in-flight requests", "timeout", opts.Worker.DrainTimeout) done := make(chan struct{}) @@ -280,17 +187,153 @@ func (r *Runner) Run(ctx context.Context) error { wg.Wait() } - impl.Shutdown() + flow.Shutdown() + + return nil +} + +func initTracer(baseCtx context.Context) (func(context.Context) error, error) { + shutdown, err := uotel.InitTracer(baseCtx) + if err != nil { + logr.FromContextOrDiscard(baseCtx).Error(err, "Failed to initialize OpenTelemetry tracer") + return nil, err + } + return shutdown, nil +} + +func loadWorkerPools(workerConfig WorkerConfig, setupLog logr.Logger) (poolsMap map[string]pipeline.WorkerPoolConfig, totalConcurrency int, err error) { + var pools []pipeline.WorkerPoolConfig + if workerConfig.PoolConfigFile != "" { + pools, err = pipeline.LoadWorkerPools(workerConfig.PoolConfigFile) + if err != nil { + return nil, -1, err + } + setupLog.Info("Loaded named pools config", "count", len(pools)) + } else { + pools = []pipeline.WorkerPoolConfig{{ + ID: "default", + Workers: workerConfig.Concurrency, + }} + setupLog.Info("No queue/pool configs set. Created default pool", "workers", workerConfig.Concurrency) + } + + poolsMap = make(map[string]pipeline.WorkerPoolConfig) + for _, p := range pools { + if p.Workers <= 0 { + p.Workers = workerConfig.Concurrency + } + poolsMap[p.ID] = p + totalConcurrency += p.Workers + metrics.SetPoolWorkerLimit(p.ID, float64(p.Workers)) + } + return poolsMap, totalConcurrency, err + +} + +func loadRequestMergePolicy(name string) (pipeline.RequestMergePolicy, error) { + switch name { + case "random-robin": + return async.NewRandomRobinPolicy(), nil + default: + return nil, fmt.Errorf("unknown request merge policy: %s", name) + } +} + +func loadFlow(opts *Options, gateFactory *flowcontrol.GateFactory, poolsMap map[string]pipeline.WorkerPoolConfig) (pipeline.Flow, error) { + workerPools := make([]pipeline.WorkerPoolConfig, 0, len(poolsMap)) + for _, p := range poolsMap { + workerPools = append(workerPools, p) + } + switch opts.Queue.Impl { + case "redis-pubsub": + return redis.NewRedisMQFlow(opts.Redis, opts.RedisConnection, redis.WithRedisTracing(opts.Observability.RedisTracing), redis.WithWorkerPools(workerPools)) + case "redis-sortedset": + return redis.NewRedisSortedSetFlow(opts.RedisSortedSet, opts.RedisConnection, redis.WithGateFactory(gateFactory), redis.WithSortedSetRedisTracing(opts.Observability.RedisTracing), redis.WithSortedSetWorkerPools(workerPools)) + case "gcp-pubsub": + return pubsub.NewGCPPubSubMQFlow(opts.PubSub, pubsub.WithWorkerPools(workerPools)) + case "gcp-pubsub-gated": + return pubsub.NewGCPPubSubMQFlow(opts.PubSub, pubsub.WithGateFactory(gateFactory), pubsub.WithWorkerPools(workerPools)) + default: + return nil, fmt.Errorf("unknown message queue implementation: %s", opts.Queue.Impl) + } +} + +func initHealthServer(impl pipeline.Flow, serverCfg ServerConfig, setupLog logr.Logger) (*health.Server, error) { + var checker health.Checker + if hc, ok := impl.(pipeline.HealthChecker); ok { + checker = hc.HealthCheck + } + healthServer := health.NewServer(serverCfg.HealthPort, checker, setupLog.WithName("health")) + healthLn, err := healthServer.ListenAndServe() + if err != nil { + setupLog.Error(err, "Failed to bind health server") + return nil, err + } + go func() { + if err := healthServer.Serve(healthLn); err != nil { + setupLog.Error(err, "Health server failed") + } + }() + return healthServer, nil +} - shutdownCtx, shutdownCancel := context.WithTimeout(context.Background(), 5*time.Second) - defer shutdownCancel() - if err := healthServer.Shutdown(shutdownCtx); err != nil { - setupLog.Error(err, "Health server shutdown error") +func startMetricsServer(ctx context.Context, serverCfg ServerConfig, setupLog logr.Logger) error { + metricsServerOptions := metricsserver.Options{ + BindAddress: fmt.Sprintf(":%d", serverCfg.MetricsPort), + FilterProvider: func() func(c *rest.Config, httpClient *http.Client) (metricsserver.Filter, error) { + if serverCfg.MetricsEndpointAuth { + return filters.WithAuthenticationAndAuthorization + } + return nil + }(), } + restConfig := ctrl.GetConfigOrDie() + msrv, err := metricsserver.NewServer(metricsServerOptions, restConfig, http.DefaultClient) + if err != nil { + setupLog.Error(err, "Failed to create metrics server") + return err + } + go msrv.Start(ctx) //nolint:errcheck return nil } +func initInferenceClient(tlsCfg TLSConfig, totalConcurrency int) (*asyncworker.HTTPInferenceClient, error) { + tlsConfig, err := buildTLSConfig(tlsCfg) + if err != nil { + return nil, fmt.Errorf("failed to build TLS configuration: %w", err) + } + + inferenceTransport := &http.Transport{ + MaxIdleConns: 100, + MaxIdleConnsPerHost: totalConcurrency, + IdleConnTimeout: 90 * time.Second, + TLSClientConfig: tlsConfig, + } + inferenceHTTPClient := &http.Client{Transport: otelhttp.NewTransport(inferenceTransport, + otelhttp.WithSpanNameFormatter(func(_ string, _ *http.Request) string { + return "http-request" + }), + )} + return asyncworker.NewHTTPInferenceClient(inferenceHTTPClient), nil +} + +func loadTransforms(ctx context.Context, configFile string, setupLog logr.Logger) (*transform.Chain, error) { + if configFile == "" { + return nil, nil + } + cfg, err := transform.LoadConfig(configFile) + if err != nil { + return nil, fmt.Errorf("failed to load transform configuration: %w", err) + } + chain, err := transform.BuildChain(cfg.RequestTransforms, plugins.NewHandle(ctx)) + if err != nil { + return nil, fmt.Errorf("failed to build request transform chain: %w", err) + } + setupLog.Info("Loaded request transform plugins", "count", chain.Len()) + return chain, nil +} + func buildTLSConfig(cfg TLSConfig) (*tls.Config, error) { if cfg.CACert == "" && cfg.Cert == "" && cfg.Key == "" && !cfg.InsecureSkipVerify { return nil, nil