436 lines
13 KiB
Go
436 lines
13 KiB
Go
package health
|
|
|
|
import (
|
|
"context"
|
|
"crypto/rand"
|
|
"encoding/hex"
|
|
"reflect"
|
|
"sort"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
healthDomain "proxy-pool/internal/domain/health"
|
|
proxyDomain "proxy-pool/internal/domain/proxy"
|
|
)
|
|
|
|
var (
|
|
ErrInvalidTaskBroker = healthDomain.ErrInvalidTaskBroker
|
|
ErrInvalidTaskClaim = healthDomain.ErrInvalidTaskClaim
|
|
ErrInvalidLeasedTask = healthDomain.ErrInvalidLeasedTask
|
|
ErrTaskNotFound = healthDomain.ErrTaskNotFound
|
|
ErrTaskLeaseExpired = healthDomain.ErrTaskLeaseExpired
|
|
ErrTaskLeaseNotOwned = healthDomain.ErrTaskLeaseNotOwned
|
|
ErrTaskObservation = healthDomain.ErrTaskObservation
|
|
)
|
|
|
|
const defaultMaxTasksPerClaim = 128
|
|
|
|
type TaskMaterial = healthDomain.TaskMaterial
|
|
type TaskMaterialResolver = healthDomain.TaskMaterialResolver
|
|
type TaskMaterialResolverFunc = healthDomain.TaskMaterialResolverFunc
|
|
type TaskClaim = healthDomain.TaskClaim
|
|
type LeasedTask = healthDomain.LeasedTask
|
|
type TaskBroker = healthDomain.TaskBroker
|
|
|
|
type MemoryTaskBrokerOptions struct {
|
|
LeaseTTL time.Duration
|
|
MaxPerClaim int
|
|
Now func() time.Time
|
|
Material TaskMaterialResolver
|
|
}
|
|
|
|
// MemoryTaskBroker is a deterministic contract implementation for unit tests
|
|
// and single-process development. Production Controller instances use Redis so
|
|
// leases and limits remain shared across replicas.
|
|
type MemoryTaskBroker struct {
|
|
mu sync.Mutex
|
|
options MemoryTaskBrokerOptions
|
|
tasks map[string]brokerTask
|
|
leases map[string]taskLease
|
|
finished map[string]taskCompletion
|
|
}
|
|
|
|
type brokerTask struct {
|
|
plan PlannedTask
|
|
}
|
|
|
|
type taskLease struct {
|
|
checkerID string
|
|
instanceID string
|
|
token string
|
|
expiresAt time.Time
|
|
}
|
|
|
|
type taskCompletion struct {
|
|
checkerID string
|
|
token string
|
|
expiresAt time.Time
|
|
}
|
|
|
|
func NewMemoryTaskBroker(options MemoryTaskBrokerOptions) (*MemoryTaskBroker, error) {
|
|
if options.LeaseTTL <= 0 || options.Now == nil || nilTaskMaterialResolver(options.Material) {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
if options.MaxPerClaim == 0 {
|
|
options.MaxPerClaim = defaultMaxTasksPerClaim
|
|
}
|
|
if options.MaxPerClaim <= 0 {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
return &MemoryTaskBroker{
|
|
options: options, tasks: make(map[string]brokerTask), leases: make(map[string]taskLease),
|
|
finished: make(map[string]taskCompletion),
|
|
}, nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) Offer(ctx context.Context, plans []PlannedTask) (int, error) {
|
|
if ctx == nil || broker == nil || broker.options.Now == nil {
|
|
return 0, ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return 0, err
|
|
}
|
|
if len(plans) == 0 {
|
|
return 0, nil
|
|
}
|
|
now := broker.options.Now().UTC()
|
|
if now.IsZero() {
|
|
return 0, ErrInvalidTaskBroker
|
|
}
|
|
for _, plan := range plans {
|
|
if err := validatePlannedTask(plan, now); err != nil {
|
|
return 0, err
|
|
}
|
|
}
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
broker.reapLocked(now)
|
|
offered := 0
|
|
for _, plan := range plans {
|
|
taskID := deterministicTaskID(plan)
|
|
if _, exists := broker.tasks[taskID]; exists {
|
|
continue
|
|
}
|
|
broker.tasks[taskID] = brokerTask{plan: plan}
|
|
offered++
|
|
}
|
|
return offered, nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) Claim(ctx context.Context, claim TaskClaim) ([]LeasedTask, error) {
|
|
if ctx == nil || broker == nil || broker.options.Now == nil || nilTaskMaterialResolver(broker.options.Material) {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, err
|
|
}
|
|
levels, err := normalizeTaskClaim(claim, broker.options.MaxPerClaim)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
now := broker.options.Now().UTC()
|
|
if now.IsZero() {
|
|
return nil, ErrInvalidTaskBroker
|
|
}
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
broker.reapLocked(now)
|
|
capacity := claim.MaxInFlight - broker.activeForCheckerLocked(claim.CheckerID)
|
|
if capacity <= 0 {
|
|
return []LeasedTask{}, nil
|
|
}
|
|
if capacity > broker.options.MaxPerClaim {
|
|
capacity = broker.options.MaxPerClaim
|
|
}
|
|
taskIDs := broker.claimableTaskIDsLocked(levels)
|
|
if len(taskIDs) > capacity {
|
|
taskIDs = taskIDs[:capacity]
|
|
}
|
|
if len(taskIDs) == 0 {
|
|
return []LeasedTask{}, nil
|
|
}
|
|
materials := make([]TaskMaterial, len(taskIDs))
|
|
tokens := make([]string, len(taskIDs))
|
|
for index, taskID := range taskIDs {
|
|
material, resolveErr := broker.options.Material.ResolveCheckTask(ctx, broker.tasks[taskID].plan.Candidate)
|
|
if resolveErr != nil {
|
|
return nil, resolveErr
|
|
}
|
|
if validateTaskMaterial(material) != nil {
|
|
return nil, ErrInvalidLeasedTask
|
|
}
|
|
materials[index] = material
|
|
leaseToken, tokenErr := newLeaseToken()
|
|
if tokenErr != nil {
|
|
return nil, tokenErr
|
|
}
|
|
tokens[index] = leaseToken
|
|
}
|
|
result := make([]LeasedTask, len(taskIDs))
|
|
for index, taskID := range taskIDs {
|
|
plan := broker.tasks[taskID].plan
|
|
leaseUntil := now.Add(broker.options.LeaseTTL)
|
|
if plan.Deadline.Before(leaseUntil) {
|
|
leaseUntil = plan.Deadline
|
|
}
|
|
leaseToken := tokens[index]
|
|
broker.leases[taskID] = taskLease{checkerID: claim.CheckerID, instanceID: claim.InstanceID, token: leaseToken, expiresAt: leaseUntil}
|
|
result[index] = leasedTask(taskID, leaseToken, plan, materials[index])
|
|
}
|
|
return result, nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) AuthorizeObservation(
|
|
ctx context.Context,
|
|
checkerID string,
|
|
leaseToken string,
|
|
observation healthDomain.Observation,
|
|
now time.Time,
|
|
) error {
|
|
if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || !validTaskIdentifier(leaseToken) || now.IsZero() {
|
|
return ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := healthDomain.NormalizeObservation(observation)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
return broker.authorizeObservationLocked(checkerID, leaseToken, normalized, now.UTC())
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) CompleteObservation(
|
|
ctx context.Context,
|
|
checkerID string,
|
|
leaseToken string,
|
|
observation healthDomain.Observation,
|
|
now time.Time,
|
|
) error {
|
|
if ctx == nil || broker == nil || !validTaskIdentifier(checkerID) || !validTaskIdentifier(leaseToken) || now.IsZero() {
|
|
return ErrInvalidTaskBroker
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
normalized, err := healthDomain.NormalizeObservation(observation)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
now = now.UTC()
|
|
broker.mu.Lock()
|
|
defer broker.mu.Unlock()
|
|
if err := broker.authorizeObservationLocked(checkerID, leaseToken, normalized, now); err != nil {
|
|
return err
|
|
}
|
|
if _, complete := broker.finished[normalized.TaskID]; complete {
|
|
return nil
|
|
}
|
|
task := broker.tasks[normalized.TaskID]
|
|
delete(broker.leases, normalized.TaskID)
|
|
broker.finished[normalized.TaskID] = taskCompletion{checkerID: checkerID, token: leaseToken, expiresAt: task.plan.Deadline}
|
|
return nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) authorizeObservationLocked(
|
|
checkerID string,
|
|
leaseToken string,
|
|
observation healthDomain.Observation,
|
|
now time.Time,
|
|
) error {
|
|
broker.reapLocked(now)
|
|
task, exists := broker.tasks[observation.TaskID]
|
|
if !exists {
|
|
return ErrTaskNotFound
|
|
}
|
|
if !taskMatchesObservation(task.plan, observation) {
|
|
return ErrTaskObservation
|
|
}
|
|
if finished, complete := broker.finished[observation.TaskID]; complete {
|
|
if finished.checkerID != checkerID || finished.token != leaseToken {
|
|
return ErrTaskLeaseNotOwned
|
|
}
|
|
return nil
|
|
}
|
|
lease, leased := broker.leases[observation.TaskID]
|
|
if !leased || !lease.expiresAt.After(now) {
|
|
return ErrTaskLeaseExpired
|
|
}
|
|
if lease.checkerID != checkerID || lease.token != leaseToken {
|
|
return ErrTaskLeaseNotOwned
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) reapLocked(now time.Time) {
|
|
for taskID, task := range broker.tasks {
|
|
if !task.plan.Deadline.After(now) {
|
|
delete(broker.tasks, taskID)
|
|
delete(broker.leases, taskID)
|
|
delete(broker.finished, taskID)
|
|
}
|
|
}
|
|
for taskID, lease := range broker.leases {
|
|
if !lease.expiresAt.After(now) {
|
|
delete(broker.leases, taskID)
|
|
}
|
|
}
|
|
for taskID, complete := range broker.finished {
|
|
if !complete.expiresAt.After(now) {
|
|
delete(broker.finished, taskID)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) activeForCheckerLocked(checkerID string) int {
|
|
active := 0
|
|
for _, lease := range broker.leases {
|
|
if lease.checkerID == checkerID {
|
|
active++
|
|
}
|
|
}
|
|
return active
|
|
}
|
|
|
|
func (broker *MemoryTaskBroker) claimableTaskIDsLocked(levels map[healthDomain.Level]struct{}) []string {
|
|
result := make([]string, 0, len(broker.tasks))
|
|
for taskID, task := range broker.tasks {
|
|
if _, leased := broker.leases[taskID]; leased {
|
|
continue
|
|
}
|
|
if _, complete := broker.finished[taskID]; complete {
|
|
continue
|
|
}
|
|
if _, supported := levels[task.plan.Candidate.Level]; supported {
|
|
result = append(result, taskID)
|
|
}
|
|
}
|
|
sort.Slice(result, func(left, right int) bool {
|
|
leftTask, rightTask := broker.tasks[result[left]].plan, broker.tasks[result[right]].plan
|
|
if leftTask.Priority != rightTask.Priority {
|
|
return leftTask.Priority < rightTask.Priority
|
|
}
|
|
if !leftTask.Candidate.DueAt.Equal(rightTask.Candidate.DueAt) {
|
|
return leftTask.Candidate.DueAt.Before(rightTask.Candidate.DueAt)
|
|
}
|
|
return result[left] < result[right]
|
|
})
|
|
return result
|
|
}
|
|
|
|
func normalizeTaskClaim(claim TaskClaim, maximum int) (map[healthDomain.Level]struct{}, error) {
|
|
if !validTaskIdentifier(claim.CheckerID) || !validTaskIdentifier(claim.InstanceID) || claim.MaxInFlight <= 0 ||
|
|
claim.MaxInFlight > maximum || len(claim.SupportedLevels) == 0 {
|
|
return nil, ErrInvalidTaskClaim
|
|
}
|
|
levels := make(map[healthDomain.Level]struct{}, len(claim.SupportedLevels))
|
|
for _, level := range claim.SupportedLevels {
|
|
switch level {
|
|
case healthDomain.LevelBasic, healthDomain.LevelEgress, healthDomain.LevelTarget:
|
|
default:
|
|
return nil, ErrInvalidTaskClaim
|
|
}
|
|
if _, duplicate := levels[level]; duplicate {
|
|
return nil, ErrInvalidTaskClaim
|
|
}
|
|
levels[level] = struct{}{}
|
|
}
|
|
return levels, nil
|
|
}
|
|
|
|
func validatePlannedTask(task PlannedTask, now time.Time) error {
|
|
if task.Deadline.IsZero() || !task.Deadline.After(now) || task.Attempts <= 0 || !validTaskIdentifier(task.Candidate.ProxyID) ||
|
|
task.Candidate.DueAt.IsZero() {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
switch task.Candidate.State {
|
|
case proxyDomain.StateFetched, proxyDomain.StateSuspect, proxyDomain.StateUnhealthy, proxyDomain.StateAvailable:
|
|
default:
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
switch task.Candidate.Level {
|
|
case healthDomain.LevelBasic, healthDomain.LevelEgress:
|
|
if task.Candidate.RoutingName != "" || task.Candidate.TargetURL != "" {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
case healthDomain.LevelTarget:
|
|
if _, err := healthDomain.NormalizeTargetProfile(healthDomain.TargetProfile{
|
|
RoutingName: task.Candidate.RoutingName, TargetURL: task.Candidate.TargetURL,
|
|
}); err != nil {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
default:
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateTaskMaterial(material TaskMaterial) error {
|
|
if strings.TrimSpace(material.Host) != material.Host || material.Host == "" || material.Port == 0 ||
|
|
(material.SecretRef == "") != (material.CredentialVersion == "") {
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
switch material.Protocol {
|
|
case proxyDomain.SchemeHTTP, proxyDomain.SchemeHTTPS, proxyDomain.SchemeSOCKS5:
|
|
return nil
|
|
default:
|
|
return ErrInvalidLeasedTask
|
|
}
|
|
}
|
|
|
|
func deterministicTaskID(task PlannedTask) string {
|
|
return healthDomain.TaskIDFor(task)
|
|
}
|
|
|
|
func leasedTask(taskID, leaseToken string, plan PlannedTask, material TaskMaterial) LeasedTask {
|
|
return LeasedTask{
|
|
TaskID: taskID, LeaseToken: leaseToken, ProxyID: plan.Candidate.ProxyID, Protocol: material.Protocol, Host: material.Host, Port: material.Port,
|
|
SecretRef: material.SecretRef, CredentialVersion: material.CredentialVersion, Username: material.Username, Password: material.Password,
|
|
Level: plan.Candidate.Level, RoutingName: plan.Candidate.RoutingName, TargetURL: plan.Candidate.TargetURL,
|
|
Deadline: plan.Deadline.UTC(), Attempts: plan.Attempts,
|
|
}
|
|
}
|
|
|
|
func newLeaseToken() (string, error) {
|
|
var entropy [24]byte
|
|
if _, err := rand.Read(entropy[:]); err != nil {
|
|
return "", err
|
|
}
|
|
return "lease_" + hex.EncodeToString(entropy[:]), nil
|
|
}
|
|
|
|
func taskMatchesObservation(task PlannedTask, observation healthDomain.Observation) bool {
|
|
return task.Candidate.ProxyID == observation.ProxyID && task.Candidate.Level == observation.Level &&
|
|
task.Candidate.RoutingName == observation.RoutingName && task.Candidate.TargetURL == observation.TargetURL
|
|
}
|
|
|
|
func validTaskIdentifier(value string) bool {
|
|
if value == "" || len(value) > 256 || strings.TrimSpace(value) != value {
|
|
return false
|
|
}
|
|
for _, character := range value {
|
|
if character <= ' ' || character == '\x7f' {
|
|
return false
|
|
}
|
|
}
|
|
return true
|
|
}
|
|
|
|
func nilTaskMaterialResolver(value TaskMaterialResolver) 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
|
|
}
|
|
}
|
|
|
|
var _ TaskBroker = (*MemoryTaskBroker)(nil)
|