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
}

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
}