wangbo
7786692d32
将 Worker 发现、路由画像、容量与执行传输抽象为平台无关接口，新增 Kubernetes 和静态容量适配器，并以 shadow 模式接入生产配置。 实现网络与容量评分、路由防抖、池队列、同步 Worker 租约、一次性执行令牌，以及提交状态不明时禁止重复分配的安全语义。 新增 0105 兼容迁移、管理接口、指标、OpenAPI 和回归测试。已执行全量 Go 测试、go vet、OpenAPI、迁移安全、Compose 与 Kustomize 验证。
100 lines
2.4 KiB
Go
100 lines
2.4 KiB
Go
package runner
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"time"
|
|
|
|
"github.com/easyai/easyai-ai-gateway/apps/api/internal/executionpool"
|
|
"github.com/jackc/pgx/v5"
|
|
"github.com/riverqueue/river"
|
|
"github.com/riverqueue/river/rivertype"
|
|
)
|
|
|
|
// riverExecutionBroker is the River adapter for the platform-neutral broker
|
|
// port. Queue names and River uniqueness semantics do not escape this adapter.
|
|
type riverExecutionBroker struct {
|
|
service *Service
|
|
}
|
|
|
|
func (broker riverExecutionBroker) Publish(
|
|
ctx context.Context,
|
|
poolID string,
|
|
taskID string,
|
|
notBefore time.Time,
|
|
) (int64, error) {
|
|
tx, err := broker.service.store.Pool().Begin(ctx)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
defer func() {
|
|
rollbackCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
defer cancel()
|
|
_ = tx.Rollback(rollbackCtx)
|
|
}()
|
|
jobID, err := broker.publishTx(ctx, tx, poolID, taskID, notBefore, 2)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if err := tx.Commit(ctx); err != nil {
|
|
return 0, err
|
|
}
|
|
return jobID, nil
|
|
}
|
|
|
|
func (broker riverExecutionBroker) publishTx(
|
|
ctx context.Context,
|
|
tx pgx.Tx,
|
|
poolID string,
|
|
taskID string,
|
|
notBefore time.Time,
|
|
priority int,
|
|
) (int64, error) {
|
|
riverClient := broker.service.asyncControlClient()
|
|
if riverClient == nil {
|
|
return 0, errors.New("River execution broker is not started")
|
|
}
|
|
opts := riverExecutionInsertOptions(poolID, notBefore, priority)
|
|
result, err := riverClient.InsertTx(ctx, tx, asyncTaskArgs{TaskID: taskID}, opts)
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
if result.Job == nil {
|
|
return 0, nil
|
|
}
|
|
if _, err := tx.Exec(ctx, `
|
|
UPDATE gateway_tasks
|
|
SET river_job_id = $2,
|
|
updated_at = now()
|
|
WHERE id = $1::uuid`, taskID, result.Job.ID); err != nil {
|
|
return 0, err
|
|
}
|
|
return result.Job.ID, nil
|
|
}
|
|
|
|
func riverExecutionInsertOptions(poolID string, notBefore time.Time, priority int) *river.InsertOpts {
|
|
if priority < 1 {
|
|
priority = 2
|
|
}
|
|
opts := &river.InsertOpts{
|
|
MaxAttempts: 1000,
|
|
Priority: priority,
|
|
Queue: executionpool.QueueName(poolID),
|
|
Tags: []string{"gateway-task"},
|
|
UniqueOpts: river.UniqueOpts{
|
|
ByArgs: true, ByQueue: true,
|
|
ByState: []rivertype.JobState{
|
|
rivertype.JobStateAvailable, rivertype.JobStatePending,
|
|
rivertype.JobStateRetryable, rivertype.JobStateRunning,
|
|
rivertype.JobStateScheduled,
|
|
},
|
|
},
|
|
}
|
|
if !notBefore.IsZero() {
|
|
opts.ScheduledAt = notBefore
|
|
}
|
|
return opts
|
|
}
|
|
|
|
var _ executionpool.ExecutionBroker = riverExecutionBroker{}
|