proxy-pool/internal/controller/bootstrap/bootstrap.go

package bootstrap

import (
	"context"
	"errors"
	"fmt"
	"reflect"
	"sort"
	"strings"
	"time"

	"github.com/prometheus/client_golang/prometheus"

	"proxy-pool/internal/config"
	"proxy-pool/internal/controller/admin"
	"proxy-pool/internal/controller/distribution"
	"proxy-pool/internal/controller/extraction"
	controllerHealth "proxy-pool/internal/controller/health"
	"proxy-pool/internal/controller/operations"
	"proxy-pool/internal/controller/pool"
	"proxy-pool/internal/controller/provider"
	controllerRouting "proxy-pool/internal/controller/routing"
	controllerRuntime "proxy-pool/internal/controller/runtime"
	"proxy-pool/internal/controller/worker"
	"proxy-pool/internal/domain/activitypool"
	extractionDomain "proxy-pool/internal/domain/extraction"
	healthDomain "proxy-pool/internal/domain/health"
	ownershipDomain "proxy-pool/internal/domain/ownership"
	"proxy-pool/internal/domain/workerruntime"
	"proxy-pool/internal/platform/admission"
	"proxy-pool/internal/platform/credentials"
	"proxy-pool/internal/platform/httpserver"
	"proxy-pool/internal/platform/lifecycle"
	platformMetrics "proxy-pool/internal/platform/metrics"
)

var (
	ErrInvalidOptions = errors.New("invalid controller bootstrap options")
	ErrStartup        = errors.New("controller startup failed")
)

const (
	checkSchedulerPollInterval = 250 * time.Millisecond
	checkSchedulerBatchSize    = 128
)

type Options struct {
	ConfigPath     string
	Resolver       config.Resolver
	Now            func() time.Time
	HTTP           httpserver.Options
	HolderID       string
	RedisNamespace string
	FingerprintKey []byte
}

type activityStore interface {
	extractionDomain.Store
	activitypool.Upserter
	activitypool.GlobalHealthStore
	activitypool.TargetHealthStore
	activitypool.ProxyUpstreamReader
	pool.InventoryReader
	activitypool.StateInventoryReader
}

type healthTaskRuntime interface {
	controllerHealth.TaskBroker
	controllerHealth.UpstreamTaskSource
}

type ports struct {
	state            admin.StateRepository
	activity         activityStore
	readiness        distribution.ReadinessChecker
	metricsReadiness platformMetrics.ReadinessChecker
	admission        admission.Admitter
	coordinator      provider.Coordinator
	credentials      credentials.Store
	providerResults  provider.ResultRecorder
	capacityMetrics  pool.CapacityObserver
	workerStore      workerruntime.ControlStore
	close            func() error
}

type infrastructure interface {
	Open(context.Context, *config.Config) (ports, error)
}

type controllerRunner interface {
	Run(context.Context) error
}

type runtimeFactory interface {
	New(*config.Config, controllerRuntime.Dependencies, controllerRuntime.Options) (controllerRunner, error)
}

type workerRuntimeFactory interface {
	New(config.ControlPlane, worker.Service, worker.ServerOptions) (controllerRunner, error)
}

func Run(ctx context.Context, options Options) error {
	return runWithWorkerFactory(ctx, options, &productionInfrastructure{
		holderID: options.HolderID, namespace: options.RedisNamespace,
	}, productionRuntimeFactory{}, productionWorkerRuntimeFactory{})
}

func run(ctx context.Context, options Options, infrastructure infrastructure, factory runtimeFactory) (resultErr error) {
	return runWithWorkerFactory(ctx, options, infrastructure, factory, productionWorkerRuntimeFactory{})
}

