proxy-pool/internal/controller/health/upstream_drain_reaper.go

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
	})
}