wangbo
e05922b0f4
实现本地三节点 K3s 同构环境、脱敏生产快照、Gemini 图片和多参考图视频协议模拟、统一验收报告及故障注入。\n\n新增 Worker 容量控制器、资源与连接预算、任务恢复保护，并将生产验收拆分为 validation 执行和人工 CAS 放量。\n\n验证包括 Go 全量测试、PostgreSQL HTTP 集成测试、go vet、OpenAPI、ShellCheck、前端检查、迁移及发布脚本测试。
533 lines
18 KiB
Go
533 lines
18 KiB
Go
package store
|
|
|
|
import (
|
|
"context"
|
|
"crypto/sha256"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"math"
|
|
"strings"
|
|
"time"
|
|
|
|
"github.com/jackc/pgx/v5"
|
|
)
|
|
|
|
type CapacityProfile struct {
|
|
ProfileKey string `json:"profileKey"`
|
|
ReleaseSHA string `json:"releaseSha"`
|
|
ConfigHash string `json:"configHash"`
|
|
AcceptanceRunID string `json:"acceptanceRunId"`
|
|
Kind string `json:"kind"`
|
|
Model string `json:"model"`
|
|
CapacityProfile string `json:"capacityProfile"`
|
|
StableThroughputPerSecond float64 `json:"stableThroughputPerSecond"`
|
|
AdmittedThroughputPerSecond float64 `json:"admittedThroughputPerSecond"`
|
|
MaxRunning int `json:"maxRunning"`
|
|
MaxQueued int `json:"maxQueued"`
|
|
MaxPredictedWaitSeconds int `json:"maxPredictedWaitSeconds"`
|
|
Enabled bool `json:"enabled"`
|
|
Metadata map[string]any `json:"metadata"`
|
|
CreatedAt time.Time `json:"createdAt"`
|
|
UpdatedAt time.Time `json:"updatedAt"`
|
|
}
|
|
|
|
type CapacityProfileInput struct {
|
|
ProfileKey string `json:"profileKey,omitempty"`
|
|
Kind string `json:"kind"`
|
|
Model string `json:"model,omitempty"`
|
|
StableThroughputPerSecond float64 `json:"stableThroughputPerSecond"`
|
|
AdmittedThroughputPerSecond float64 `json:"admittedThroughputPerSecond"`
|
|
MaxRunning int `json:"maxRunning"`
|
|
MaxQueued int `json:"maxQueued"`
|
|
MaxPredictedWaitSeconds int `json:"maxPredictedWaitSeconds"`
|
|
Metadata map[string]any `json:"metadata,omitempty"`
|
|
}
|
|
|
|
type StageAcceptanceCapacityProfilesInput struct {
|
|
RunID string `json:"-"`
|
|
ConfigHash string `json:"configHash"`
|
|
CapacityProfiles []CapacityProfileInput `json:"capacityProfiles"`
|
|
}
|
|
|
|
func (s *Store) ListCapacityProfiles(ctx context.Context) ([]CapacityProfile, error) {
|
|
rows, err := s.pool.Query(ctx, `
|
|
SELECT profile_key, release_sha, config_hash, COALESCE(acceptance_run_id::text, ''),
|
|
kind, model, capacity_profile,
|
|
stable_throughput_per_second::float8,
|
|
admitted_throughput_per_second::float8,
|
|
max_running, max_queued, max_predicted_wait_seconds, enabled,
|
|
metadata, created_at, updated_at
|
|
FROM gateway_capacity_profiles
|
|
ORDER BY enabled DESC, kind, model, updated_at DESC`)
|
|
if err != nil {
|
|
if IsUndefinedDatabaseObject(err) {
|
|
return []CapacityProfile{}, nil
|
|
}
|
|
return nil, err
|
|
}
|
|
defer rows.Close()
|
|
items := make([]CapacityProfile, 0)
|
|
for rows.Next() {
|
|
item, scanErr := scanCapacityProfile(rows)
|
|
if scanErr != nil {
|
|
return nil, scanErr
|
|
}
|
|
items = append(items, item)
|
|
}
|
|
return items, rows.Err()
|
|
}
|
|
|
|
func replaceCapacityProfilesTx(
|
|
ctx context.Context,
|
|
tx pgx.Tx,
|
|
run AcceptanceRun,
|
|
configHash string,
|
|
inputs []CapacityProfileInput,
|
|
) error {
|
|
configHash = strings.ToLower(strings.TrimSpace(configHash))
|
|
if !capacityConfigHashPattern.MatchString(configHash) {
|
|
return errors.New("configHash must be a lowercase SHA-256 value")
|
|
}
|
|
if err := validateCapacityProfileInputs(inputs); err != nil {
|
|
return err
|
|
}
|
|
if _, err := tx.Exec(ctx, `
|
|
UPDATE gateway_capacity_profiles
|
|
SET enabled = false, updated_at = now()
|
|
WHERE enabled = true`); err != nil {
|
|
return err
|
|
}
|
|
for _, input := range inputs {
|
|
kind := strings.ToLower(strings.TrimSpace(input.Kind))
|
|
model := strings.TrimSpace(input.Model)
|
|
if model == "" {
|
|
model = "*"
|
|
}
|
|
profileKey := strings.TrimSpace(input.ProfileKey)
|
|
if profileKey == "" {
|
|
profileKey = fmt.Sprintf("%s:%s:%s:%s", run.ReleaseSHA[:12], strings.ToLower(run.CapacityProfile), kind, model)
|
|
}
|
|
metadata, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(input.Metadata)))
|
|
if _, err := tx.Exec(ctx, `
|
|
INSERT INTO gateway_capacity_profiles (
|
|
profile_key, release_sha, config_hash, acceptance_run_id, kind, model,
|
|
capacity_profile, stable_throughput_per_second,
|
|
admitted_throughput_per_second, max_running, max_queued,
|
|
max_predicted_wait_seconds, enabled, metadata
|
|
)
|
|
VALUES ($1, $2, $3, $4::uuid, $5, $6, $7, $8, $9, $10, $11, $12, true, $13::jsonb)
|
|
ON CONFLICT (profile_key) DO UPDATE
|
|
SET release_sha = EXCLUDED.release_sha,
|
|
config_hash = EXCLUDED.config_hash,
|
|
acceptance_run_id = EXCLUDED.acceptance_run_id,
|
|
kind = EXCLUDED.kind,
|
|
model = EXCLUDED.model,
|
|
capacity_profile = EXCLUDED.capacity_profile,
|
|
stable_throughput_per_second = EXCLUDED.stable_throughput_per_second,
|
|
admitted_throughput_per_second = EXCLUDED.admitted_throughput_per_second,
|
|
max_running = EXCLUDED.max_running,
|
|
max_queued = EXCLUDED.max_queued,
|
|
max_predicted_wait_seconds = EXCLUDED.max_predicted_wait_seconds,
|
|
enabled = true,
|
|
metadata = EXCLUDED.metadata,
|
|
updated_at = now()`,
|
|
profileKey, run.ReleaseSHA, configHash, run.ID, kind, model,
|
|
run.CapacityProfile, input.StableThroughputPerSecond,
|
|
input.AdmittedThroughputPerSecond, input.MaxRunning, input.MaxQueued,
|
|
input.MaxPredictedWaitSeconds, metadata,
|
|
); err != nil {
|
|
return err
|
|
}
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func validateCapacityProfileInputs(inputs []CapacityProfileInput) error {
|
|
if len(inputs) == 0 {
|
|
return errors.New("at least one certified capacity profile is required")
|
|
}
|
|
hasImage := false
|
|
hasVideo := false
|
|
keys := make(map[string]struct{}, len(inputs))
|
|
for _, input := range inputs {
|
|
kind := strings.ToLower(strings.TrimSpace(input.Kind))
|
|
switch kind {
|
|
case "image":
|
|
hasImage = true
|
|
case "video":
|
|
hasVideo = true
|
|
case "mixed":
|
|
hasImage = true
|
|
hasVideo = true
|
|
default:
|
|
return fmt.Errorf("unsupported capacity profile kind %q", input.Kind)
|
|
}
|
|
if input.StableThroughputPerSecond <= 0 ||
|
|
input.AdmittedThroughputPerSecond <= 0 ||
|
|
input.AdmittedThroughputPerSecond > input.StableThroughputPerSecond ||
|
|
input.MaxRunning <= 0 ||
|
|
input.MaxQueued < 0 ||
|
|
input.MaxPredictedWaitSeconds <= 0 {
|
|
return fmt.Errorf("invalid capacity limits for %s/%s", kind, input.Model)
|
|
}
|
|
if input.AdmittedThroughputPerSecond > input.StableThroughputPerSecond*0.800001 {
|
|
return fmt.Errorf("admitted throughput for %s/%s exceeds the 80%% safety envelope", kind, input.Model)
|
|
}
|
|
model := strings.TrimSpace(input.Model)
|
|
if model == "" {
|
|
model = "*"
|
|
}
|
|
key := kind + "\x00" + model
|
|
if _, exists := keys[key]; exists {
|
|
return fmt.Errorf("duplicate capacity profile for %s/%s", kind, model)
|
|
}
|
|
keys[key] = struct{}{}
|
|
}
|
|
if !hasImage || !hasVideo {
|
|
return errors.New("certified capacity profiles must cover both image and video traffic")
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *Store) StageAcceptanceCapacityProfiles(
|
|
ctx context.Context,
|
|
input StageAcceptanceCapacityProfilesInput,
|
|
) (AcceptanceRun, error) {
|
|
input.RunID = strings.TrimSpace(input.RunID)
|
|
input.ConfigHash = strings.ToLower(strings.TrimSpace(input.ConfigHash))
|
|
if !capacityConfigHashPattern.MatchString(input.ConfigHash) {
|
|
return AcceptanceRun{}, errors.New("configHash must be a lowercase SHA-256 value")
|
|
}
|
|
if err := validateCapacityProfileInputs(input.CapacityProfiles); err != nil {
|
|
return AcceptanceRun{}, err
|
|
}
|
|
rawProfiles, _ := json.Marshal(input.CapacityProfiles)
|
|
var profileValues any
|
|
_ = json.Unmarshal(rawProfiles, &profileValues)
|
|
profiles, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(profileValues)))
|
|
profilesHash := fmt.Sprintf("%x", sha256.Sum256(profiles))
|
|
return scanAcceptanceRun(s.pool.QueryRow(ctx, `
|
|
UPDATE gateway_acceptance_runs
|
|
SET config = jsonb_set(
|
|
jsonb_set(
|
|
jsonb_set(config, '{validationCapacityProfiles}', $2::jsonb, true),
|
|
'{validationConfigHash}', to_jsonb($3::text), true
|
|
),
|
|
'{validationProfilesHash}', to_jsonb($4::text), true
|
|
),
|
|
updated_at = now()
|
|
WHERE id = $1::uuid
|
|
AND status = 'running'
|
|
RETURNING `+acceptanceRunColumns,
|
|
input.RunID, profiles, input.ConfigHash, profilesHash,
|
|
))
|
|
}
|
|
|
|
func capacityProfileInputsHash(inputs []CapacityProfileInput) string {
|
|
raw, _ := json.Marshal(inputs)
|
|
var value any
|
|
_ = json.Unmarshal(raw, &value)
|
|
canonical, _ := json.Marshal(sanitizeAcceptanceMetadata(sanitizeJSONForStorage(value)))
|
|
return fmt.Sprintf("%x", sha256.Sum256(canonical))
|
|
}
|
|
|
|
func enforceTaskCapacityTx(ctx context.Context, tx pgx.Tx, input CreateTaskInput, runMode string) error {
|
|
if runMode != "production" && runMode != "acceptance" {
|
|
return nil
|
|
}
|
|
kind := capacityKindForTask(input.Kind)
|
|
if kind == "" {
|
|
return nil
|
|
}
|
|
model := strings.TrimSpace(input.Model)
|
|
profile, applies, err := capacityProfileForTaskTx(ctx, tx, input, runMode, kind, model)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if !applies {
|
|
return nil
|
|
}
|
|
var admissionLock bool
|
|
if err := tx.QueryRow(ctx, `
|
|
SELECT pg_try_advisory_xact_lock(
|
|
hashtextextended('capacity-profile:' || $1, 0)
|
|
)`, profile.ProfileKey).Scan(&admissionLock); err != nil {
|
|
return err
|
|
}
|
|
if !admissionLock {
|
|
return capacityExceededError(profile, "capacity_admission_busy", "admission_lock", 1, 1, 0, 0)
|
|
}
|
|
|
|
var queued int
|
|
var running int
|
|
var recent int
|
|
if err := tx.QueryRow(ctx, `
|
|
SELECT
|
|
count(*) FILTER (WHERE status = 'queued')::int,
|
|
count(*) FILTER (WHERE status = 'running')::int,
|
|
count(*) FILTER (WHERE created_at >= now() - interval '60 seconds')::int
|
|
FROM gateway_tasks
|
|
WHERE (
|
|
($3 = 'production' AND run_mode = 'production')
|
|
OR (
|
|
$3 = 'acceptance'
|
|
AND acceptance_run_id = NULLIF($4, '')::uuid
|
|
AND run_mode = 'acceptance'
|
|
)
|
|
)
|
|
AND (
|
|
($1 = 'image' AND kind IN ('images.generations', 'images.edits', 'images.vectorize'))
|
|
OR ($1 = 'video' AND kind IN ('videos.generations', 'videos.upscales'))
|
|
OR ($1 = 'mixed' AND kind IN (
|
|
'images.generations', 'images.edits', 'images.vectorize',
|
|
'videos.generations', 'videos.upscales'
|
|
))
|
|
)
|
|
AND ($2 = '*' OR model = $2)
|
|
AND (status IN ('queued', 'running') OR created_at >= now() - interval '60 seconds')`,
|
|
profile.Kind, profile.Model, runMode, strings.TrimSpace(input.AcceptanceRunID),
|
|
).Scan(&queued, &running, &recent); err != nil {
|
|
return err
|
|
}
|
|
|
|
nextQueued := queued + 1
|
|
nextRecentRate := float64(recent+1) / 60
|
|
predictedWait := float64(nextQueued) / profile.AdmittedThroughputPerSecond
|
|
reason := ""
|
|
metric := ""
|
|
current := float64(nextQueued)
|
|
limit := float64(profile.MaxQueued)
|
|
switch {
|
|
case nextQueued > profile.MaxQueued:
|
|
reason, metric = "capacity_queue_limit", "queue_depth"
|
|
case predictedWait > float64(profile.MaxPredictedWaitSeconds):
|
|
reason, metric = "capacity_predicted_wait_limit", "predicted_wait_seconds"
|
|
current, limit = predictedWait, float64(profile.MaxPredictedWaitSeconds)
|
|
case nextRecentRate > profile.AdmittedThroughputPerSecond:
|
|
reason, metric = "capacity_throughput_limit", "throughput_per_second"
|
|
current, limit = nextRecentRate, profile.AdmittedThroughputPerSecond
|
|
}
|
|
if reason == "" {
|
|
return nil
|
|
}
|
|
return capacityExceededError(profile, reason, metric, current, limit, queued, running)
|
|
}
|
|
|
|
func capacityProfileForTaskTx(
|
|
ctx context.Context,
|
|
tx pgx.Tx,
|
|
input CreateTaskInput,
|
|
runMode string,
|
|
kind string,
|
|
model string,
|
|
) (CapacityProfile, bool, error) {
|
|
if runMode == "acceptance" {
|
|
var profilesJSON []byte
|
|
var releaseSHA string
|
|
var capacityProfile string
|
|
var configHash string
|
|
err := tx.QueryRow(ctx, `
|
|
SELECT COALESCE(config->'validationCapacityProfiles', 'null'::jsonb),
|
|
release_sha,
|
|
capacity_profile,
|
|
COALESCE(config->>'validationConfigHash', '')
|
|
FROM gateway_acceptance_runs
|
|
WHERE id = NULLIF($1, '')::uuid
|
|
AND status = 'running'`, strings.TrimSpace(input.AcceptanceRunID)).Scan(
|
|
&profilesJSON, &releaseSHA, &capacityProfile, &configHash,
|
|
)
|
|
if err != nil {
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
return CapacityProfile{}, false, nil
|
|
}
|
|
return CapacityProfile{}, false, err
|
|
}
|
|
if string(profilesJSON) == "null" {
|
|
return CapacityProfile{}, false, nil
|
|
}
|
|
var inputs []CapacityProfileInput
|
|
if err := json.Unmarshal(profilesJSON, &inputs); err != nil {
|
|
return CapacityProfile{}, false, fmt.Errorf("decode staged capacity profiles: %w", err)
|
|
}
|
|
best := -1
|
|
bestRank := 100
|
|
for index, candidate := range inputs {
|
|
candidateKind := strings.ToLower(strings.TrimSpace(candidate.Kind))
|
|
candidateModel := strings.TrimSpace(candidate.Model)
|
|
if candidateModel == "" {
|
|
candidateModel = "*"
|
|
}
|
|
rank := 100
|
|
switch {
|
|
case candidateKind == kind && candidateModel == model:
|
|
rank = 0
|
|
case candidateKind == kind && candidateModel == "*":
|
|
rank = 1
|
|
case candidateKind == "mixed" && candidateModel == "*":
|
|
rank = 2
|
|
}
|
|
if rank < bestRank {
|
|
best = index
|
|
bestRank = rank
|
|
}
|
|
}
|
|
if best < 0 {
|
|
return CapacityProfile{}, true, capacityProfileMissingError(kind, model)
|
|
}
|
|
selected := inputs[best]
|
|
selectedModel := strings.TrimSpace(selected.Model)
|
|
if selectedModel == "" {
|
|
selectedModel = "*"
|
|
}
|
|
return CapacityProfile{
|
|
ProfileKey: "acceptance:" + strings.TrimSpace(input.AcceptanceRunID) + ":" + kind + ":" + selectedModel,
|
|
ReleaseSHA: releaseSHA,
|
|
ConfigHash: configHash,
|
|
AcceptanceRunID: strings.TrimSpace(input.AcceptanceRunID),
|
|
Kind: strings.ToLower(strings.TrimSpace(selected.Kind)),
|
|
Model: selectedModel,
|
|
CapacityProfile: capacityProfile,
|
|
StableThroughputPerSecond: selected.StableThroughputPerSecond,
|
|
AdmittedThroughputPerSecond: selected.AdmittedThroughputPerSecond,
|
|
MaxRunning: selected.MaxRunning,
|
|
MaxQueued: selected.MaxQueued,
|
|
MaxPredictedWaitSeconds: selected.MaxPredictedWaitSeconds,
|
|
Metadata: selected.Metadata,
|
|
}, true, nil
|
|
}
|
|
|
|
profile, err := scanCapacityProfile(tx.QueryRow(ctx, `
|
|
SELECT profile_key, release_sha, config_hash, COALESCE(acceptance_run_id::text, ''),
|
|
kind, model, capacity_profile,
|
|
stable_throughput_per_second::float8,
|
|
admitted_throughput_per_second::float8,
|
|
max_running, max_queued, max_predicted_wait_seconds, enabled,
|
|
metadata, created_at, updated_at
|
|
FROM gateway_capacity_profiles
|
|
WHERE enabled = true
|
|
AND (
|
|
(kind = $1 AND model IN ($2, '*'))
|
|
OR (kind = 'mixed' AND model = '*')
|
|
)
|
|
ORDER BY CASE
|
|
WHEN kind = $1 AND model = $2 THEN 0
|
|
WHEN kind = $1 AND model = '*' THEN 1
|
|
ELSE 2
|
|
END
|
|
LIMIT 1`, kind, model))
|
|
if err != nil {
|
|
if IsUndefinedDatabaseObject(err) {
|
|
return CapacityProfile{}, false, nil
|
|
}
|
|
if errors.Is(err, pgx.ErrNoRows) {
|
|
var enabledProfiles int
|
|
if countErr := tx.QueryRow(ctx, `
|
|
SELECT count(*)::int
|
|
FROM gateway_capacity_profiles
|
|
WHERE enabled = true
|
|
AND kind IN ($1, 'mixed')`, kind).Scan(&enabledProfiles); countErr != nil {
|
|
return CapacityProfile{}, false, countErr
|
|
}
|
|
if enabledProfiles == 0 {
|
|
return CapacityProfile{}, false, nil
|
|
}
|
|
return CapacityProfile{}, true, capacityProfileMissingError(kind, model)
|
|
}
|
|
return CapacityProfile{}, false, err
|
|
}
|
|
return profile, true, nil
|
|
}
|
|
|
|
func capacityProfileMissingError(kind string, model string) error {
|
|
return &RateLimitExceededError{
|
|
ScopeType: "capacity_profile",
|
|
ScopeKey: kind + ":" + model,
|
|
ScopeName: "uncertified_model",
|
|
ScopeMetadata: map[string]any{
|
|
"kind": kind, "model": model,
|
|
},
|
|
Metric: "certified_capacity",
|
|
Limit: 0,
|
|
Amount: 1,
|
|
Current: 0,
|
|
Projected: 1,
|
|
Message: "model has no certified production capacity profile",
|
|
RetryAfter: 60 * time.Second,
|
|
Retryable: true,
|
|
Reason: "capacity_profile_missing",
|
|
}
|
|
}
|
|
|
|
func capacityExceededError(
|
|
profile CapacityProfile,
|
|
reason string,
|
|
metric string,
|
|
current float64,
|
|
limit float64,
|
|
queued int,
|
|
running int,
|
|
) error {
|
|
predictedWait := float64(queued) / profile.AdmittedThroughputPerSecond
|
|
retrySeconds := int(math.Ceil(predictedWait))
|
|
if retrySeconds < 1 {
|
|
retrySeconds = 1
|
|
}
|
|
return &RateLimitExceededError{
|
|
ScopeType: "capacity_profile",
|
|
ScopeKey: profile.ProfileKey,
|
|
ScopeName: profile.CapacityProfile,
|
|
ScopeMetadata: map[string]any{
|
|
"kind": profile.Kind,
|
|
"model": profile.Model,
|
|
"releaseSha": profile.ReleaseSHA,
|
|
"configHash": profile.ConfigHash,
|
|
"maxRunning": profile.MaxRunning,
|
|
"maxQueued": profile.MaxQueued,
|
|
"running": running,
|
|
"predictedWaitSeconds": predictedWait,
|
|
},
|
|
Metric: metric,
|
|
Limit: limit,
|
|
Amount: 1,
|
|
Current: current,
|
|
Projected: current + 1,
|
|
Message: "certified production capacity has been reached",
|
|
RetryAfter: time.Duration(retrySeconds) * time.Second,
|
|
Retryable: true,
|
|
Reason: reason,
|
|
QueueDepth: queued,
|
|
QueueLimit: profile.MaxQueued,
|
|
}
|
|
}
|
|
|
|
func capacityKindForTask(kind string) string {
|
|
switch strings.TrimSpace(kind) {
|
|
case "images.generations", "images.edits", "images.vectorize":
|
|
return "image"
|
|
case "videos.generations", "videos.upscales":
|
|
return "video"
|
|
default:
|
|
return ""
|
|
}
|
|
}
|
|
|
|
func scanCapacityProfile(scanner taskScanner) (CapacityProfile, error) {
|
|
var item CapacityProfile
|
|
var metadata []byte
|
|
if err := scanner.Scan(
|
|
&item.ProfileKey, &item.ReleaseSHA, &item.ConfigHash, &item.AcceptanceRunID,
|
|
&item.Kind, &item.Model, &item.CapacityProfile,
|
|
&item.StableThroughputPerSecond, &item.AdmittedThroughputPerSecond,
|
|
&item.MaxRunning, &item.MaxQueued, &item.MaxPredictedWaitSeconds,
|
|
&item.Enabled, &metadata, &item.CreatedAt, &item.UpdatedAt,
|
|
); err != nil {
|
|
return CapacityProfile{}, err
|
|
}
|
|
_ = json.Unmarshal(metadata, &item.Metadata)
|
|
if item.Metadata == nil {
|
|
item.Metadata = map[string]any{}
|
|
}
|
|
return item, nil
|
|
}
|