proxy-pool/internal/controller/provider/upstream_runtime.go

package provider

import (
	"context"
	"errors"
	"fmt"
	"hash/fnv"
	"reflect"
	"sync/atomic"
	"time"

	"proxy-pool/internal/controller/pool"
	"proxy-pool/internal/domain/activitypool"
)

var (
	ErrInvalidUpstreamRuntime = errors.New("invalid Provider upstream runtime")
	ErrUpstreamRuntimeStopped = errors.New("Provider upstream runtime stopped")
)

type UpstreamRuntimeConfig struct {
	Provider          Config
	ReconcilePolicy   pool.ReconcilePolicy
	ReconcileInterval time.Duration
}

type UpstreamRuntimeDependencies struct {
	Coordinator Coordinator
	Inventory   pool.InventoryReader
	Adapter     ProviderAdapter
	Parser      Parser
	Activity    activitypool.Upserter
	Results     ResultRecorder
	Capacity    pool.CapacityObserver
}

// UpstreamRuntime owns every leader-scoped object for one Upstream. A fresh
// local budget and coalescing signal are created for each leadership term.
type UpstreamRuntime struct {
	config       UpstreamRuntimeConfig
	dependencies UpstreamRuntimeDependencies
	sleeper      Sleeper
	termSequence atomic.Uint64
}

func NewUpstreamRuntime(
	config UpstreamRuntimeConfig,
	dependencies UpstreamRuntimeDependencies,
) (*UpstreamRuntime, error) {
	if config.ReconcileInterval <= 0 || nilRuntimeDependency(dependencies.Coordinator) ||
		nilRuntimeDependency(dependencies.Inventory) || nilRuntimeDependency(dependencies.Adapter) ||
		nilRuntimeDependency(dependencies.Parser) || nilRuntimeDependency(dependencies.Activity) ||
		nilRuntimeDependency(dependencies.Results) {
		return nil, ErrInvalidUpstreamRuntime
	}
	runtime := &UpstreamRuntime{
		config: config, dependencies: dependencies, sleeper: timerSleeper{},
	}
	if _, err := runtime.newLeaderTerm(); err != nil {
		return nil, errors.Join(ErrInvalidUpstreamRuntime, err)
	}
	return runtime, nil
}

func (runtime *UpstreamRuntime) Run(ctx context.Context) error {
	if runtime == nil || ctx == nil {
		return ErrInvalidUpstreamRuntime
	}
	limits := CoordinationLimits{
		RequestInterval:    runtime.config.Provider.RequestInterval,
		MaxInFlight:        runtime.config.Provider.MaxInFlight,
		MaxAttemptDuration: runtime.config.Provider.Timeout,
		MaxTotal:           runtime.config.Provider.MaxTotal,
	}
	err := runtime.dependencies.Coordinator.RunLeader(
		ctx,
		runtime.config.Provider.UpstreamID,
		limits,
		func(leaderCtx context.Context, session LeaderSession) error {
			term, buildErr := runtime.newLeaderTerm()
			if buildErr != nil {
				return buildErr
			}
			return term.run(leaderCtx, session)
		},
	)
	if ctx.Err() != nil {
		return ctx.Err()
	}
	if err == nil {
		return ErrUpstreamRuntimeStopped
	}
	return err
}

func (runtime *UpstreamRuntime) ID() string {
	if runtime == nil {
		return ""
	}
	return runtime.config.Provider.UpstreamID
}

