easyai-ai-gateway/apps/api/internal/store/async_worker_capacity.go

package store

import (
	"context"
	"fmt"
)

type AsyncWorkerCapacitySnapshot struct {
	Capacity             int
	GlobalCapacity       int
	Desired              int
	HardLimit            int
	Capped               bool
	EnabledModels        int
	UnlimitedModels      int
	EnabledGroups        int
	UnlimitedGroups      int
	ModelDesired         int
	GroupDesired         int
	ActiveInstances      int
	InstanceID           string
	LoadMode             string
	LocalSafeCapacity    int
	LocalHeavyCapacity   int
	LocalActiveTasks     int
	LocalPreparingTasks  int
	LocalWaitingTasks    int
	LocalFinalizingTasks int
	LocalPressureState   string
	LocalPressureReason  string
}

func (s *Store) AsyncWorkerCapacity(ctx context.Context, hardLimit int) (AsyncWorkerCapacitySnapshot, error) {
	if hardLimit < 1 {
		return AsyncWorkerCapacitySnapshot{}, fmt.Errorf("async worker hard limit must be positive")
	}
	rows, err := s.pool.Query(ctx, `
SELECT COALESCE(b.default_rate_limit_policy, '{}'::jsonb),
       p.rate_limit_policy,
       COALESCE(rp.rate_limit_policy, '{}'::jsonb),
       (m.runtime_policy_set_id IS NOT NULL),
       COALESCE(m.runtime_policy_override, '{}'::jsonb),
       m.rate_limit_policy,
       m.rate_limit_policy_mode
FROM platform_models m
JOIN integration_platforms p ON p.id = m.platform_id
LEFT JOIN base_model_catalog b ON b.id = m.base_model_id
LEFT JOIN model_runtime_policy_sets rp ON rp.id = COALESCE(m.runtime_policy_set_id, b.runtime_policy_set_id)
WHERE p.status = 'enabled'
  AND p.deleted_at IS NULL
  AND m.enabled = true`)
	if err != nil {
		return AsyncWorkerCapacitySnapshot{}, err
	}
	defer rows.Close()

	modelPolicies := make([]map[string]any, 0)
	for rows.Next() {
		var basePolicyBytes, platformPolicyBytes, runtimePolicyBytes []byte
		var runtimeOverrideBytes, modelPolicyBytes []byte
		var runtimeExplicit bool
		var modelPolicyMode string
		if err := rows.Scan(
			&basePolicyBytes,
			&platformPolicyBytes,
			&runtimePolicyBytes,
			&runtimeExplicit,
			&runtimeOverrideBytes,
			&modelPolicyBytes,
			&modelPolicyMode,
		); err != nil {
			return AsyncWorkerCapacitySnapshot{}, err
		}
		modelPolicies = append(modelPolicies, EffectiveRateLimitPolicy(EffectiveRateLimitPolicyInput{
			BasePolicy:            decodeObject(basePolicyBytes),
			PlatformPolicy:        decodeObject(platformPolicyBytes),
			RuntimePolicy:         decodeObject(runtimePolicyBytes),
			RuntimePolicyExplicit: runtimeExplicit,
			RuntimePolicyOverride: decodeObject(runtimeOverrideBytes),
			ModelPolicy:           decodeObject(modelPolicyBytes),
			ModelPolicyMode:       modelPolicyMode,
		}))
	}
	if err := rows.Err(); err != nil {
		return AsyncWorkerCapacitySnapshot{}, err
	}
	groupRows, err := s.pool.Query(ctx, `
SELECT rate_limit_policy
FROM gateway_user_groups
WHERE status = 'active'`)
	if err != nil {
		return AsyncWorkerCapacitySnapshot{}, err
	}
	defer groupRows.Close()
	groupPolicies := make([]map[string]any, 0)
	for groupRows.Next() {
		var policyBytes []byte
		if err := groupRows.Scan(&policyBytes); err != nil {
			return AsyncWorkerCapacitySnapshot{}, err
		}
		groupPolicies = append(groupPolicies, NormalizeRateLimitPolicy(decodeObject(policyBytes)))
	}
	if err := groupRows.Err(); err != nil {
		return AsyncWorkerCapacitySnapshot{}, err
	}
	return asyncWorkerCapacityFromPolicySets(modelPolicies, groupPolicies, hardLimit), nil
}

func asyncWorkerCapacityFromPolicies(policies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot {
	return asyncWorkerCapacityFromPolicySets(policies, nil, hardLimit)
}

func asyncWorkerCapacityFromPolicySets(modelPolicies []map[string]any, groupPolicies []map[string]any, hardLimit int) AsyncWorkerCapacitySnapshot {
	snapshot := AsyncWorkerCapacitySnapshot{
		HardLimit:     hardLimit,
		EnabledModels: len(modelPolicies),
		EnabledGroups: len(groupPolicies),
	}
	modelDesired, modelFinite, unlimitedModels := concurrentPolicySetCapacity(modelPolicies)
	groupDesired, groupFinite, unlimitedGroups := concurrentPolicySetCapacity(groupPolicies)
	snapshot.ModelDesired = modelDesired
	snapshot.GroupDesired = groupDesired
	snapshot.UnlimitedModels = unlimitedModels
	snapshot.UnlimitedGroups = unlimitedGroups

	desired := hardLimit
	switch {
	case snapshot.EnabledModels == 0:
		desired = 1
	case modelFinite && groupFinite:
		desired = min(modelDesired, groupDesired)
	case modelFinite:
		desired = modelDesired
	case groupFinite:
		desired = groupDesired
	}
	if desired < 1 {
		desired = 1
	}
	snapshot.Desired = desired
	snapshot.Capacity = desired
	if snapshot.Capacity > hardLimit {
		snapshot.Capacity = hardLimit
		snapshot.Capped = true
	}
	snapshot.GlobalCapacity = snapshot.Capacity
	return snapshot
}

func concurrentPolicySetCapacity(policies []map[string]any) (total int, finite bool, unlimited int) {
	if len(policies) == 0 {
		return 0, false, 0
	}
	finite = true
	for _, policy := range policies {
		capacity, policyFinite := ConcurrentPolicyCapacity(policy)
		if !policyFinite {
			unlimited++
			finite = false
			continue
		}
		total += capacity
	}
	return total, finite, unlimited
}