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

package bootstrap

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

	"proxy-pool/internal/config"
	"proxy-pool/internal/controller/admin"
	"proxy-pool/internal/domain/adminstate"
	"proxy-pool/internal/platform/lifecycle"
)

const providerSupervisorInterval = time.Second

var (
	ErrProviderSupervisor                 = errors.New("invalid Provider supervisor")
	errProviderManagementStateUnavailable = errors.New("Provider management state unavailable")
	errProviderManagementSnapshotStale    = errors.New("Provider management snapshot stale")
	errProviderConfigurationPending       = errors.New("Provider configuration synchronization pending")
)

type providerConfigurationStore interface {
	Current() *config.Config
	Revision() uint64
	PublishRevision(*config.Config, uint64) bool
}

type providerConfigurationSource interface {
	LoadConfiguration(context.Context) (admin.LoadedConfiguration, error)
}

type providerStateReader interface {
	Snapshot(context.Context) (adminstate.Snapshot, error)
}

type providerRunnerBuilder func(string, config.Upstream) (lifecycle.Runner, error)
type providerConfigurationPreparer func(context.Context, *config.Config) error
type providerConfigurationObserver func(*config.Config)

type providerSupervisor struct {
	configuration  providerConfigurationStore
	state          providerStateReader
	build          providerRunnerBuilder
	prepare        providerConfigurationPreparer
	observe        providerConfigurationObserver
	source         providerConfigurationSource
	fingerprintKey []byte
	interval       time.Duration
	notify         chan struct{}
}

type runningProvider struct {
	configuration config.Upstream
	cancel        context.CancelFunc
	done          chan struct{}
}

func newProviderSupervisor(
	configuration providerConfigurationStore,
	state providerStateReader,
	build providerRunnerBuilder,
	prepare providerConfigurationPreparer,
	observe providerConfigurationObserver,
	source providerConfigurationSource,
	fingerprintKey []byte,
	interval time.Duration,
) (*providerSupervisor, error) {
	if nilInterface(configuration) || build == nil || interval <= 0 ||
		(!nilInterface(state) && len(fingerprintKey) < config.MinimumFingerprintKeyBytes) {
		return nil, ErrProviderSupervisor
	}
	return &providerSupervisor{
		configuration:  configuration,
		state:          state,
		build:          build,
		prepare:        prepare,
		observe:        observe,
		source:         source,
		fingerprintKey: append([]byte(nil), fingerprintKey...),
		interval:       interval,
		notify:         make(chan struct{}, 1),
	}, nil
}

func (supervisor *providerSupervisor) Notify() {
	if supervisor == nil || supervisor.notify == nil {
		return
	}
	select {
	case supervisor.notify <- struct{}{}:
	default:
	}
}

func (supervisor *providerSupervisor) ValidateConfiguration(ctx context.Context, configuration *config.Config) error {
	if supervisor == nil || ctx == nil || configuration == nil || supervisor.build == nil {
		return ErrProviderSupervisor
	}
	if err := supervisor.validateConfiguration(ctx, configuration); err != nil {
		return errors.Join(ErrProviderSupervisor, err)
	}
	return nil
}

func (supervisor *providerSupervisor) validateConfiguration(ctx context.Context, configuration *config.Config) error {
	if err := ctx.Err(); err != nil {
		return err
	}
	if err := config.Validate(configuration); err != nil {
		return err
	}
	if supervisor.prepare != nil {
		if err := supervisor.prepare(ctx, configuration); err != nil {
			return err
		}
	}
	names := make([]string, 0, len(configuration.Upstreams))
	for name, upstream := range configuration.Upstreams {
		if upstream.Enabled {
			names = append(names, name)
		}
	}
	sort.Strings(names)
	for _, name := range names {
		if err := ctx.Err(); err != nil {
			return err
		}
		if _, err := supervisor.build(name, configuration.Upstreams[name]); err != nil {
			return err
		}
	}
	return nil
}

func (supervisor *providerSupervisor) ValidateUpstream(ctx context.Context, name string) error {
	if supervisor == nil || ctx == nil || name == "" || supervisor.build == nil {
		return ErrProviderSupervisor
	}
	if err := ctx.Err(); err != nil {
		return err
	}
	configuration := supervisor.configuration.Current()
	if configuration == nil {
		return ErrProviderSupervisor
	}
	if supervisor.prepare != nil {
		if err := supervisor.prepare(ctx, configuration); err != nil {
			return errors.Join(ErrProviderSupervisor, err)
		}
	}
	upstream, exists := configuration.Upstreams[name]
	if !exists {
		return ErrProviderSupervisor
	}
	upstream.Enabled = true
	if _, err := supervisor.build(name, upstream); err != nil {
		return errors.Join(ErrProviderSupervisor, err)
	}
	return nil
}

func (supervisor *providerSupervisor) Run(ctx context.Context) error {
	if supervisor == nil || ctx == nil || nilInterface(supervisor.configuration) ||
		supervisor.build == nil || supervisor.interval <= 0 || supervisor.notify == nil {
		return ErrProviderSupervisor
	}
	active := make(map[string]*runningProvider)
	failures := make(chan error, 1)
	defer stopAllProviders(active)
	ticker := time.NewTicker(supervisor.interval)
	defer ticker.Stop()

	for {
		if err := supervisor.reconcile(ctx, active, failures); err != nil {
			return err
		}
		select {
		case <-ctx.Done():
			return ctx.Err()
		case err := <-failures:
			return err
		case <-supervisor.notify:
		case <-ticker.C:
		}
	}
}

