perf(queue): 按策略动态扩缩异步 Worker

将 River 执行容量按平台模型与用户组的有效并发策略动态调整，并为长任务续租、并发租约原子抢占和限流退避补充保护。\n\n统一平台模型限流继承语义，兼容历史 platformLimits/modelLimits，并为三个迁移图像模型建立独立并发租约。补充管理端显式继承/覆盖配置、指标、单元测试及隔离 PostgreSQL 验收。\n\n验证：go test ./... -count=1；pnpm lint；pnpm test；pnpm build；隔离 PostgreSQL 并发原子性/续租测试；128 任务与三分钟长任务动态 Worker 验收。
2026-07-24 12:26:56 +08:00
parent 290b8c1854
commit 6c5daf29ca
37 changed files with 2365 additions and 205 deletions
@@ -0,0 +1,189 @@
+package runner
+
+import (
+	"context"
+	"errors"
+	"io"
+	"log/slog"
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
+	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
+)
+
+type fakeAsyncExecutionClient struct {
+	startErr      error
+	stopGate      <-chan struct{}
+	started       atomic.Int64
+	stopped       atomic.Int64
+	stopCancelled atomic.Int64
+}
+
+func (c *fakeAsyncExecutionClient) Start(context.Context) error {
+	c.started.Add(1)
+	return c.startErr
+}
+
+func (c *fakeAsyncExecutionClient) Stop(ctx context.Context) error {
+	c.stopped.Add(1)
+	if c.stopGate == nil {
+		return nil
+	}
+	select {
+	case <-c.stopGate:
+		return nil
+	case <-ctx.Done():
+		return ctx.Err()
+	}
+}
+
+func (c *fakeAsyncExecutionClient) StopAndCancel(context.Context) error {
+	c.stopCancelled.Add(1)
+	return nil
+}
+
+func TestResizeAsyncWorkerCapacityStartsReplacementBeforeGracefulDrain(t *testing.T) {
+	oldClient := &fakeAsyncExecutionClient{}
+	drainGate := make(chan struct{})
+	oldClient.stopGate = drainGate
+	newClient := &fakeAsyncExecutionClient{}
+	service := asyncWorkerManagerTestService(1)
+	service.riverExecutionClient = oldClient
+	service.asyncCapacityLoader = fixedAsyncCapacity(3)
+	service.asyncClientFactory = func(capacity int) (asyncExecutionClient, error) {
+		if capacity != 3 {
+			t.Fatalf("factory capacity=%d, want=3", capacity)
+		}
+		return newClient, nil
+	}
+
+	service.resizeAsyncWorkerCapacity(context.Background())
+	if newClient.started.Load() != 1 {
+		t.Fatalf("replacement start count=%d, want=1", newClient.started.Load())
+	}
+	if service.riverExecutionClient != newClient || service.riverWorkerCapacity != 3 {
+		t.Fatalf("replacement was not installed: capacity=%d client=%T", service.riverWorkerCapacity, service.riverExecutionClient)
+	}
+	waitForAtomicValue(t, &oldClient.stopped, 1)
+	if oldClient.stopCancelled.Load() != 0 {
+		t.Fatal("graceful drain cancelled an already running old task")
+	}
+	close(drainGate)
+	waitForDrainingClients(t, service, 0)
+}
+
+func TestResizeAsyncWorkerCapacityFailuresKeepCurrentClient(t *testing.T) {
+	t.Run("refresh failure", func(t *testing.T) {
+		oldClient := &fakeAsyncExecutionClient{}
+		service := asyncWorkerManagerTestService(4)
+		service.riverExecutionClient = oldClient
+		service.asyncCapacityLoader = func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
+			return store.AsyncWorkerCapacitySnapshot{}, errors.New("database unavailable")
+		}
+		service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
+			t.Fatal("factory must not run after refresh failure")
+			return nil, nil
+		}
+		service.resizeAsyncWorkerCapacity(context.Background())
+		if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
+			t.Fatal("refresh failure replaced the current client")
+		}
+	})
+
+	t.Run("create failure", func(t *testing.T) {
+		oldClient := &fakeAsyncExecutionClient{}
+		service := asyncWorkerManagerTestService(4)
+		service.riverExecutionClient = oldClient
+		service.asyncCapacityLoader = fixedAsyncCapacity(8)
+		service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
+			return nil, errors.New("factory failed")
+		}
+		service.resizeAsyncWorkerCapacity(context.Background())
+		if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
+			t.Fatal("create failure replaced the current client")
+		}
+	})
+
+	t.Run("start failure", func(t *testing.T) {
+		oldClient := &fakeAsyncExecutionClient{}
+		failedClient := &fakeAsyncExecutionClient{startErr: errors.New("start failed")}
+		service := asyncWorkerManagerTestService(4)
+		service.riverExecutionClient = oldClient
+		service.asyncCapacityLoader = fixedAsyncCapacity(8)
+		service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
+			return failedClient, nil
+		}
+		service.resizeAsyncWorkerCapacity(context.Background())
+		if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
+			t.Fatal("start failure replaced the current client")
+		}
+		if failedClient.stopCancelled.Load() != 1 {
+			t.Fatal("failed replacement client was not cleaned up")
+		}
+	})
+}
+
+func TestStopAsyncWorkersIncludesCurrentAndDrainingClients(t *testing.T) {
+	current := &fakeAsyncExecutionClient{}
+	draining := &fakeAsyncExecutionClient{}
+	service := asyncWorkerManagerTestService(2)
+	service.riverExecutionClient = current
+	service.riverDrainingClients[draining] = struct{}{}
+	ctx, cancel := context.WithCancel(context.Background())
+	cancel()
+	service.stopAsyncWorkersOnShutdown(ctx)
+	if current.stopCancelled.Load() != 1 || draining.stopCancelled.Load() != 1 {
+		t.Fatalf("shutdown cancellations current=%d draining=%d, want 1/1", current.stopCancelled.Load(), draining.stopCancelled.Load())
+	}
+}
+
+func asyncWorkerManagerTestService(capacity int) *Service {
+	return &Service{
+		cfg: config.Config{
+			AsyncWorkerHardLimit:              2048,
+			AsyncWorkerRefreshIntervalSeconds: 5,
+		},
+		logger:               slog.New(slog.NewTextHandler(io.Discard, nil)),
+		riverDrainingClients: make(map[asyncExecutionClient]struct{}),
+		riverWorkerCapacity:  capacity,
+		riverExecutionClient: nil,
+	}
+}
+
+func fixedAsyncCapacity(capacity int) func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
+	return func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
+		return store.AsyncWorkerCapacitySnapshot{Capacity: capacity, Desired: capacity, HardLimit: 2048}, nil
+	}
+}
+
+func waitForAtomicValue(t *testing.T, value *atomic.Int64, want int64) {
+	t.Helper()
+	deadline := time.Now().Add(time.Second)
+	for time.Now().Before(deadline) {
+		if value.Load() == want {
+			return
+		}
+		time.Sleep(time.Millisecond)
+	}
+	t.Fatalf("atomic value=%d, want=%d", value.Load(), want)
+}
+
+func waitForDrainingClients(t *testing.T, service *Service, want int) {
+	t.Helper()
+	deadline := time.Now().Add(time.Second)
+	for time.Now().Before(deadline) {
+		service.riverMu.RLock()
+		count := len(service.riverDrainingClients)
+		service.riverMu.RUnlock()
+		if count == want {
+			return
+		}
+		time.Sleep(time.Millisecond)
+	}
+	service.riverMu.RLock()
+	count := len(service.riverDrainingClients)
+	service.riverMu.RUnlock()
+	t.Fatalf("draining client count=%d, want=%d", count, want)
+}