package health import ( "context" "errors" "sort" "time" "proxy-pool/internal/config" "proxy-pool/internal/domain/activitypool" "proxy-pool/internal/domain/adminstate" healthDomain "proxy-pool/internal/domain/health" ) var ErrInvalidUpstreamDrainReaper = errors.New("invalid upstream drain reaper") // UpstreamDrainConfigurationSource returns one configuration and revision from // the same atomic publication. config.Store implements this interface. type UpstreamDrainConfigurationSource interface { Snapshot() (*config.Config, uint64) } type UpstreamDrainStateSource interface { Snapshot(context.Context) (adminstate.Snapshot, error) } // EffectiveUpstreamDrainPolicySource merges static configuration with optional // Admin state. With Admin enabled it uses the same revision equality rule as // Worker routing, so a torn config/state view produces no policy publication. type EffectiveUpstreamDrainPolicySource struct { configuration UpstreamDrainConfigurationSource state UpstreamDrainStateSource } func NewEffectiveUpstreamDrainPolicySource( configuration UpstreamDrainConfigurationSource, states ...UpstreamDrainStateSource, ) (*EffectiveUpstreamDrainPolicySource, error) { if nilInterface(configuration) || len(states) > 1 || (len(states) == 1 && nilInterface(states[0])) { return nil, ErrInvalidUpstreamDrainReaper } source := &EffectiveUpstreamDrainPolicySource{configuration: configuration} if len(states) == 1 { source.state = states[0] } return source, nil } func (source *EffectiveUpstreamDrainPolicySource) ReadUpstreamDrainPolicies( ctx context.Context, ) ([]activitypool.UpstreamDrainPolicy, error) { if ctx == nil || source == nil || nilInterface(source.configuration) { return nil, ErrInvalidUpstreamDrainReaper } if err := ctx.Err(); err != nil { return nil, err } configuration, revision := source.configuration.Snapshot() if configuration == nil { return nil, ErrInvalidUpstreamDrainReaper } states := make(map[string]adminstate.UpstreamState, len(configuration.Upstreams)) if source.state != nil { snapshot, err := source.state.Snapshot(ctx) if err != nil { return nil, err } if snapshot.Config == nil || snapshot.Config.Revision != revision { return nil, ErrInvalidUpstreamDrainReaper } for _, upstream := range snapshot.Upstreams { if upstream.Name == "" || upstream.Revision == 0 { return nil, ErrInvalidUpstreamDrainReaper } if _, duplicate := states[upstream.Name]; duplicate { return nil, ErrInvalidUpstreamDrainReaper } states[upstream.Name] = upstream } } if revision == 0 { // A config-only Controller has no persisted management revision. Its // configuration cannot be hot-reloaded, so one stable non-zero fence is enough. revision = 1 } names := make([]string, 0, len(configuration.Upstreams)) for name := range configuration.Upstreams { names = append(names, name) } sort.Strings(names) policies := make([]activitypool.UpstreamDrainPolicy, 0, len(names)) for _, name := range names { upstream := configuration.Upstreams[name] policy := activitypool.UpstreamDrainPolicy{UpstreamID: name, Revision: revision, Enabled: upstream.Enabled} if source.state != nil { admin, exists := states[name] if !exists { return nil, ErrInvalidUpstreamDrainReaper } policy.Enabled = policy.Enabled && admin.Enabled policy.Revision = admin.Revision } policies = append(policies, policy) } return policies, nil } type UpstreamDrainPolicySource interface { ReadUpstreamDrainPolicies(context.Context) ([]activitypool.UpstreamDrainPolicy, error) } // UpstreamDrainStore is the narrow activity-pool capability required by the // configuration-driven Drain reaper. type UpstreamDrainStore interface { activitypool.UpstreamDrainPolicyWriter activitypool.DisabledUpstreamDrainLister activitypool.DisabledUpstreamDrainStarter } type UpstreamDrainReaperOptions struct { PollInterval time.Duration BatchSize int Now func() time.Time Metrics healthDomain.DrainMetricsObserver } type UpstreamDrainResult struct { Candidates int Started int } // ConfiguredUpstreamDrainReaper continually publishes an effective, complete // policy view before reading one bounded batch of disabled Upstream ownership. // The Redis ownership operation performs the final policy-revision fence. type ConfiguredUpstreamDrainReaper struct { policies UpstreamDrainPolicySource store UpstreamDrainStore options UpstreamDrainReaperOptions } func NewConfiguredUpstreamDrainReaper( policies UpstreamDrainPolicySource, store UpstreamDrainStore, options UpstreamDrainReaperOptions, ) (*ConfiguredUpstreamDrainReaper, error) { if nilInterface(policies) || nilInterface(store) || options.PollInterval <= 0 || options.BatchSize <= 0 || options.Now == nil { return nil, ErrInvalidUpstreamDrainReaper } return &ConfiguredUpstreamDrainReaper{policies: policies, store: store, options: options}, nil } func (reaper *ConfiguredUpstreamDrainReaper) Tick(ctx context.Context) (UpstreamDrainResult, error) { if ctx == nil || reaper == nil || nilInterface(reaper.policies) || nilInterface(reaper.store) || reaper.options.PollInterval <= 0 || reaper.options.BatchSize <= 0 || reaper.options.Now == nil { return UpstreamDrainResult{}, ErrInvalidUpstreamDrainReaper } if err := ctx.Err(); err != nil { return UpstreamDrainResult{}, err } now := reaper.options.Now() if now.IsZero() { return UpstreamDrainResult{}, ErrInvalidUpstreamDrainReaper } policies, err := reaper.policies.ReadUpstreamDrainPolicies(ctx) if err != nil { return UpstreamDrainResult{}, err } if err := reaper.store.ReplaceUpstreamDrainPolicies(ctx, policies); err != nil { return UpstreamDrainResult{}, err } result := UpstreamDrainResult{} remaining := reaper.options.BatchSize utcNow := now.UTC() for _, policy := range policies { if policy.Enabled || remaining == 0 { continue } candidates, err := reaper.store.ListDisabledUpstreamDrainCandidates(ctx, utcNow, policy, remaining) if err != nil { return result, err } if len(candidates) > remaining { return result, ErrInvalidUpstreamDrainReaper } for _, candidate := range candidates { started, err := reaper.store.BeginDisabledUpstreamDrain(ctx, utcNow, candidate) if err != nil { return result, err } result.Candidates++ if started { result.Started++ } } remaining -= len(candidates) } if reaper.options.Metrics != nil { reaper.options.Metrics.ObserveDrain(healthDomain.DrainReasonUpstreamDisabled, result.Candidates, result.Started) } return result, nil } func (reaper *ConfiguredUpstreamDrainReaper) Run(ctx context.Context) error { if reaper == nil { return ErrInvalidUpstreamDrainReaper } return runScheduler(ctx, reaper.options.PollInterval, func(tickCtx context.Context) (TickResult, error) { _, err := reaper.Tick(tickCtx) return TickResult{}, err }) }