mirror of
https://github.com/multica-ai/multica.git
synced 2026-07-26 12:35:35 +02:00
251 lines
8.0 KiB
Go
251 lines
8.0 KiB
Go
package handler
|
|
|
|
import (
|
|
"context"
|
|
"encoding/json"
|
|
"errors"
|
|
"fmt"
|
|
"time"
|
|
|
|
"github.com/redis/go-redis/v9"
|
|
)
|
|
|
|
// Redis-backed implementation of ModelListStore. The wire layout matches
|
|
// runtime_local_skills_redis_store.go (which solves the same multi-node
|
|
// dispatch problem for skill lists/imports) so the operational story is
|
|
// identical: namespaced keys, ZSET-backed pending queue, atomic claim via
|
|
// the shared Lua script.
|
|
//
|
|
// Key layout:
|
|
//
|
|
// mul:{runtime_pending}:model_list:req:<request_id> → JSON-encoded ModelListRequest, TTL = retention
|
|
// mul:{runtime_pending}:model_list:pending:<runtime_id> → ZSET { member = request_id, score = created_at UnixNano }
|
|
// TTL = retention*2 (kept alive long enough for
|
|
// lazy sweep on PopPending)
|
|
//
|
|
// PopPending uses claimPendingScript (defined in
|
|
// runtime_local_skills_redis_store.go) to atomically ZREM the pending entry
|
|
// and SET the record to "running" — splitting those two writes would strand
|
|
// requests on a transient Redis hiccup between them.
|
|
|
|
const (
|
|
// Namespaced under mul:*:model_list:* so the key set doesn't collide with
|
|
// the realtime relay (ws:*) or the local-skill stores.
|
|
modelListKeyPrefix = "mul:" + runtimePendingRedisHashTag + ":model_list:req:"
|
|
modelListPendingPrefix = "mul:" + runtimePendingRedisHashTag + ":model_list:pending:"
|
|
modelListRedisPopMaxRetries = 5
|
|
)
|
|
|
|
func modelListKey(id string) string { return modelListKeyPrefix + id }
|
|
func modelListPendingKey(runtimeID string) string { return modelListPendingPrefix + runtimeID }
|
|
|
|
// RedisModelListStore stores model list requests in Redis so every API node
|
|
// agrees on the same pending / running / terminal state.
|
|
type RedisModelListStore struct {
|
|
rdb *redis.Client
|
|
}
|
|
|
|
func NewRedisModelListStore(rdb *redis.Client) *RedisModelListStore {
|
|
return &RedisModelListStore{rdb: rdb}
|
|
}
|
|
|
|
func (s *RedisModelListStore) Create(ctx context.Context, runtimeID string) (*ModelListRequest, error) {
|
|
now := time.Now()
|
|
req := &ModelListRequest{
|
|
ID: randomID(),
|
|
RuntimeID: runtimeID,
|
|
Status: ModelListPending,
|
|
Supported: true,
|
|
CreatedAt: now,
|
|
UpdatedAt: now,
|
|
}
|
|
data, err := s.marshalRequest(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
pipe := s.rdb.TxPipeline()
|
|
pipe.Set(ctx, modelListKey(req.ID), data, modelListStoreRetention)
|
|
pipe.ZAdd(ctx, modelListPendingKey(runtimeID), redis.Z{
|
|
Score: float64(now.UnixNano()),
|
|
Member: req.ID,
|
|
})
|
|
// Keep the pending zset alive past the per-record retention so stale
|
|
// members can be lazily swept on PopPending.
|
|
pipe.Expire(ctx, modelListPendingKey(runtimeID), modelListStoreRetention*2)
|
|
if _, err := pipe.Exec(ctx); err != nil {
|
|
return nil, fmt.Errorf("persist model list request: %w", err)
|
|
}
|
|
return req, nil
|
|
}
|
|
|
|
func (s *RedisModelListStore) Get(ctx context.Context, id string) (*ModelListRequest, error) {
|
|
return s.loadRequest(ctx, id)
|
|
}
|
|
|
|
// loadRequest fetches a single record, applies timeout transitions if the
|
|
// stored state has aged past the threshold, and persists the transition so
|
|
// sibling nodes observe the same terminal state.
|
|
func (s *RedisModelListStore) loadRequest(ctx context.Context, id string) (*ModelListRequest, error) {
|
|
raw, err := s.rdb.Get(ctx, modelListKey(id)).Bytes()
|
|
if errors.Is(err, redis.Nil) {
|
|
return nil, nil
|
|
}
|
|
if err != nil {
|
|
return nil, fmt.Errorf("get model list request: %w", err)
|
|
}
|
|
req, err := s.unmarshalRequest(raw)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if applyModelListTimeout(req, time.Now()) {
|
|
if err := s.persistRequest(ctx, req); err != nil {
|
|
return nil, err
|
|
}
|
|
// Drop from pending zset on terminal transition. PopPending would
|
|
// also do this, but doing it here keeps the set clean for readers
|
|
// that never call PopPending.
|
|
s.rdb.ZRem(ctx, modelListPendingKey(req.RuntimeID), req.ID)
|
|
}
|
|
return req, nil
|
|
}
|
|
|
|
func (s *RedisModelListStore) persistRequest(ctx context.Context, req *ModelListRequest) error {
|
|
data, err := s.marshalRequest(req)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if err := s.rdb.Set(ctx, modelListKey(req.ID), data, modelListStoreRetention).Err(); err != nil {
|
|
return fmt.Errorf("persist model list request: %w", err)
|
|
}
|
|
return nil
|
|
}
|
|
|
|
// ModelListRequest tags RunStartedAt as `json:"-"` so the server-side
|
|
// bookkeeping field doesn't leak into the HTTP response (the UI only
|
|
// needs Status / UpdatedAt to drive its polling loop). Redis persistence
|
|
// has to keep that field, otherwise the running-timeout escape hatch
|
|
// silently breaks across nodes — every reader sees RunStartedAt=nil and
|
|
// applyModelListTimeout's running branch becomes a no-op. Wrap in an
|
|
// internal envelope that re-promotes the field on the wire.
|
|
type redisModelListEnvelope struct {
|
|
Public *ModelListRequest `json:"r"`
|
|
RunStartedAt *time.Time `json:"s,omitempty"`
|
|
}
|
|
|
|
func (s *RedisModelListStore) marshalRequest(req *ModelListRequest) ([]byte, error) {
|
|
env := redisModelListEnvelope{Public: req, RunStartedAt: req.RunStartedAt}
|
|
data, err := json.Marshal(env)
|
|
if err != nil {
|
|
return nil, fmt.Errorf("marshal model list request: %w", err)
|
|
}
|
|
return data, nil
|
|
}
|
|
|
|
func (s *RedisModelListStore) unmarshalRequest(raw []byte) (*ModelListRequest, error) {
|
|
var env redisModelListEnvelope
|
|
if err := json.Unmarshal(raw, &env); err != nil {
|
|
return nil, fmt.Errorf("decode model list request: %w", err)
|
|
}
|
|
if env.Public == nil {
|
|
return nil, fmt.Errorf("decode model list request: missing payload")
|
|
}
|
|
env.Public.RunStartedAt = env.RunStartedAt
|
|
return env.Public, nil
|
|
}
|
|
|
|
// HasPending is a cheap read-only ZCARD probe used by the heartbeat hot path
|
|
// to decide whether to invoke the side-effecting PopPending.
|
|
func (s *RedisModelListStore) HasPending(ctx context.Context, runtimeID string) (bool, error) {
|
|
cnt, err := s.rdb.ZCard(ctx, modelListPendingKey(runtimeID)).Result()
|
|
if err != nil {
|
|
return false, fmt.Errorf("zcard pending: %w", err)
|
|
}
|
|
return cnt > 0, nil
|
|
}
|
|
|
|
func (s *RedisModelListStore) PopPending(ctx context.Context, runtimeID string) (*ModelListRequest, error) {
|
|
pendingKey := modelListPendingKey(runtimeID)
|
|
|
|
for attempt := 0; attempt < modelListRedisPopMaxRetries; attempt++ {
|
|
ids, err := s.rdb.ZRange(ctx, pendingKey, 0, 0).Result()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("zrange pending: %w", err)
|
|
}
|
|
if len(ids) == 0 {
|
|
return nil, nil
|
|
}
|
|
id := ids[0]
|
|
|
|
req, err := s.loadRequest(ctx, id)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
if req == nil {
|
|
// Record expired but the zset still references it — drop and retry.
|
|
s.rdb.ZRem(ctx, pendingKey, id)
|
|
continue
|
|
}
|
|
if req.Status != ModelListPending {
|
|
// Either the timeout fired inside loadRequest or another node
|
|
// already picked it up. Unlink from the pending set and retry.
|
|
s.rdb.ZRem(ctx, pendingKey, id)
|
|
continue
|
|
}
|
|
|
|
now := time.Now()
|
|
req.Status = ModelListRunning
|
|
req.RunStartedAt = &now
|
|
req.UpdatedAt = now
|
|
data, err := s.marshalRequest(req)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
result, err := claimPendingScript.Run(
|
|
ctx, s.rdb,
|
|
[]string{pendingKey, modelListKey(id)},
|
|
id, data, int(modelListStoreRetention.Seconds()),
|
|
).Int64()
|
|
if err != nil {
|
|
return nil, fmt.Errorf("claim pending: %w", err)
|
|
}
|
|
if result == 0 {
|
|
// Another node won the race. The record is owned by the winner;
|
|
// retry to pick up whatever else is queued (or nothing).
|
|
continue
|
|
}
|
|
return req, nil
|
|
}
|
|
return nil, nil
|
|
}
|
|
|
|
func (s *RedisModelListStore) Complete(ctx context.Context, id string, models []ModelEntry, supported bool) error {
|
|
req, err := s.loadRequest(ctx, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if req == nil {
|
|
return nil
|
|
}
|
|
req.Status = ModelListCompleted
|
|
req.Models = models
|
|
req.Supported = supported
|
|
req.UpdatedAt = time.Now()
|
|
return s.persistRequest(ctx, req)
|
|
}
|
|
|
|
func (s *RedisModelListStore) Fail(ctx context.Context, id string, errMsg string) error {
|
|
req, err := s.loadRequest(ctx, id)
|
|
if err != nil {
|
|
return err
|
|
}
|
|
if req == nil {
|
|
return nil
|
|
}
|
|
req.Status = ModelListFailed
|
|
req.Error = errMsg
|
|
req.UpdatedAt = time.Now()
|
|
return s.persistRequest(ctx, req)
|
|
}
|