func (supervisor *providerSupervisor) reconcile(
	ctx context.Context,
	active map[string]*runningProvider,
	failures chan<- error,
) error {
	desired, err := supervisor.desired(ctx)
	if err != nil {
		if errors.Is(err, errProviderManagementStateUnavailable) ||
			errors.Is(err, errProviderManagementSnapshotStale) {
			return nil
		}
		if errors.Is(err, errProviderConfigurationPending) {
			stopAllProviders(active)
			clear(active)
			return nil
		}
		return err
	}
	names := make([]string, 0, len(desired))
	prepared := make(map[string]lifecycle.Runner)
	for name, upstream := range desired {
		names = append(names, name)
		current := active[name]
		if current != nil && reflect.DeepEqual(current.configuration, upstream) {
			continue
		}
		runner, buildErr := supervisor.build(name, upstream)
		if buildErr != nil || nilInterface(runner) {
			return errors.Join(ErrProviderSupervisor, buildErr)
		}
		prepared[name] = runner
	}
	sort.Strings(names)

	for name, current := range active {
		if _, keep := desired[name]; !keep {
			stopProvider(current)
			delete(active, name)
		}
	}
	for _, name := range names {
		runner := prepared[name]
		if runner == nil {
			continue
		}
		if current := active[name]; current != nil {
			stopProvider(current)
		}
		active[name] = startProvider(ctx, name, desired[name], runner, failures)
	}
	return nil
}

func (supervisor *providerSupervisor) desired(ctx context.Context) (map[string]config.Upstream, error) {
	configuration := supervisor.configuration.Current()
	if configuration == nil {
		return nil, ErrProviderSupervisor
	}
	enabled := make(map[string]bool, len(configuration.Upstreams))
	var snapshot adminstate.Snapshot
	if !nilInterface(supervisor.state) {
		var err error
		snapshot, err = supervisor.state.Snapshot(ctx)
		if err != nil {
			return nil, errors.Join(errProviderManagementStateUnavailable, err)
		}
		configuration, err = supervisor.synchronizeConfiguration(ctx, configuration, snapshot)
		if err != nil {
			return nil, err
		}
		for _, upstream := range snapshot.Upstreams {
			enabled[upstream.Name] = upstream.Enabled
		}
	}
	if supervisor.prepare != nil {
		if err := supervisor.prepare(ctx, configuration); err != nil {
			return nil, errors.Join(ErrProviderSupervisor, err)
		}
	}
	if supervisor.observe != nil {
		supervisor.observe(configuration)
	}
	desired := make(map[string]config.Upstream)
	for name, upstream := range configuration.Upstreams {
		isEnabled := upstream.Enabled
		if !nilInterface(supervisor.state) {
			isEnabled = enabled[name]
		}
		if isEnabled {
			upstream.Enabled = true
			desired[name] = upstream
		}
	}
	return desired, nil
}

func (supervisor *providerSupervisor) synchronizeConfiguration(
	ctx context.Context,
	current *config.Config,
	snapshot adminstate.Snapshot,
) (*config.Config, error) {
	if snapshot.Config == nil || snapshot.Config.Checksum == "" {
		return current, nil
	}
	if supervisor.configuration.Revision() > snapshot.Config.Revision {
		return nil, errProviderManagementSnapshotStale
	}
	checksum, err := config.Fingerprint(current, supervisor.fingerprintKey)
	if err != nil {
		return nil, errors.Join(ErrProviderSupervisor, err)
	}
	if checksum == snapshot.Config.Checksum {
		supervisor.configuration.PublishRevision(current, snapshot.Config.Revision)
		return current, nil
	}
	if nilInterface(supervisor.source) {
		return nil, errProviderConfigurationPending
	}
	loaded, err := supervisor.source.LoadConfiguration(ctx)
	if err != nil || loaded.Value == nil {
		return nil, errors.Join(errProviderConfigurationPending, err)
	}
	checksum, err = config.Fingerprint(loaded.Value, supervisor.fingerprintKey)
	if err != nil || checksum != snapshot.Config.Checksum {
		return nil, errors.Join(errProviderConfigurationPending, err)
	}
	if err := supervisor.validateConfiguration(ctx, loaded.Value); err != nil {
		return nil, errors.Join(errProviderConfigurationPending, err)
	}
	if supervisor.configuration.PublishRevision(loaded.Value, snapshot.Config.Revision) {
		return loaded.Value, nil
	}
	if supervisor.configuration.Revision() > snapshot.Config.Revision {
		return nil, errProviderManagementSnapshotStale
	}
	return nil, errProviderConfigurationPending
}

func startProvider(
	ctx context.Context,
	name string,
	configuration config.Upstream,
	runner lifecycle.Runner,
	failures chan<- error,
) *runningProvider {
	runCtx, cancel := context.WithCancel(ctx)
	running := &runningProvider{configuration: configuration, cancel: cancel, done: make(chan struct{})}
	go func() {
		defer close(running.done)
		err := runner.Run(runCtx)
		if runCtx.Err() != nil {
			return
		}
		if err == nil {
			err = lifecycle.ErrRunnerStopped
		}
		select {
		case failures <- fmt.Errorf("Provider %s runtime: %w", name, err):
		default:
		}
	}()
	return running
}

func stopProvider(running *runningProvider) {
	if running == nil {
		return
	}
	running.cancel()
	<-running.done
}

func stopAllProviders(active map[string]*runningProvider) {
	var wait sync.WaitGroup
	for _, running := range active {
		wait.Add(1)
		go func(current *runningProvider) {
			defer wait.Done()
			stopProvider(current)
		}(running)
	}
	wait.Wait()
}