func runWithWorkerFactory(
	ctx context.Context,
	options Options,
	infrastructure infrastructure,
	factory runtimeFactory,
	workerFactory workerRuntimeFactory,
) (resultErr error) {
	if ctx == nil || strings.TrimSpace(options.ConfigPath) != options.ConfigPath || options.ConfigPath == "" ||
		nilInterface(options.Resolver) || nilInterface(infrastructure) || nilInterface(factory) || nilInterface(workerFactory) {
		return ErrInvalidOptions
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if options.Now == nil {
		options.Now = time.Now
	}

	loader, err := admin.NewFileConfigurationLoader(options.ConfigPath, options.Resolver)
	if err != nil {
		return errors.Join(ErrInvalidOptions, err)
	}
	loaded, err := loader.LoadConfiguration(ctx)
	if err != nil {
		return fmt.Errorf("%w: load configuration: %w", ErrStartup, err)
	}
	if loaded.Value.Admin.Enabled && len(options.FingerprintKey) < config.MinimumFingerprintKeyBytes {
		return errors.Join(ErrInvalidOptions, config.ErrInvalidFingerprint)
	}
	configurationStore, err := config.NewStore(loaded.Value)
	if err != nil {
		return fmt.Errorf("%w: initialize configuration store: %w", ErrStartup, err)
	}

	opened, err := infrastructure.Open(ctx, loaded.Value)
	if err != nil {
		return fmt.Errorf("%w: open infrastructure: %w", ErrStartup, err)
	}
	if opened.close == nil {
		return errors.Join(ErrStartup, ErrInvalidOptions)
	}
	defer func() {
		resultErr = errors.Join(resultErr, opened.close())
	}()
	var providerState providerStateReader
	if loaded.Value.Admin.Enabled {
		providerState = opened.state
	}
	supervisor, err := newProviderSupervisor(
		configurationStore,
		providerState,
		func(name string, upstream config.Upstream) (lifecycle.Runner, error) {
			buildRuntime, err := providerRuntimeBuilder(opened)
			if err != nil {
				return nil, err
			}
			return buildRuntime(name, upstream)
		},
		func(ctx context.Context, configuration *config.Config) error {
			return prepareProviderConfiguration(ctx, configuration, opened.credentials)
		},
		func(configuration *config.Config) { retainProviderStats(configuration, opened.providerResults) },
		loader,
		options.FingerprintKey,
		providerSupervisorInterval,
	)
	if err != nil {
		return fmt.Errorf("%w: build Provider supervisor: %w", ErrStartup, err)
	}
	snapshotRefresh := worker.NewSnapshotRefreshBroker()
	var sequentialCoordinator lifecycle.Runner

	dependencies := controllerRuntime.Dependencies{}
	var checkerMetrics healthDomain.TaskMetricsObserver
	var drainMetrics healthDomain.DrainMetricsObserver
	var extractionMetrics extraction.MetricsObserver
	if loaded.Value.Metrics.Enabled {
		if nilInterface(opened.metricsReadiness) {
			return errors.Join(ErrStartup, ErrInvalidOptions)
		}
		collector, collectorErr := platformMetrics.NewCheckerCollector(prometheus.DefaultRegisterer)
		if collectorErr != nil {
			return fmt.Errorf("%w: build Checker metrics: %w", ErrStartup, collectorErr)
		}
		checkerMetrics = collector
		drainCollector, drainCollectorErr := platformMetrics.NewDrainCollector(prometheus.DefaultRegisterer)
		if drainCollectorErr != nil {
			return fmt.Errorf("%w: build Drain metrics: %w", ErrStartup, drainCollectorErr)
		}
		drainMetrics = drainCollector
		providerCollector, providerCollectorErr := platformMetrics.NewProviderCollector(prometheus.DefaultRegisterer)
		if providerCollectorErr != nil {
			return fmt.Errorf("%w: build Provider metrics: %w", ErrStartup, providerCollectorErr)
		}
		if registrar, ok := opened.providerResults.(provider.ResultObserverRegistrar); ok && !nilInterface(registrar) {
			registrar.AddResultObserver(providerCollector)
		}
		capacityCollector, capacityCollectorErr := platformMetrics.NewCapacityCollector(prometheus.DefaultRegisterer)
		if capacityCollectorErr != nil {
			return fmt.Errorf("%w: build capacity metrics: %w", ErrStartup, capacityCollectorErr)
		}
		opened.capacityMetrics = capacityCollector
		extractionCollector, extractionCollectorErr := platformMetrics.NewExtractionCollector(prometheus.DefaultRegisterer)
		if extractionCollectorErr != nil {
			return fmt.Errorf("%w: build extraction metrics: %w", ErrStartup, extractionCollectorErr)
		}
		extractionMetrics = extractionCollector
		handler, handlerErr := platformMetrics.NewHandler(platformMetrics.Dependencies{
			Gatherer: prometheus.DefaultGatherer, Readiness: opened.metricsReadiness,
		})
		if handlerErr != nil {
			return fmt.Errorf("%w: build metrics handler: %w", ErrStartup, handlerErr)
		}
		dependencies.MetricsHandler = handler
	}
	if loaded.Value.Distribution.Enabled {
		if nilInterface(opened.activity) || nilInterface(opened.readiness) || nilInterface(opened.admission) {
			return errors.Join(ErrStartup, ErrInvalidOptions)
		}
		service, serviceErr := extraction.NewService(
			opened.activity,
			extractionPolicy(loaded.Value),
			opened.admission,
			options.Now,
			extractionMetrics,
		)
		if serviceErr != nil {
			return fmt.Errorf("%w: build extraction service: %w", ErrStartup, serviceErr)
		}
		dependencies.Extractor = service
		dependencies.Readiness = opened.readiness
	}
	if loaded.Value.Admin.Enabled {
		if nilInterface(opened.state) || nilInterface(opened.activity) {
			return errors.Join(ErrStartup, ErrInvalidOptions)
		}
		var providerStats []provider.StatsReader
		if stats, ok := opened.providerResults.(provider.StatsReader); ok && !nilInterface(stats) {
			providerStats = append(providerStats, stats)
		}
		statusReader, statusErr := operations.NewReader(
			configurationStore,
			opened.activity,
			options.Now,
			providerStats...,
		)
		if statusErr != nil {
			return fmt.Errorf("%w: build operational status reader: %w", ErrStartup, statusErr)
		}
		service, serviceErr := admin.NewApplicationService(admin.ApplicationDependencies{
			State: opened.state, Operations: statusReader, Configuration: loader, Publisher: configurationStore,
			Runtime: supervisor, SnapshotRefresh: snapshotRefresh,
		}, admin.ApplicationOptions{Now: options.Now, FingerprintKey: options.FingerprintKey})
		if serviceErr != nil {
			return fmt.Errorf("%w: build admin service: %w", ErrStartup, serviceErr)
		}
		if _, applyErr := service.ApplyConfiguration(ctx, admin.ReloadCommand{
			RequestID: "controller-startup", ActorID: "proxy-controller",
		}, loaded); applyErr != nil {
			return fmt.Errorf("%w: commit startup configuration: %w", ErrStartup, applyErr)
		}
		dependencies.AdminService = service
		if stats, ok := opened.providerResults.(provider.StatsReader); ok && !nilInterface(stats) {
			coordinator, coordinatorErr := controllerRouting.NewSequentialCoordinator(
				configurationStore,
				opened.state,
				stats,
				snapshotRefresh,
				controllerRouting.Options{Now: options.Now},
			)
			if coordinatorErr != nil {
				return fmt.Errorf("%w: build Sequential routing coordinator: %w", ErrStartup, coordinatorErr)
			}
			if registrar, ok := opened.providerResults.(provider.ResultObserverRegistrar); ok && !nilInterface(registrar) {
				registrar.AddResultObserver(coordinator)
			}
			sequentialCoordinator = coordinator
		}
	}
	runners := make([]lifecycle.Runner, 0, 3+len(loaded.Value.Upstreams))
	if hasHTTPRuntime(loaded.Value) {
		runner, err := factory.New(configurationStore.Current(), dependencies, controllerRuntime.Options{HTTP: options.HTTP})
		if err != nil {
			return fmt.Errorf("%w: build HTTP runtime: %w", ErrStartup, err)
		}
		if nilInterface(runner) {
			return errors.Join(ErrStartup, ErrInvalidOptions)
		}
		runners = append(runners, runner)
	}
	if loaded.Value.ControlPlane.Enabled {
		if nilInterface(opened.workerStore) || nilInterface(opened.activity) {
			return errors.Join(ErrStartup, ErrInvalidOptions)
		}
		var routingSources []worker.RoutingStateReader
		if loaded.Value.Admin.Enabled {
			if nilInterface(opened.state) {
				return errors.Join(ErrStartup, ErrInvalidOptions)
			}
			routingSources = append(routingSources, opened.state)
		}
		routingSource, routingErr := worker.NewGatewayRoutingSource(configurationStore, routingSources...)
		if routingErr != nil {
			return fmt.Errorf("%w: build Worker routing source: %w", ErrStartup, routingErr)
		}
		var snapshotReader ownershipDomain.SnapshotReader
		if reader, ok := opened.workerStore.(ownershipDomain.SnapshotReader); ok {
			snapshotReader = reader
		}
		var drainTickets ownershipDomain.DrainTicketStore
		if tickets, ok := opened.workerStore.(ownershipDomain.DrainTicketStore); ok {
			drainTickets = tickets
		}
		service, serviceErr := worker.NewService(opened.workerStore, worker.Options{
			ProtocolVersion:    loaded.Value.ControlPlane.ProtocolVersion,
			HeartbeatInterval:  loaded.Value.ControlPlane.HeartbeatInterval.Value(),
			SessionTTL:         loaded.Value.ControlPlane.SessionTTL.Value(),
			MaxStaleAge:        loaded.Value.ControlPlane.MaxStaleAge.Value(),
			MaxRuntimeCounters: loaded.Value.ControlPlane.MaxRuntimeCounters,
			MaxSnapshotBytes:   loaded.Value.ControlPlane.MaxMessageBytes,
			SnapshotReader:     snapshotReader,
			DrainTickets:       drainTickets,
			RoutingSource:      routingSource,
			Credentials:        opened.credentials,
			SnapshotRefresh:    snapshotRefresh,
		})
		if serviceErr != nil {
			return fmt.Errorf("%w: build Worker control service: %w", ErrStartup, serviceErr)
		}
		threshold, thresholdErr := controllerHealth.NewConfiguredFailureThresholdResolver(configurationStore, opened.activity, options.Now)
		if thresholdErr != nil {
			return fmt.Errorf("%w: build Checker failure threshold resolver: %w", ErrStartup, thresholdErr)
		}
		reducer, reducerErr := controllerHealth.NewReducer(opened.activity, opened.activity, threshold)
		if reducerErr != nil {
			return fmt.Errorf("%w: build Checker health reducer: %w", ErrStartup, reducerErr)
		}
		checkerIdentity, identityErr := newCheckerIdentity(loaded.Value.ControlPlane)
		if identityErr != nil {
			return fmt.Errorf("%w: build Checker identity authorizer: %w", ErrStartup, identityErr)
		}
		checkerOptions := controllerHealth.DefaultGRPCHandlerOptions()
		checkerOptions.Metrics = checkerMetrics
		if tasks, ok := opened.activity.(healthTaskRuntime); ok && !nilInterface(tasks) {
			checkerOptions.TaskBroker = tasks
		}
		checkerHandler, handlerErr := controllerHealth.NewGRPCHandler(reducer, checkerIdentity, checkerOptions)
		if handlerErr != nil {
			return fmt.Errorf("%w: build Checker control handler: %w", ErrStartup, handlerErr)
		}
		serverOptions := worker.DefaultServerOptions()
		serverOptions.Checker = checkerHandler
		runner, runnerErr := workerFactory.New(loaded.Value.ControlPlane, service, serverOptions)
		if runnerErr != nil {
			return fmt.Errorf("%w: build Worker control server: %w", ErrStartup, runnerErr)
		}
		if nilInterface(runner) {
			return errors.Join(ErrStartup, ErrInvalidOptions)
		}
		runners = append(runners, runner)
		if tasks, ok := opened.activity.(healthTaskRuntime); ok && !nilInterface(tasks) {
			schedulers, schedulerErr := newHealthSchedulers(configurationStore, tasks, options.Now)
			if schedulerErr != nil {
				return fmt.Errorf("%w: build Checker health schedulers: %w", ErrStartup, schedulerErr)
			}
			runners = append(runners, schedulers...)
		}
		if unhealthy, ok := opened.activity.(activitypool.UnhealthyRemover); ok && !nilInterface(unhealthy) {
			reaper, reaperErr := controllerHealth.NewConfiguredUnhealthyReaper(
				configurationStore,
				unhealthy,
				controllerHealth.UnhealthyReaperOptions{
					PollInterval: checkSchedulerPollInterval, BatchSize: checkSchedulerBatchSize, Now: options.Now, Metrics: drainMetrics,
				},
			)
			if reaperErr != nil {
				return fmt.Errorf("%w: build unhealthy reaper: %w", ErrStartup, reaperErr)
			}
			runners = append(runners, reaper)
		}
		if drainStore, ok := opened.activity.(controllerHealth.UpstreamDrainStore); ok && !nilInterface(drainStore) {
			var policyStates []controllerHealth.UpstreamDrainStateSource
			if loaded.Value.Admin.Enabled {
				if nilInterface(opened.state) {
					return errors.Join(ErrStartup, ErrInvalidOptions)
				}
				policyStates = append(policyStates, opened.state)
			}
			policySource, policyErr := controllerHealth.NewEffectiveUpstreamDrainPolicySource(configurationStore, policyStates...)
			if policyErr != nil {
				return fmt.Errorf("%w: build upstream Drain policy source: %w", ErrStartup, policyErr)
			}
			reaper, reaperErr := controllerHealth.NewConfiguredUpstreamDrainReaper(
				policySource,
				drainStore,
				controllerHealth.UpstreamDrainReaperOptions{
					PollInterval: checkSchedulerPollInterval, BatchSize: checkSchedulerBatchSize, Now: options.Now, Metrics: drainMetrics,
				},
			)
			if reaperErr != nil {
				return fmt.Errorf("%w: build upstream Drain reaper: %w", ErrStartup, reaperErr)
			}
			runners = append(runners, reaper)
		}
	}
	if sequentialCoordinator != nil {
		runners = append(runners, sequentialCoordinator)
	}
	runners = append(runners, supervisor)
	group, err := lifecycle.NewGroup(runners...)
	if err != nil {
		return fmt.Errorf("%w: build process lifecycle: %w", ErrStartup, err)
	}
	return group.Run(ctx)
}

func newHealthSchedulers(
	configuration controllerHealth.ConfigurationSource,
	tasks healthTaskRuntime,
	now func() time.Time,
) ([]lifecycle.Runner, error) {
	if nilInterface(configuration) || nilInterface(tasks) || now == nil {
		return nil, ErrInvalidOptions
	}
	current := configuration.Current()
	if current == nil {
		return nil, ErrInvalidOptions
	}
	hasEnabledUpstream := false
	for _, upstream := range current.Upstreams {
		if upstream.Enabled {
			hasEnabledUpstream = true
			break
		}
	}
	if !hasEnabledUpstream {
		return nil, ErrInvalidOptions
	}
	supervisor, err := controllerHealth.NewConfiguredSchedulerSupervisor(configuration, tasks, tasks, controllerHealth.SchedulerRunnerOptions{
		PollInterval: checkSchedulerPollInterval, BatchSize: checkSchedulerBatchSize, Now: now,
	})
	if err != nil {
		return nil, err
	}
	return []lifecycle.Runner{supervisor}, nil
}

func newCheckerIdentity(controlPlane config.ControlPlane) (controllerHealth.CheckerIdentityAuthorizer, error) {
	switch controlPlane.TLS.Mode {
	case "disabled":
		return worker.AllowLoopbackIdentity{}, nil
	case "mtls":
		return worker.NewSPIFFEIdentityAuthorizer(controlPlane.TLS.TrustDomain, controlPlane.TLS.Environment)
	default:
		return nil, ErrInvalidOptions
	}
}

func prepareProviderConfiguration(
	ctx context.Context,
	configuration *config.Config,
	credentialStore credentials.Store,
) error {
	if ctx == nil || configuration == nil {
		return ErrProviderRuntime
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	if !configuration.Admin.Enabled && !hasEnabledUpstream(configuration) {
		return nil
	}
	ensurer, ok := credentialStore.(credentials.CapacityEnsurer)
	if !ok || nilInterface(credentialStore) {
		return ErrProviderRuntime
	}
	if err := ensurer.EnsureCapacity(ctx, providerCredentialCapacity(configuration)); err != nil {
		return err
	}
	return nil
}

func retainProviderStats(configuration *config.Config, results provider.ResultRecorder) {
	if configuration == nil {
		return
	}
	retainer, ok := results.(provider.StatsRetainer)
	if !ok || nilInterface(retainer) {
		return
	}
	names := make([]string, 0, len(configuration.Upstreams))
	for name := range configuration.Upstreams {
		names = append(names, name)
	}
	sort.Strings(names)
	retainer.RetainProviderStats(names)
}

func hasHTTPRuntime(configuration *config.Config) bool {
	return configuration != nil &&
		(configuration.Distribution.Enabled || configuration.Admin.Enabled || configuration.Metrics.Enabled)
}

func extractionPolicy(configuration *config.Config) extraction.Policy {
	configured := configuration.Distribution.Extraction
	return extraction.Policy{
		MaxCountPerRequest: configured.MaxCountPerRequest,
		DefaultFulfillment: extractionDomain.Fulfillment(configured.Fulfillment),
		MinRemainingTTL:    configured.MinRemainingTTL.Value(),
		MaxHealthCheckAge:  configured.MaxHealthCheckAge.Value(),
		ReserveForGateway:  configured.ReserveForGateway,
		IdempotencyTTL:     configured.IdempotencyTTL.Value(),
	}
}

func nilInterface(value any) bool {
	if value == nil {
		return true
	}
	reflected := reflect.ValueOf(value)
	switch reflected.Kind() {
	case reflect.Chan, reflect.Func, reflect.Interface, reflect.Map, reflect.Pointer, reflect.Slice:
		return reflected.IsNil()
	default:
		return false
	}
}

type productionRuntimeFactory struct{}

func (productionRuntimeFactory) New(
	configuration *config.Config,
	dependencies controllerRuntime.Dependencies,
	options controllerRuntime.Options,
) (controllerRunner, error) {
	return controllerRuntime.New(configuration, dependencies, options)
}

type productionWorkerRuntimeFactory struct{}

func (productionWorkerRuntimeFactory) New(
	controlPlane config.ControlPlane,
	service worker.Service,
	options worker.ServerOptions,
) (controllerRunner, error) {
	return worker.NewServer(controlPlane, service, options)
}