func (runtime *UpstreamRuntime) newLeaderTerm() (*upstreamLeaderTerm, error) {
	budget, err := pool.NewFetchBudget(pool.FetchBudgetConfig{
		UpstreamID:       runtime.config.Provider.UpstreamID,
		MaxSize:          runtime.config.Provider.MaxSize,
		ExpectedPerFetch: runtime.config.ReconcilePolicy.ExpectedPerFetch,
	})
	if err != nil {
		return nil, err
	}
	providerConfig := runtime.config.Provider
	providerConfig.RequestInterval = 0
	providerReconciler, err := NewReconciler(providerConfig, Ports{
		Adapter: runtime.dependencies.Adapter, Parser: runtime.dependencies.Parser,
		Activity: runtime.dependencies.Activity, Results: runtime.dependencies.Results,
		Capacity: budget,
	})
	if err != nil {
		return nil, err
	}
	poolReconciler, err := pool.NewReconciler(
		runtime.config.ReconcilePolicy,
		budget,
		providerReconciler,
	)
	if err != nil {
		return nil, err
	}
	return &upstreamLeaderTerm{
		upstreamID:   runtime.config.Provider.UpstreamID,
		sourceID:     fmt.Sprintf("%p:%d", runtime, runtime.termSequence.Add(1)),
		interval:     runtime.config.ReconcileInterval,
		safetyMargin: runtime.config.ReconcilePolicy.SafetyMargin,
		inventory:    runtime.dependencies.Inventory,
		pool:         poolReconciler,
		provider:     providerReconciler,
		capacity:     runtime.dependencies.Capacity,
		sleeper:      runtime.sleeper,
	}, nil
}

type upstreamLeaderTerm struct {
	upstreamID   string
	sourceID     string
	interval     time.Duration
	safetyMargin time.Duration
	inventory    pool.InventoryReader
	pool         *pool.Reconciler
	provider     *Reconciler
	capacity     pool.CapacityObserver
	sleeper      Sleeper
}

func (term *upstreamLeaderTerm) run(ctx context.Context, session LeaderSession) error {
	if !nilRuntimeDependency(term.capacity) {
		defer term.capacity.RemoveCapacityUpstream(term.upstreamID, term.sourceID)
	}
	termCtx, cancel := context.WithCancel(ctx)
	defer cancel()
	done := make(chan error, 2)
	go func() { done <- term.provider.RunLeader(termCtx, session) }()
	go func() { done <- term.reconcileInventory(termCtx) }()

	first := <-done
	cancel()
	second := <-done
	if ctx.Err() != nil {
		return nil
	}
	if first == nil {
		first = ErrUpstreamRuntimeStopped
	}
	if second != nil && !errors.Is(second, context.Canceled) {
		return errors.Join(first, second)
	}
	return first
}

func (term *upstreamLeaderTerm) reconcileInventory(ctx context.Context) error {
	if delay := initialReconcileDelay(term.upstreamID, term.interval); delay > 0 {
		if err := term.sleeper.Sleep(ctx, delay); err != nil {
			if ctx.Err() != nil {
				return nil
			}
			return fmt.Errorf("wait for initial Provider inventory reconciliation: %w", err)
		}
	}
	for ctx.Err() == nil {
		inventory, err := term.inventory.ReadInventory(
			ctx,
			term.upstreamID,
			term.safetyMargin,
		)
		if err == nil && inventory.Managed >= 0 && inventory.AvailableSlots >= 0 {
			decision := term.pool.ReconcileSnapshot(inventory)
			term.observeCapacity(pool.CapacityObservation{
				UpstreamID: term.upstreamID, SourceID: term.sourceID, Result: pool.CapacityReadSuccess,
				Managed: inventory.Managed, AvailableSlots: inventory.AvailableSlots,
				EffectiveSlots: decision.EffectiveSlots, PendingExpected: decision.PendingExpected,
			})
		} else {
			term.observeCapacity(pool.CapacityObservation{
				UpstreamID: term.upstreamID, SourceID: term.sourceID, Result: pool.CapacityReadError,
			})
		}
		if err := term.sleeper.Sleep(ctx, term.interval); err != nil {
			if ctx.Err() != nil {
				return nil
			}
			return fmt.Errorf("wait for Provider inventory reconciliation: %w", err)
		}
	}
	return nil
}

func (term *upstreamLeaderTerm) observeCapacity(observation pool.CapacityObservation) {
	if term == nil || nilRuntimeDependency(term.capacity) {
		return
	}
	term.capacity.ObserveCapacity(observation)
}

func initialReconcileDelay(upstreamID string, interval time.Duration) time.Duration {
	window := min(interval/4, 250*time.Millisecond)
	if window <= 1 {
		return 0
	}
	digest := fnv.New64a()
	_, _ = digest.Write([]byte(upstreamID))
	return time.Duration(digest.Sum64() % uint64(window))
}

func nilRuntimeDependency(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
	}
}