new-api-custom/service/group_availability.go

package service

import (
	"sync"
	"time"
)

// Group availability is derived from recent real request outcomes instead of an
// active probe. A group whose attempts keep failing is skipped during group
// selection until a cooldown elapses, after which it is eligible again so the
// next request acts as a probe.
const (
	groupAvailabilityMaxSamples     = 20
	groupAvailabilityMinSamples     = 4
	groupAvailabilityFailRatio      = 0.6
	groupAvailabilityFailureBurst   = 3
	groupAvailabilitySampleLifetime = 120 * time.Second
	groupAvailabilityCooldown       = 20 * time.Second
)

type groupOutcome struct {
	at      time.Time
	success bool
}

type groupAvailabilityState struct {
	mu       sync.Mutex
	outcomes []groupOutcome
}

var (
	groupAvailabilityMu    sync.RWMutex
	groupAvailabilityStates = make(map[string]*groupAvailabilityState)
)

func groupAvailabilityStateFor(group string) *groupAvailabilityState {
	groupAvailabilityMu.RLock()
	state, ok := groupAvailabilityStates[group]
	groupAvailabilityMu.RUnlock()
	if ok {
		return state
	}
	groupAvailabilityMu.Lock()
	defer groupAvailabilityMu.Unlock()
	if state, ok := groupAvailabilityStates[group]; ok {
		return state
	}
	state = &groupAvailabilityState{}
	groupAvailabilityStates[group] = state
	return state
}

// RecordGroupOutcome records the result of one upstream attempt for a group.
func RecordGroupOutcome(group string, success bool) {
	if group == "" {
		return
	}
	state := groupAvailabilityStateFor(group)
	state.mu.Lock()
	defer state.mu.Unlock()
	now := time.Now()
	state.outcomes = append(state.outcomes, groupOutcome{at: now, success: success})
	trimmed := state.outcomes[:0]
	for _, outcome := range state.outcomes {
		if now.Sub(outcome.at) <= groupAvailabilitySampleLifetime {
			trimmed = append(trimmed, outcome)
		}
	}
	state.outcomes = trimmed
	if len(state.outcomes) > groupAvailabilityMaxSamples {
		state.outcomes = state.outcomes[len(state.outcomes)-groupAvailabilityMaxSamples:]
	}
}

// IsGroupAvailable reports whether a group should be preferred by selection.
// A group without recorded traffic is available, so a fresh group is never
// skipped before it has been tried.
func IsGroupAvailable(group string) bool {
	if group == "" {
		return true
	}
	groupAvailabilityMu.RLock()
	state, ok := groupAvailabilityStates[group]
	groupAvailabilityMu.RUnlock()
	if !ok {
		return true
	}
	state.mu.Lock()
	defer state.mu.Unlock()
	now := time.Now()
	samples := make([]groupOutcome, 0, len(state.outcomes))
	for _, outcome := range state.outcomes {
		if now.Sub(outcome.at) <= groupAvailabilitySampleLifetime {
			samples = append(samples, outcome)
		}
	}
	if len(samples) == 0 {
		return true
	}
	last := samples[len(samples)-1]
	if last.success {
		return true
	}
	if now.Sub(last.at) >= groupAvailabilityCooldown {
		return true
	}
	burst := 0
	for i := len(samples) - 1; i >= 0; i-- {
		if samples[i].success {
			break
		}
		burst++
	}
	if burst >= groupAvailabilityFailureBurst {
		return false
	}
	if len(samples) < groupAvailabilityMinSamples {
		return true
	}
	failures := 0
	for _, sample := range samples {
		if !sample.success {
			failures++
		}
	}
	return float64(failures)/float64(len(samples)) < groupAvailabilityFailRatio
}

// ResetGroupAvailability clears all tracked state. Intended for tests.
func ResetGroupAvailability() {
	groupAvailabilityMu.Lock()
	defer groupAvailabilityMu.Unlock()
	groupAvailabilityStates = make(map[string]*groupAvailabilityState)
}