diff --git a/deploy/config.yaml b/deploy/config.yaml index 3d76b57..ac316fc 100644 --- a/deploy/config.yaml +++ b/deploy/config.yaml @@ -188,4 +188,10 @@ generation: worker_concurrency: 8 stream_enabled: false article_timeout: 8m + task_lease_ttl: 10m + task_recovery_interval: 1m + task_recovery_timeout: 30s + task_recovery_batch_size: 100 + task_queued_stale_after: 2m + task_max_attempts: 3 diff --git a/deploy/k3s/config/config.yaml b/deploy/k3s/config/config.yaml index f45126c..1d85641 100644 --- a/deploy/k3s/config/config.yaml +++ b/deploy/k3s/config/config.yaml @@ -145,3 +145,9 @@ generation: worker_concurrency: 8 stream_enabled: false article_timeout: 8m + task_lease_ttl: 10m + task_recovery_interval: 1m + task_recovery_timeout: 30s + task_recovery_batch_size: 100 + task_queued_stale_after: 2m + task_max_attempts: 3 diff --git a/server/cmd/worker-generate/main.go b/server/cmd/worker-generate/main.go index 5d94092..0bc37e5 100644 --- a/server/cmd/worker-generate/main.go +++ b/server/cmd/worker-generate/main.go @@ -107,7 +107,7 @@ func main() { kolWorker, app.Logger, app.Config().Generation, - ).WithConfigProvider(app.ConfigStore).Start(ctx) + ).WithCache(app.Cache).WithConfigProvider(app.ConfigStore).Start(ctx) generateworker.NewTemplateAssistWorker( app.RabbitMQ, diff --git a/server/configs/config.yaml b/server/configs/config.yaml index 0c798da..fe11cb4 100644 --- a/server/configs/config.yaml +++ b/server/configs/config.yaml @@ -187,4 +187,10 @@ generation: worker_concurrency: 8 stream_enabled: false article_timeout: 8m + task_lease_ttl: 10m + task_recovery_interval: 1m + task_recovery_timeout: 30s + task_recovery_batch_size: 100 + task_queued_stale_after: 2m + task_max_attempts: 3 diff --git a/server/internal/bootstrap/bootstrap.go b/server/internal/bootstrap/bootstrap.go index 47f702c..02951d1 100644 --- a/server/internal/bootstrap/bootstrap.go +++ b/server/internal/bootstrap/bootstrap.go @@ -126,7 +126,18 @@ func New(configPath string) (*App, error) { cfg.RabbitMQ.DesktopDispatchMonitorPrefix + ".*.*", } desktopDispatch := stream.NewDesktopDispatchHub(mqClient, logger, desktopDispatchBindings) - appCache := cache.New(cfg.Cache.Driver, rdb) + appCache := cache.NewWithOptions(cfg.Cache.Driver, rdb, cache.Options{ + Namespace: cfg.Cache.Namespace, + JitterRatio: cfg.Cache.JitterRatio, + MetricsEnabled: cfg.Cache.MetricsEnabled, + AsyncFillEnabled: cfg.Cache.AsyncFillEnabled, + AsyncFillWorkers: cfg.Cache.AsyncFillWorkers, + AsyncFillBuffer: cfg.Cache.AsyncFillBuffer, + AsyncFillTimeout: cfg.Cache.AsyncFillTimeout, + L1Enabled: cfg.Cache.L1Enabled, + L1TTL: cfg.Cache.L1TTL, + DeleteScanCount: cfg.Cache.DeleteScanCount, + }) brandService := tenantapp.NewBrandService(pool, monitoringPool, auditLogs, cfg.BrandLibrary).WithCache(appCache) monitoringService := tenantapp.NewMonitoringService(pool, monitoringPool, mqClient, cfg.MonitoringDispatch, cfg.BrandLibrary, logger).WithRedis(rdb) kolProfiles := repository.NewKolProfileRepository(pool) diff --git a/server/internal/ops/app/kol_subscription_cache_test.go b/server/internal/ops/app/kol_subscription_cache_test.go index d7c96ae..e9e910f 100644 --- a/server/internal/ops/app/kol_subscription_cache_test.go +++ b/server/internal/ops/app/kol_subscription_cache_test.go @@ -18,15 +18,28 @@ func (c *recordingCache) Get(context.Context, string) ([]byte, error) { return nil, sharedcache.ErrNotFound } +func (c *recordingCache) GetMulti(context.Context, []string) (map[string][]byte, error) { + return map[string][]byte{}, nil +} + func (c *recordingCache) Set(context.Context, string, []byte, time.Duration) error { return nil } +func (c *recordingCache) SetMulti(context.Context, []sharedcache.Item) error { + return nil +} + func (c *recordingCache) Delete(_ context.Context, key string) error { c.deletedKeys = append(c.deletedKeys, key) return nil } +func (c *recordingCache) DeleteMany(_ context.Context, keys []string) error { + c.deletedKeys = append(c.deletedKeys, keys...) + return nil +} + func (c *recordingCache) DeletePrefix(_ context.Context, prefix string) error { c.deletedPrefixes = append(c.deletedPrefixes, prefix) return nil diff --git a/server/internal/ops/config/config.go b/server/internal/ops/config/config.go index 4607e9d..7ba1834 100644 --- a/server/internal/ops/config/config.go +++ b/server/internal/ops/config/config.go @@ -272,6 +272,7 @@ func load(path string) (*Config, error) { if err := resolved.Scan(&settings); err != nil { return nil, fmt.Errorf("scan config: %w", err) } + applySharedDefaults(settings) var cfg Config decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{ @@ -303,6 +304,25 @@ func load(path string) (*Config, error) { return &cfg, nil } +func applySharedDefaults(settings map[string]any) { + cacheSettings := ensureMapSetting(settings, "cache") + if _, ok := cacheSettings["metrics_enabled"]; !ok { + cacheSettings["metrics_enabled"] = true + } +} + +func ensureMapSetting(settings map[string]any, key string) map[string]any { + if settings == nil { + return map[string]any{} + } + if existing, ok := settings[key].(map[string]any); ok { + return existing + } + next := make(map[string]any) + settings[key] = next + return next +} + func normalizeConfig(cfg *Config) { if cfg.Server.Port == 0 { cfg.Server.Port = 8090 @@ -313,9 +333,11 @@ func normalizeConfig(cfg *Config) { if strings.TrimSpace(cfg.Redis.Addr) == "" { cfg.Redis.Addr = "localhost:6379" } + sharedconfig.NormalizeRedisConfig(&cfg.Redis) if strings.TrimSpace(cfg.Cache.Driver) == "" { cfg.Cache.Driver = "redis" } + sharedconfig.NormalizeCacheConfig(&cfg.Cache) if strings.TrimSpace(cfg.Log.Level) == "" { cfg.Log.Level = "info" } diff --git a/server/internal/shared/cache/async.go b/server/internal/shared/cache/async.go new file mode 100644 index 0000000..7991d22 --- /dev/null +++ b/server/internal/shared/cache/async.go @@ -0,0 +1,178 @@ +package cache + +import ( + "context" + "errors" + "runtime/debug" + "sync" + "time" +) + +var ErrAsyncQueueFull = errors.New("cache: async fill queue full") + +type asyncCache struct { + inner Cache + fanout *fanout + timeout time.Duration +} + +func newAsyncCache(inner Cache, workers, buffer int, timeout time.Duration) Cache { + return &asyncCache{ + inner: inner, + fanout: newFanout("cache_fill", workers, buffer), + timeout: timeout, + } +} + +func (c *asyncCache) Get(ctx context.Context, key string) ([]byte, error) { + return c.inner.Get(ctx, key) +} + +func (c *asyncCache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) { + return c.inner.GetMulti(ctx, keys) +} + +func (c *asyncCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { + if isCacheControlKey(key) { + return c.inner.Set(ctx, key, value, ttl) + } + err := c.fanout.Do(ctx, func(taskCtx context.Context) { + cctx, cancel := context.WithTimeout(taskCtx, c.timeout) + defer cancel() + if setErr := c.inner.Set(cctx, key, value, ttl); setErr != nil { + recordCacheOperation("set", "error", time.Duration(0)) + } + }) + recordCacheAsyncQueue(c.fanout.Len()) + if err != nil { + recordCacheAsyncDrop("set") + return err + } + return nil +} + +func (c *asyncCache) SetMulti(ctx context.Context, items []Item) error { + err := c.fanout.Do(ctx, func(taskCtx context.Context) { + cctx, cancel := context.WithTimeout(taskCtx, c.timeout) + defer cancel() + if setErr := c.inner.SetMulti(cctx, items); setErr != nil { + recordCacheOperation("set_multi", "error", time.Duration(0)) + } + }) + recordCacheAsyncQueue(c.fanout.Len()) + if err != nil { + recordCacheAsyncDrop("set_multi") + return err + } + return nil +} + +func (c *asyncCache) Delete(ctx context.Context, key string) error { + return c.inner.Delete(ctx, key) +} + +func (c *asyncCache) DeleteMany(ctx context.Context, keys []string) error { + return c.inner.DeleteMany(ctx, keys) +} + +func (c *asyncCache) DeletePrefix(ctx context.Context, prefix string) error { + return c.inner.DeletePrefix(ctx, prefix) +} + +type fanoutTask struct { + ctx context.Context + fn func(context.Context) +} + +type fanout struct { + name string + ch chan fanoutTask + ctx context.Context + cancel context.CancelFunc + wg sync.WaitGroup +} + +func newFanout(name string, workers, buffer int) *fanout { + if workers <= 0 { + workers = defaultAsyncFillWorkers + } + if buffer <= 0 { + buffer = defaultAsyncFillBuffer + } + ctx, cancel := context.WithCancel(context.Background()) + f := &fanout{ + name: name, + ch: make(chan fanoutTask, buffer), + ctx: ctx, + cancel: cancel, + } + f.wg.Add(workers) + for idx := 0; idx < workers; idx++ { + go f.run() + } + return f +} + +func (f *fanout) Do(ctx context.Context, fn func(context.Context)) error { + if fn == nil { + return nil + } + if err := f.ctx.Err(); err != nil { + return err + } + if ctx == nil { + ctx = context.Background() + } + + select { + case f.ch <- fanoutTask{ctx: contextWithoutCancel(ctx), fn: fn}: + return nil + default: + return ErrAsyncQueueFull + } +} + +func (f *fanout) Len() int { + if f == nil { + return 0 + } + return len(f.ch) +} + +func (f *fanout) Close() { + if f == nil { + return + } + f.cancel() + f.wg.Wait() +} + +func (f *fanout) run() { + defer f.wg.Done() + for { + select { + case task := <-f.ch: + f.safeRun(task) + recordCacheAsyncQueue(len(f.ch)) + case <-f.ctx.Done(): + return + } + } +} + +func (f *fanout) safeRun(task fanoutTask) { + defer func() { + if recover() != nil { + recordCacheAsyncPanic() + _ = debug.Stack() + } + }() + task.fn(task.ctx) +} + +func contextWithoutCancel(ctx context.Context) context.Context { + if ctx == nil { + return context.Background() + } + return context.WithoutCancel(ctx) +} diff --git a/server/internal/shared/cache/cache.go b/server/internal/shared/cache/cache.go index 1d9712f..018aaaf 100644 --- a/server/internal/shared/cache/cache.go +++ b/server/internal/shared/cache/cache.go @@ -8,9 +8,18 @@ import ( var ErrNotFound = errors.New("cache: key not found") +type Item struct { + Key string + Value []byte + TTL time.Duration +} + type Cache interface { Get(ctx context.Context, key string) ([]byte, error) + GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error + SetMulti(ctx context.Context, items []Item) error Delete(ctx context.Context, key string) error + DeleteMany(ctx context.Context, keys []string) error DeletePrefix(ctx context.Context, prefix string) error } diff --git a/server/internal/shared/cache/factory.go b/server/internal/shared/cache/factory.go index 1a24b64..67d4f3a 100644 --- a/server/internal/shared/cache/factory.go +++ b/server/internal/shared/cache/factory.go @@ -2,11 +2,39 @@ package cache import goredis "github.com/redis/go-redis/v9" -func New(driver string, rdb *goredis.Client) Cache { +func New(driver string, rdb *goredis.Client, opts ...Option) Cache { + options := DefaultOptions() + for _, opt := range opts { + if opt != nil { + opt(&options) + } + } + return NewWithOptions(driver, rdb, options) +} + +func NewWithOptions(driver string, rdb *goredis.Client, options Options) Cache { + options = normalizeOptions(options) + + var backend Cache switch driver { case "redis": - return newRedisCache(rdb) + backend = newRedisCache(rdb, options.DeleteScanCount) default: - return newMemoryCache() + backend = newMemoryCache() } + + if options.Namespace != "" { + backend = newPrefixCache(backend, options.Namespace) + } + if options.L1Enabled && driver == "redis" { + backend = newL1Cache(backend, options.L1TTL) + } + if options.MetricsEnabled { + backend = newObservedCache(backend, driver) + } + if options.AsyncFillEnabled { + backend = newAsyncCache(backend, options.AsyncFillWorkers, options.AsyncFillBuffer, options.AsyncFillTimeout) + } + + return newOptionsCache(backend, options) } diff --git a/server/internal/shared/cache/l1.go b/server/internal/shared/cache/l1.go new file mode 100644 index 0000000..af4bda5 --- /dev/null +++ b/server/internal/shared/cache/l1.go @@ -0,0 +1,117 @@ +package cache + +import ( + "context" + "strings" + "time" +) + +type l1Cache struct { + inner Cache + local Cache + ttl time.Duration +} + +func newL1Cache(inner Cache, ttl time.Duration) Cache { + if ttl <= 0 { + return inner + } + return &l1Cache{inner: inner, local: newMemoryCache(), ttl: ttl} +} + +func (c *l1Cache) Get(ctx context.Context, key string) ([]byte, error) { + if isCacheControlKey(key) { + return c.inner.Get(ctx, key) + } + if value, err := c.local.Get(ctx, key); err == nil { + recordCacheL1("hit") + return value, nil + } + recordCacheL1("miss") + + value, err := c.inner.Get(ctx, key) + if err != nil { + return nil, err + } + _ = c.local.Set(ctx, key, value, c.ttl) + return value, nil +} + +func (c *l1Cache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) { + if len(keys) == 0 { + return map[string][]byte{}, nil + } + + result := make(map[string][]byte, len(keys)) + miss := make([]string, 0, len(keys)) + for _, key := range keys { + value, err := c.local.Get(ctx, key) + if err == nil { + recordCacheL1("hit") + result[key] = value + continue + } + recordCacheL1("miss") + miss = append(miss, key) + } + if len(miss) == 0 { + return result, nil + } + + values, err := c.inner.GetMulti(ctx, miss) + if err != nil { + return nil, err + } + for key, value := range values { + result[key] = value + _ = c.local.Set(ctx, key, value, c.ttl) + } + return result, nil +} + +func (c *l1Cache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { + if isCacheControlKey(key) { + return c.inner.Set(ctx, key, value, ttl) + } + if err := c.inner.Set(ctx, key, value, ttl); err != nil { + return err + } + return c.local.Set(ctx, key, value, c.l1TTL(ttl)) +} + +func (c *l1Cache) SetMulti(ctx context.Context, items []Item) error { + if err := c.inner.SetMulti(ctx, items); err != nil { + return err + } + for _, item := range items { + _ = c.local.Set(ctx, item.Key, item.Value, c.l1TTL(item.TTL)) + } + return nil +} + +func (c *l1Cache) Delete(ctx context.Context, key string) error { + _ = c.local.Delete(ctx, key) + return c.inner.Delete(ctx, key) +} + +func (c *l1Cache) DeleteMany(ctx context.Context, keys []string) error { + _ = c.local.DeleteMany(ctx, keys) + return c.inner.DeleteMany(ctx, keys) +} + +func (c *l1Cache) DeletePrefix(ctx context.Context, prefix string) error { + _ = c.local.DeletePrefix(ctx, prefix) + return c.inner.DeletePrefix(ctx, prefix) +} + +func (c *l1Cache) l1TTL(ttl time.Duration) time.Duration { + if ttl > 0 && ttl < c.ttl { + return ttl + } + return c.ttl +} + +func isCacheControlKey(key string) bool { + key = strings.TrimSpace(key) + return strings.HasPrefix(key, "cache:version:") || strings.Contains(key, ":cache:version:") +} diff --git a/server/internal/shared/cache/memory.go b/server/internal/shared/cache/memory.go index e3ea05f..2108c32 100644 --- a/server/internal/shared/cache/memory.go +++ b/server/internal/shared/cache/memory.go @@ -41,6 +41,25 @@ func (c *memoryCache) Get(_ context.Context, key string) ([]byte, error) { return dst, nil } +func (c *memoryCache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) { + if len(keys) == 0 { + return map[string][]byte{}, nil + } + + result := make(map[string][]byte, len(keys)) + for _, key := range keys { + value, err := c.Get(ctx, key) + if err != nil { + if err == ErrNotFound { + continue + } + return nil, err + } + result[key] = value + } + return result, nil +} + func (c *memoryCache) Set(_ context.Context, key string, value []byte, ttl time.Duration) error { dst := make([]byte, len(value)) copy(dst, value) @@ -51,6 +70,15 @@ func (c *memoryCache) Set(_ context.Context, key string, value []byte, ttl time. return nil } +func (c *memoryCache) SetMulti(ctx context.Context, items []Item) error { + for _, item := range items { + if err := c.Set(ctx, item.Key, item.Value, item.TTL); err != nil { + return err + } + } + return nil +} + func (c *memoryCache) Delete(_ context.Context, key string) error { c.mu.Lock() delete(c.entries, key) @@ -58,6 +86,15 @@ func (c *memoryCache) Delete(_ context.Context, key string) error { return nil } +func (c *memoryCache) DeleteMany(ctx context.Context, keys []string) error { + for _, key := range keys { + if err := c.Delete(ctx, key); err != nil { + return err + } + } + return nil +} + func (c *memoryCache) DeletePrefix(_ context.Context, prefix string) error { if prefix == "" { return nil diff --git a/server/internal/shared/cache/memory_test.go b/server/internal/shared/cache/memory_test.go new file mode 100644 index 0000000..b1cc2cc --- /dev/null +++ b/server/internal/shared/cache/memory_test.go @@ -0,0 +1,44 @@ +package cache + +import ( + "context" + "testing" + "time" +) + +func TestMemoryCacheMultiAndDeleteMany(t *testing.T) { + ctx := context.Background() + c := newMemoryCache() + + if err := c.SetMulti(ctx, []Item{ + {Key: "a", Value: []byte("one"), TTL: time.Minute}, + {Key: "b", Value: []byte("two"), TTL: time.Minute}, + }); err != nil { + t.Fatalf("SetMulti() error = %v", err) + } + + values, err := c.GetMulti(ctx, []string{"a", "missing", "b"}) + if err != nil { + t.Fatalf("GetMulti() error = %v", err) + } + if got := string(values["a"]); got != "one" { + t.Fatalf("values[a] = %q, want one", got) + } + if got := string(values["b"]); got != "two" { + t.Fatalf("values[b] = %q, want two", got) + } + if _, ok := values["missing"]; ok { + t.Fatalf("missing key should not be returned") + } + + if err := c.DeleteMany(ctx, []string{"a", "b"}); err != nil { + t.Fatalf("DeleteMany() error = %v", err) + } + values, err = c.GetMulti(ctx, []string{"a", "b"}) + if err != nil { + t.Fatalf("GetMulti() after DeleteMany error = %v", err) + } + if len(values) != 0 { + t.Fatalf("values after DeleteMany = %v, want empty", values) + } +} diff --git a/server/internal/shared/cache/metrics.go b/server/internal/shared/cache/metrics.go new file mode 100644 index 0000000..c39d032 --- /dev/null +++ b/server/internal/shared/cache/metrics.go @@ -0,0 +1,263 @@ +package cache + +import ( + "context" + "errors" + "sort" + "strings" + "sync" + "sync/atomic" + "time" + + "github.com/prometheus/client_golang/prometheus" +) + +type metricsKey struct { + Operation string + Result string +} + +type cacheMetricsRecorder struct { + operations sync.Map + l1HitTotal atomic.Int64 + l1MissTotal atomic.Int64 + asyncDropTotal atomic.Int64 + asyncPanicTotal atomic.Int64 + asyncQueueLength atomic.Int64 +} + +type cacheOperationMetrics struct { + Count atomic.Int64 + DurationNanos atomic.Int64 +} + +var cacheMetrics = &cacheMetricsRecorder{} + +type observedCache struct { + inner Cache + driver string +} + +func newObservedCache(inner Cache, driver string) Cache { + driver = strings.TrimSpace(driver) + if driver == "" { + driver = "unknown" + } + return &observedCache{inner: inner, driver: driver} +} + +func (c *observedCache) Get(ctx context.Context, key string) ([]byte, error) { + start := time.Now() + value, err := c.inner.Get(ctx, key) + recordCacheOperation("get", cacheResult(err), time.Since(start)) + return value, err +} + +func (c *observedCache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) { + start := time.Now() + values, err := c.inner.GetMulti(ctx, keys) + result := cacheResult(err) + if err == nil { + recordCacheOperation("get_multi_hit", "ok", time.Duration(0)) + missCount := len(keys) - len(values) + if missCount > 0 { + for i := 0; i < missCount; i++ { + recordCacheOperation("get_multi_miss", "not_found", time.Duration(0)) + } + } + } + recordCacheOperation("get_multi", result, time.Since(start)) + return values, err +} + +func (c *observedCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { + start := time.Now() + err := c.inner.Set(ctx, key, value, ttl) + recordCacheOperation("set", cacheResult(err), time.Since(start)) + return err +} + +func (c *observedCache) SetMulti(ctx context.Context, items []Item) error { + start := time.Now() + err := c.inner.SetMulti(ctx, items) + recordCacheOperation("set_multi", cacheResult(err), time.Since(start)) + return err +} + +func (c *observedCache) Delete(ctx context.Context, key string) error { + start := time.Now() + err := c.inner.Delete(ctx, key) + recordCacheOperation("delete", cacheResult(err), time.Since(start)) + return err +} + +func (c *observedCache) DeleteMany(ctx context.Context, keys []string) error { + start := time.Now() + err := c.inner.DeleteMany(ctx, keys) + recordCacheOperation("delete_many", cacheResult(err), time.Since(start)) + return err +} + +func (c *observedCache) DeletePrefix(ctx context.Context, prefix string) error { + start := time.Now() + err := c.inner.DeletePrefix(ctx, prefix) + recordCacheOperation("delete_prefix", cacheResult(err), time.Since(start)) + return err +} + +func cacheResult(err error) string { + switch { + case err == nil: + return "ok" + case errors.Is(err, ErrNotFound): + return "not_found" + case errors.Is(err, ErrAsyncQueueFull): + return "async_queue_full" + default: + return "error" + } +} + +func recordCacheOperation(operation, result string, duration time.Duration) { + if operation == "" { + operation = "unknown" + } + if result == "" { + result = "unknown" + } + valueAny, _ := cacheMetrics.operations.LoadOrStore(metricsKey{ + Operation: operation, + Result: result, + }, &cacheOperationMetrics{}) + value := valueAny.(*cacheOperationMetrics) + value.Count.Add(1) + if duration > 0 { + value.DurationNanos.Add(duration.Nanoseconds()) + } +} + +func recordCacheL1(result string) { + if result == "hit" { + cacheMetrics.l1HitTotal.Add(1) + return + } + cacheMetrics.l1MissTotal.Add(1) +} + +func recordCacheAsyncDrop(operation string) { + cacheMetrics.asyncDropTotal.Add(1) + recordCacheOperation("async_drop_"+operation, "dropped", time.Duration(0)) +} + +func recordCacheAsyncPanic() { + cacheMetrics.asyncPanicTotal.Add(1) +} + +func recordCacheAsyncQueue(length int) { + cacheMetrics.asyncQueueLength.Store(int64(length)) +} + +type MetricsSnapshot struct { + Operations []OperationMetric `json:"operations"` + L1HitTotal int64 `json:"l1_hit_total"` + L1MissTotal int64 `json:"l1_miss_total"` + AsyncDropTotal int64 `json:"async_drop_total"` + AsyncPanicTotal int64 `json:"async_panic_total"` + AsyncQueueLength int64 `json:"async_queue_length"` +} + +type OperationMetric struct { + Operation string `json:"operation"` + Result string `json:"result"` + Count int64 `json:"count"` + DurationSeconds float64 `json:"duration_seconds"` +} + +func MetricsSnapshotValue() MetricsSnapshot { + snapshot := MetricsSnapshot{ + L1HitTotal: cacheMetrics.l1HitTotal.Load(), + L1MissTotal: cacheMetrics.l1MissTotal.Load(), + AsyncDropTotal: cacheMetrics.asyncDropTotal.Load(), + AsyncPanicTotal: cacheMetrics.asyncPanicTotal.Load(), + AsyncQueueLength: cacheMetrics.asyncQueueLength.Load(), + } + cacheMetrics.operations.Range(func(keyAny, valueAny any) bool { + key, ok := keyAny.(metricsKey) + if !ok { + return true + } + value, ok := valueAny.(*cacheOperationMetrics) + if !ok { + return true + } + snapshot.Operations = append(snapshot.Operations, OperationMetric{ + Operation: key.Operation, + Result: key.Result, + Count: value.Count.Load(), + DurationSeconds: float64(value.DurationNanos.Load()) / float64(time.Second), + }) + return true + }) + sort.Slice(snapshot.Operations, func(i, j int) bool { + if snapshot.Operations[i].Operation == snapshot.Operations[j].Operation { + return snapshot.Operations[i].Result < snapshot.Operations[j].Result + } + return snapshot.Operations[i].Operation < snapshot.Operations[j].Operation + }) + return snapshot +} + +func ResetMetricsForTest() { + cacheMetrics = &cacheMetricsRecorder{} +} + +type MetricsCollector struct{} + +func (MetricsCollector) Describe(chan<- *prometheus.Desc) {} + +func (MetricsCollector) Collect(ch chan<- prometheus.Metric) { + snapshot := MetricsSnapshotValue() + for _, item := range snapshot.Operations { + ch <- prometheus.MustNewConstMetric( + prometheus.NewDesc("cache_operations_total", "Total cache operations by operation and result.", []string{"operation", "result"}, nil), + prometheus.CounterValue, + float64(item.Count), + item.Operation, + item.Result, + ) + ch <- prometheus.MustNewConstMetric( + prometheus.NewDesc("cache_operation_duration_seconds_total", "Total cache operation duration by operation and result.", []string{"operation", "result"}, nil), + prometheus.CounterValue, + item.DurationSeconds, + item.Operation, + item.Result, + ) + } + ch <- prometheus.MustNewConstMetric( + prometheus.NewDesc("cache_l1_access_total", "Total L1 cache access by result.", []string{"result"}, nil), + prometheus.CounterValue, + float64(snapshot.L1HitTotal), + "hit", + ) + ch <- prometheus.MustNewConstMetric( + prometheus.NewDesc("cache_l1_access_total", "Total L1 cache access by result.", []string{"result"}, nil), + prometheus.CounterValue, + float64(snapshot.L1MissTotal), + "miss", + ) + ch <- prometheus.MustNewConstMetric( + prometheus.NewDesc("cache_async_drops_total", "Total async cache fill tasks dropped because the queue was full.", nil, nil), + prometheus.CounterValue, + float64(snapshot.AsyncDropTotal), + ) + ch <- prometheus.MustNewConstMetric( + prometheus.NewDesc("cache_async_panics_total", "Total panics recovered in async cache fill workers.", nil, nil), + prometheus.CounterValue, + float64(snapshot.AsyncPanicTotal), + ) + ch <- prometheus.MustNewConstMetric( + prometheus.NewDesc("cache_async_queue_length", "Current async cache fill queue length.", nil, nil), + prometheus.GaugeValue, + float64(snapshot.AsyncQueueLength), + ) +} diff --git a/server/internal/shared/cache/options.go b/server/internal/shared/cache/options.go new file mode 100644 index 0000000..ca2e1d4 --- /dev/null +++ b/server/internal/shared/cache/options.go @@ -0,0 +1,159 @@ +package cache + +import ( + "context" + "strings" + "time" +) + +const ( + defaultNamespace = "geo" + defaultCacheJitterRatio = 0.1 + defaultDeleteScanCount = int64(500) + defaultAsyncFillWorkers = 2 + defaultAsyncFillBuffer = 1024 + defaultAsyncFillTimeout = 2 * time.Second + defaultL1TTL = 30 * time.Second +) + +type Options struct { + Namespace string + JitterRatio float64 + MetricsEnabled bool + AsyncFillEnabled bool + AsyncFillWorkers int + AsyncFillBuffer int + AsyncFillTimeout time.Duration + L1Enabled bool + L1TTL time.Duration + DeleteScanCount int64 +} + +func DefaultOptions() Options { + return Options{ + Namespace: defaultNamespace, + JitterRatio: defaultCacheJitterRatio, + MetricsEnabled: true, + AsyncFillWorkers: defaultAsyncFillWorkers, + AsyncFillBuffer: defaultAsyncFillBuffer, + AsyncFillTimeout: defaultAsyncFillTimeout, + L1TTL: defaultL1TTL, + DeleteScanCount: defaultDeleteScanCount, + } +} + +type Option func(*Options) + +func WithNamespace(namespace string) Option { + return func(o *Options) { + o.Namespace = namespace + } +} + +func WithJitterRatio(ratio float64) Option { + return func(o *Options) { + o.JitterRatio = ratio + } +} + +func WithMetrics(enabled bool) Option { + return func(o *Options) { + o.MetricsEnabled = enabled + } +} + +func WithAsyncFill(enabled bool, workers, buffer int, timeout time.Duration) Option { + return func(o *Options) { + o.AsyncFillEnabled = enabled + o.AsyncFillWorkers = workers + o.AsyncFillBuffer = buffer + o.AsyncFillTimeout = timeout + } +} + +func WithL1(enabled bool, ttl time.Duration) Option { + return func(o *Options) { + o.L1Enabled = enabled + o.L1TTL = ttl + } +} + +func WithDeleteScanCount(count int64) Option { + return func(o *Options) { + o.DeleteScanCount = count + } +} + +func normalizeOptions(options Options) Options { + options.Namespace = strings.TrimSpace(options.Namespace) + if options.Namespace == "" { + options.Namespace = defaultNamespace + } + if options.JitterRatio < 0 { + options.JitterRatio = 0 + } + if options.JitterRatio > 1 { + options.JitterRatio = 1 + } + if options.AsyncFillWorkers <= 0 { + options.AsyncFillWorkers = defaultAsyncFillWorkers + } + if options.AsyncFillBuffer <= 0 { + options.AsyncFillBuffer = defaultAsyncFillBuffer + } + if options.AsyncFillTimeout <= 0 { + options.AsyncFillTimeout = defaultAsyncFillTimeout + } + if options.L1TTL <= 0 { + options.L1TTL = defaultL1TTL + } + if options.DeleteScanCount <= 0 { + options.DeleteScanCount = defaultDeleteScanCount + } + return options +} + +type optionsProvider interface { + CacheOptions() Options +} + +type optionsCache struct { + inner Cache + options Options +} + +func newOptionsCache(inner Cache, options Options) Cache { + return &optionsCache{inner: inner, options: normalizeOptions(options)} +} + +func (c *optionsCache) CacheOptions() Options { + return c.options +} + +func (c *optionsCache) Get(ctx context.Context, key string) ([]byte, error) { + return c.inner.Get(ctx, key) +} + +func (c *optionsCache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) { + return c.inner.GetMulti(ctx, keys) +} + +func (c *optionsCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { + return c.inner.Set(ctx, key, value, ttl) +} + +func (c *optionsCache) SetMulti(ctx context.Context, items []Item) error { + return c.inner.SetMulti(ctx, items) +} + +func (c *optionsCache) Delete(ctx context.Context, key string) error { + return c.inner.Delete(ctx, key) +} + +func (c *optionsCache) DeleteMany(ctx context.Context, keys []string) error { + return c.inner.DeleteMany(ctx, keys) +} + +func (c *optionsCache) DeletePrefix(ctx context.Context, prefix string) error { + return c.inner.DeletePrefix(ctx, prefix) +} diff --git a/server/internal/shared/cache/prefix.go b/server/internal/shared/cache/prefix.go new file mode 100644 index 0000000..6d76f7f --- /dev/null +++ b/server/internal/shared/cache/prefix.go @@ -0,0 +1,91 @@ +package cache + +import ( + "context" + "strings" + "time" +) + +type prefixCache struct { + inner Cache + prefix string +} + +func newPrefixCache(inner Cache, namespace string) Cache { + namespace = strings.Trim(strings.TrimSpace(namespace), ":") + if namespace == "" { + return inner + } + return &prefixCache{inner: inner, prefix: namespace + ":"} +} + +func (c *prefixCache) Get(ctx context.Context, key string) ([]byte, error) { + return c.inner.Get(ctx, c.key(key)) +} + +func (c *prefixCache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) { + if len(keys) == 0 { + return map[string][]byte{}, nil + } + + prefixed := make([]string, len(keys)) + originalByPrefixed := make(map[string]string, len(keys)) + for idx, key := range keys { + prefixedKey := c.key(key) + prefixed[idx] = prefixedKey + originalByPrefixed[prefixedKey] = key + } + + values, err := c.inner.GetMulti(ctx, prefixed) + if err != nil { + return nil, err + } + result := make(map[string][]byte, len(values)) + for prefixedKey, value := range values { + key, ok := originalByPrefixed[prefixedKey] + if !ok { + continue + } + result[key] = value + } + return result, nil +} + +func (c *prefixCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { + return c.inner.Set(ctx, c.key(key), value, ttl) +} + +func (c *prefixCache) SetMulti(ctx context.Context, items []Item) error { + if len(items) == 0 { + return nil + } + prefixed := make([]Item, len(items)) + for idx, item := range items { + item.Key = c.key(item.Key) + prefixed[idx] = item + } + return c.inner.SetMulti(ctx, prefixed) +} + +func (c *prefixCache) Delete(ctx context.Context, key string) error { + return c.inner.Delete(ctx, c.key(key)) +} + +func (c *prefixCache) DeleteMany(ctx context.Context, keys []string) error { + if len(keys) == 0 { + return nil + } + prefixed := make([]string, len(keys)) + for idx, key := range keys { + prefixed[idx] = c.key(key) + } + return c.inner.DeleteMany(ctx, prefixed) +} + +func (c *prefixCache) DeletePrefix(ctx context.Context, prefix string) error { + return c.inner.DeletePrefix(ctx, c.key(prefix)) +} + +func (c *prefixCache) key(key string) string { + return c.prefix + strings.TrimPrefix(key, c.prefix) +} diff --git a/server/internal/shared/cache/readthrough.go b/server/internal/shared/cache/readthrough.go index 43453bf..cc8926a 100644 --- a/server/internal/shared/cache/readthrough.go +++ b/server/internal/shared/cache/readthrough.go @@ -3,6 +3,7 @@ package cache import ( "context" "encoding/json" + "errors" "math/rand" "sync" "time" @@ -11,7 +12,7 @@ import ( ) const ( - defaultJitterRatio = 0.1 + readthroughDefaultJitterRatio = 0.1 ) var ( @@ -62,6 +63,7 @@ func loadJSONInternal[T any]( ) (T, bool, error) { var zero T if c == nil { + recordCacheOperation("readthrough_load", "cache_disabled", time.Duration(0)) return loader(ctx) } @@ -69,14 +71,23 @@ func loadJSONInternal[T any]( var envelope payloadEnvelope[T] if unmarshalErr := json.Unmarshal(raw, &envelope); unmarshalErr == nil { if envelope.Empty { + recordCacheOperation("readthrough_empty", "hit", time.Duration(0)) return zero, false, nil } + recordCacheOperation("readthrough_hit", "ok", time.Duration(0)) return envelope.Value, true, nil } + recordCacheOperation("readthrough_unmarshal", "error", time.Duration(0)) + } else if errors.Is(err, ErrNotFound) { + recordCacheOperation("readthrough_miss", "not_found", time.Duration(0)) + } else { + recordCacheOperation("readthrough_get", "error", time.Duration(0)) } result, err, _ := group.Do(key, func() (interface{}, error) { + start := time.Now() value, found, loadErr := loader(ctx) + recordCacheOperation("readthrough_load", cacheResult(loadErr), time.Since(start)) if loadErr != nil { return nil, loadErr } @@ -86,13 +97,18 @@ func loadJSONInternal[T any]( Value: value, } if raw, marshalErr := json.Marshal(envelope); marshalErr == nil { + recordCacheOperation("readthrough_marshal", "ok", time.Duration(0)) cacheTTL := ttl if !found { cacheTTL = emptyTTL } if cacheTTL > 0 { - _ = c.Set(ctx, key, raw, ApplyJitter(cacheTTL)) + setStart := time.Now() + setErr := c.Set(ctx, key, raw, ApplyJitterWithRatio(cacheTTL, cacheJitterRatio(c))) + recordCacheOperation("readthrough_set", cacheResult(setErr), time.Since(setStart)) } + } else { + recordCacheOperation("readthrough_marshal", "error", time.Duration(0)) } return envelope, nil @@ -112,11 +128,21 @@ func loadJSONInternal[T any]( } func ApplyJitter(ttl time.Duration) time.Duration { + return ApplyJitterWithRatio(ttl, readthroughDefaultJitterRatio) +} + +func ApplyJitterWithRatio(ttl time.Duration, ratio float64) time.Duration { if ttl <= 0 { return ttl } + if ratio <= 0 { + return ttl + } + if ratio > 1 { + ratio = 1 + } - maxJitter := int64(float64(ttl) * defaultJitterRatio) + maxJitter := int64(float64(ttl) * ratio) if maxJitter <= 0 { return ttl } @@ -126,3 +152,13 @@ func ApplyJitter(ttl time.Duration) time.Duration { return ttl + time.Duration(jitterRand.Int63n(maxJitter+1)) } + +func cacheJitterRatio(c Cache) float64 { + if c == nil { + return readthroughDefaultJitterRatio + } + if provider, ok := c.(optionsProvider); ok { + return provider.CacheOptions().JitterRatio + } + return readthroughDefaultJitterRatio +} diff --git a/server/internal/shared/cache/readthrough_test.go b/server/internal/shared/cache/readthrough_test.go new file mode 100644 index 0000000..ea12f52 --- /dev/null +++ b/server/internal/shared/cache/readthrough_test.go @@ -0,0 +1,63 @@ +package cache + +import ( + "context" + "testing" + "time" + + "golang.org/x/sync/singleflight" +) + +func TestLoadJSONWithEmptyCachesEmptyAndRecordsMetrics(t *testing.T) { + ResetMetricsForTest() + ctx := context.Background() + c := NewWithOptions("memory", nil, Options{ + Namespace: "test", + JitterRatio: 0, + MetricsEnabled: true, + }) + var group singleflight.Group + loads := 0 + + value, found, err := LoadJSONWithEmpty(ctx, c, &group, "empty", time.Minute, time.Minute, func(context.Context) (string, bool, error) { + loads++ + return "", false, nil + }) + if err != nil { + t.Fatalf("LoadJSONWithEmpty() error = %v", err) + } + if found || value != "" { + t.Fatalf("first load value=%q found=%v, want empty not found", value, found) + } + + value, found, err = LoadJSONWithEmpty(ctx, c, &group, "empty", time.Minute, time.Minute, func(context.Context) (string, bool, error) { + loads++ + return "unexpected", true, nil + }) + if err != nil { + t.Fatalf("LoadJSONWithEmpty() cached error = %v", err) + } + if found || value != "" { + t.Fatalf("cached value=%q found=%v, want empty not found", value, found) + } + if loads != 1 { + t.Fatalf("loader calls = %d, want 1", loads) + } + + snapshot := MetricsSnapshotValue() + if metricCount(snapshot, "readthrough_empty", "hit") != 1 { + t.Fatalf("readthrough_empty hit metric missing in %#v", snapshot.Operations) + } + if metricCount(snapshot, "readthrough_load", "ok") != 1 { + t.Fatalf("readthrough_load ok metric missing in %#v", snapshot.Operations) + } +} + +func metricCount(snapshot MetricsSnapshot, operation, result string) int64 { + for _, item := range snapshot.Operations { + if item.Operation == operation && item.Result == result { + return item.Count + } + } + return 0 +} diff --git a/server/internal/shared/cache/redis.go b/server/internal/shared/cache/redis.go index cf82bab..a938c58 100644 --- a/server/internal/shared/cache/redis.go +++ b/server/internal/shared/cache/redis.go @@ -3,17 +3,19 @@ package cache import ( "context" "errors" + "fmt" "time" goredis "github.com/redis/go-redis/v9" ) type redisCache struct { - client *goredis.Client + client *goredis.Client + deleteScanCount int64 } -func newRedisCache(client *goredis.Client) Cache { - return &redisCache{client: client} +func newRedisCache(client *goredis.Client, deleteScanCount int64) Cache { + return &redisCache{client: client, deleteScanCount: deleteScanCount} } func (c *redisCache) Get(ctx context.Context, key string) ([]byte, error) { @@ -24,28 +26,81 @@ func (c *redisCache) Get(ctx context.Context, key string) ([]byte, error) { return val, err } +func (c *redisCache) GetMulti(ctx context.Context, keys []string) (map[string][]byte, error) { + if len(keys) == 0 { + return map[string][]byte{}, nil + } + + values, err := c.client.MGet(ctx, keys...).Result() + if err != nil { + return nil, err + } + + result := make(map[string][]byte, len(keys)) + for idx, value := range values { + if value == nil { + continue + } + switch typed := value.(type) { + case string: + result[keys[idx]] = []byte(typed) + case []byte: + dst := make([]byte, len(typed)) + copy(dst, typed) + result[keys[idx]] = dst + default: + result[keys[idx]] = []byte(fmt.Sprint(value)) + } + } + return result, nil +} + func (c *redisCache) Set(ctx context.Context, key string, value []byte, ttl time.Duration) error { return c.client.Set(ctx, key, value, ttl).Err() } +func (c *redisCache) SetMulti(ctx context.Context, items []Item) error { + if len(items) == 0 { + return nil + } + + pipe := c.client.Pipeline() + for _, item := range items { + pipe.Set(ctx, item.Key, item.Value, item.TTL) + } + _, err := pipe.Exec(ctx) + return err +} + func (c *redisCache) Delete(ctx context.Context, key string) error { return c.client.Del(ctx, key).Err() } +func (c *redisCache) DeleteMany(ctx context.Context, keys []string) error { + if len(keys) == 0 { + return nil + } + return c.client.Del(ctx, keys...).Err() +} + func (c *redisCache) DeletePrefix(ctx context.Context, prefix string) error { if prefix == "" { return nil } pattern := prefix + "*" + scanCount := c.deleteScanCount + if scanCount <= 0 { + scanCount = defaultDeleteScanCount + } var cursor uint64 for { - keys, nextCursor, err := c.client.Scan(ctx, cursor, pattern, 100).Result() + keys, nextCursor, err := c.client.Scan(ctx, cursor, pattern, scanCount).Result() if err != nil { return err } if len(keys) > 0 { - if err := c.client.Del(ctx, keys...).Err(); err != nil { + if err := c.DeleteMany(ctx, keys); err != nil { return err } } diff --git a/server/internal/shared/config/config.go b/server/internal/shared/config/config.go index ab694d1..1d9f9db 100644 --- a/server/internal/shared/config/config.go +++ b/server/internal/shared/config/config.go @@ -1,6 +1,7 @@ package config import ( + "crypto/tls" "encoding/json" "fmt" "os" @@ -57,8 +58,24 @@ func (d DatabaseConfig) DSN() string { } type RedisConfig struct { - Addr string `mapstructure:"addr"` - DB int `mapstructure:"db"` + Addr string `mapstructure:"addr"` + Password string `mapstructure:"password"` + DB int `mapstructure:"db"` + DialTimeout time.Duration `mapstructure:"dial_timeout"` + ReadTimeout time.Duration `mapstructure:"read_timeout"` + WriteTimeout time.Duration `mapstructure:"write_timeout"` + PoolSize int `mapstructure:"pool_size"` + MinIdleConns int `mapstructure:"min_idle_conns"` + MaxRetries int `mapstructure:"max_retries"` + PoolTimeout time.Duration `mapstructure:"pool_timeout"` + TLSEnabled bool `mapstructure:"tls_enabled"` +} + +func (r RedisConfig) TLSConfig() *tls.Config { + if !r.TLSEnabled { + return nil + } + return &tls.Config{MinVersion: tls.VersionTLS12} } type RabbitMQConfig struct { @@ -272,14 +289,30 @@ type RetrievalConfig struct { } type CacheConfig struct { - Driver string `mapstructure:"driver"` // "redis" or "memory" + Driver string `mapstructure:"driver"` // "redis" or "memory" + Namespace string `mapstructure:"namespace"` + JitterRatio float64 `mapstructure:"jitter_ratio"` + MetricsEnabled bool `mapstructure:"metrics_enabled"` + AsyncFillEnabled bool `mapstructure:"async_fill_enabled"` + AsyncFillWorkers int `mapstructure:"async_fill_workers"` + AsyncFillBuffer int `mapstructure:"async_fill_buffer"` + AsyncFillTimeout time.Duration `mapstructure:"async_fill_timeout"` + L1Enabled bool `mapstructure:"l1_enabled"` + L1TTL time.Duration `mapstructure:"l1_ttl"` + DeleteScanCount int64 `mapstructure:"delete_scan_count"` } type GenerationConfig struct { - QueueSize int `mapstructure:"queue_size"` - WorkerConcurrency int `mapstructure:"worker_concurrency"` - StreamEnabled bool `mapstructure:"stream_enabled"` - ArticleTimeout time.Duration `mapstructure:"article_timeout"` + QueueSize int `mapstructure:"queue_size"` + WorkerConcurrency int `mapstructure:"worker_concurrency"` + StreamEnabled bool `mapstructure:"stream_enabled"` + ArticleTimeout time.Duration `mapstructure:"article_timeout"` + TaskLeaseTTL time.Duration `mapstructure:"task_lease_ttl"` + TaskRecoveryInterval time.Duration `mapstructure:"task_recovery_interval"` + TaskRecoveryTimeout time.Duration `mapstructure:"task_recovery_timeout"` + TaskRecoveryBatchSize int `mapstructure:"task_recovery_batch_size"` + TaskQueuedStaleAfter time.Duration `mapstructure:"task_queued_stale_after"` + TaskMaxAttempts int `mapstructure:"task_max_attempts"` } func Load(configPath string) (*Config, error) { @@ -305,12 +338,15 @@ func loadWithFiles(configPath string) (*Config, []string, error) { } applyEnvOverrides(cfg) + NormalizeRedisConfig(&cfg.Redis) + NormalizeCacheConfig(&cfg.Cache) normalizeRabbitMQConfig(&cfg.RabbitMQ) normalizeSchedulerConfig(&cfg.Scheduler) normalizeMonitoringConfig(&cfg.MonitoringWorkers) normalizeMonitoringDispatchConfig(&cfg.MonitoringDispatch) normalizeMembershipConfig(&cfg.Membership) normalizeBrandLibraryConfig(&cfg.BrandLibrary) + NormalizeGenerationConfig(&cfg.Generation) files := []string{configFile} if localConfigFile != "" { @@ -360,6 +396,7 @@ func decodeResolvedConfig(configFile, localConfigFile string) (*Config, error) { if err := resolved.Scan(&settings); err != nil { return nil, err } + applyConfigDefaults(settings) var cfg Config decoder, err := mapstructure.NewDecoder(&mapstructure.DecoderConfig{ @@ -380,6 +417,25 @@ func decodeResolvedConfig(configFile, localConfigFile string) (*Config, error) { return &cfg, nil } +func applyConfigDefaults(settings map[string]any) { + cacheSettings := ensureMapSetting(settings, "cache") + if _, ok := cacheSettings["metrics_enabled"]; !ok { + cacheSettings["metrics_enabled"] = true + } +} + +func ensureMapSetting(settings map[string]any, key string) map[string]any { + if settings == nil { + return map[string]any{} + } + if existing, ok := settings[key].(map[string]any); ok { + return existing + } + next := make(map[string]any) + settings[key] = next + return next +} + func normalizeMonitoringDispatchConfig(cfg *MonitoringDispatchConfig) { if cfg == nil { return @@ -392,6 +448,114 @@ func normalizeMonitoringDispatchConfig(cfg *MonitoringDispatchConfig) { } } +func NormalizeRedisConfig(cfg *RedisConfig) { + if cfg == nil { + return + } + cfg.Addr = strings.TrimSpace(cfg.Addr) + if cfg.Addr == "" { + cfg.Addr = "localhost:6379" + } + if cfg.DialTimeout <= 0 { + cfg.DialTimeout = 5 * time.Second + } + if cfg.ReadTimeout <= 0 { + cfg.ReadTimeout = 3 * time.Second + } + if cfg.WriteTimeout <= 0 { + cfg.WriteTimeout = 3 * time.Second + } + if cfg.PoolSize <= 0 { + cfg.PoolSize = 16 + } + if cfg.MinIdleConns < 0 { + cfg.MinIdleConns = 0 + } + if cfg.MinIdleConns > cfg.PoolSize { + cfg.MinIdleConns = cfg.PoolSize + } + if cfg.MaxRetries < 0 { + cfg.MaxRetries = 0 + } + if cfg.PoolTimeout <= 0 { + cfg.PoolTimeout = 4 * time.Second + } +} + +func NormalizeCacheConfig(cfg *CacheConfig) { + if cfg == nil { + return + } + cfg.Driver = strings.ToLower(strings.TrimSpace(cfg.Driver)) + if cfg.Driver == "" { + cfg.Driver = "redis" + } + if cfg.Driver != "redis" && cfg.Driver != "memory" { + cfg.Driver = "memory" + } + cfg.Namespace = strings.Trim(strings.TrimSpace(cfg.Namespace), ":") + if cfg.Namespace == "" { + cfg.Namespace = "geo" + } + if cfg.JitterRatio <= 0 { + cfg.JitterRatio = 0.1 + } + if cfg.JitterRatio > 1 { + cfg.JitterRatio = 1 + } + if cfg.AsyncFillWorkers <= 0 { + cfg.AsyncFillWorkers = 2 + } + if cfg.AsyncFillBuffer <= 0 { + cfg.AsyncFillBuffer = 1024 + } + if cfg.AsyncFillTimeout <= 0 { + cfg.AsyncFillTimeout = 2 * time.Second + } + if cfg.L1TTL <= 0 { + cfg.L1TTL = 30 * time.Second + } + if cfg.DeleteScanCount <= 0 { + cfg.DeleteScanCount = 500 + } +} + +func NormalizeGenerationConfig(cfg *GenerationConfig) { + if cfg == nil { + return + } + if cfg.QueueSize <= 0 { + cfg.QueueSize = 128 + } + if cfg.WorkerConcurrency <= 0 { + cfg.WorkerConcurrency = 1 + } + if cfg.ArticleTimeout <= 0 { + cfg.ArticleTimeout = 8 * time.Minute + } + if cfg.TaskLeaseTTL <= 0 { + cfg.TaskLeaseTTL = cfg.ArticleTimeout + 2*time.Minute + } + if cfg.TaskLeaseTTL < time.Minute { + cfg.TaskLeaseTTL = time.Minute + } + if cfg.TaskRecoveryInterval <= 0 { + cfg.TaskRecoveryInterval = time.Minute + } + if cfg.TaskRecoveryTimeout <= 0 { + cfg.TaskRecoveryTimeout = 30 * time.Second + } + if cfg.TaskRecoveryBatchSize <= 0 { + cfg.TaskRecoveryBatchSize = 100 + } + if cfg.TaskQueuedStaleAfter <= 0 { + cfg.TaskQueuedStaleAfter = 2 * time.Minute + } + if cfg.TaskMaxAttempts <= 0 { + cfg.TaskMaxAttempts = 3 + } +} + func candidateConfigPaths(configPath string, local bool) []string { trimmed := strings.TrimSpace(configPath) if trimmed == "" { diff --git a/server/internal/shared/config/config_test.go b/server/internal/shared/config/config_test.go index 1af78e0..7313776 100644 --- a/server/internal/shared/config/config_test.go +++ b/server/internal/shared/config/config_test.go @@ -145,6 +145,117 @@ scheduler: } } +func TestLoadAppliesRedisAndCacheDefaults(t *testing.T) { + configPath := writeTestConfig(t, ` +redis: {} +cache: {} +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.Redis.Addr != "localhost:6379" { + t.Fatalf("expected default redis addr, got %q", cfg.Redis.Addr) + } + if cfg.Redis.DialTimeout != 5*time.Second { + t.Fatalf("expected default redis dial timeout 5s, got %s", cfg.Redis.DialTimeout) + } + if cfg.Redis.ReadTimeout != 3*time.Second { + t.Fatalf("expected default redis read timeout 3s, got %s", cfg.Redis.ReadTimeout) + } + if cfg.Redis.WriteTimeout != 3*time.Second { + t.Fatalf("expected default redis write timeout 3s, got %s", cfg.Redis.WriteTimeout) + } + if cfg.Redis.PoolSize != 16 { + t.Fatalf("expected default redis pool size 16, got %d", cfg.Redis.PoolSize) + } + if cfg.Redis.PoolTimeout != 4*time.Second { + t.Fatalf("expected default redis pool timeout 4s, got %s", cfg.Redis.PoolTimeout) + } + if cfg.Cache.Driver != "redis" { + t.Fatalf("expected default cache driver redis, got %q", cfg.Cache.Driver) + } + if cfg.Cache.Namespace != "geo" { + t.Fatalf("expected default cache namespace geo, got %q", cfg.Cache.Namespace) + } + if cfg.Cache.JitterRatio != 0.1 { + t.Fatalf("expected default jitter ratio 0.1, got %f", cfg.Cache.JitterRatio) + } + if !cfg.Cache.MetricsEnabled { + t.Fatalf("expected cache metrics enabled by default") + } + if cfg.Cache.AsyncFillWorkers != 2 || cfg.Cache.AsyncFillBuffer != 1024 || cfg.Cache.AsyncFillTimeout != 2*time.Second { + t.Fatalf("unexpected async fill defaults: workers=%d buffer=%d timeout=%s", cfg.Cache.AsyncFillWorkers, cfg.Cache.AsyncFillBuffer, cfg.Cache.AsyncFillTimeout) + } + if cfg.Cache.L1TTL != 30*time.Second { + t.Fatalf("expected default l1 ttl 30s, got %s", cfg.Cache.L1TTL) + } + if cfg.Cache.DeleteScanCount != 500 { + t.Fatalf("expected default delete scan count 500, got %d", cfg.Cache.DeleteScanCount) + } +} + +func TestLoadAppliesRedisAndCacheOverrides(t *testing.T) { + configPath := writeTestConfig(t, ` +redis: + addr: redis.internal:6380 + password: secret + db: 2 + dial_timeout: 7s + read_timeout: 8s + write_timeout: 9s + pool_size: 32 + min_idle_conns: 4 + max_retries: 5 + pool_timeout: 6s + tls_enabled: true +cache: + driver: memory + namespace: tenant-cache + jitter_ratio: 0.25 + metrics_enabled: false + async_fill_enabled: true + async_fill_workers: 3 + async_fill_buffer: 64 + async_fill_timeout: 1500ms + l1_enabled: true + l1_ttl: 5s + delete_scan_count: 1200 +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.Redis.Addr != "redis.internal:6380" || cfg.Redis.Password != "secret" || cfg.Redis.DB != 2 { + t.Fatalf("unexpected redis endpoint config: %#v", cfg.Redis) + } + if cfg.Redis.DialTimeout != 7*time.Second || cfg.Redis.ReadTimeout != 8*time.Second || cfg.Redis.WriteTimeout != 9*time.Second { + t.Fatalf("unexpected redis timeout config: %#v", cfg.Redis) + } + if cfg.Redis.PoolSize != 32 || cfg.Redis.MinIdleConns != 4 || cfg.Redis.MaxRetries != 5 || cfg.Redis.PoolTimeout != 6*time.Second { + t.Fatalf("unexpected redis pool config: %#v", cfg.Redis) + } + if cfg.Redis.TLSConfig() == nil { + t.Fatalf("expected redis TLS config") + } + if cfg.Cache.Driver != "memory" || cfg.Cache.Namespace != "tenant-cache" || cfg.Cache.JitterRatio != 0.25 { + t.Fatalf("unexpected cache base config: %#v", cfg.Cache) + } + if cfg.Cache.MetricsEnabled { + t.Fatalf("expected cache metrics override false") + } + if !cfg.Cache.AsyncFillEnabled || cfg.Cache.AsyncFillWorkers != 3 || cfg.Cache.AsyncFillBuffer != 64 || cfg.Cache.AsyncFillTimeout != 1500*time.Millisecond { + t.Fatalf("unexpected cache async config: %#v", cfg.Cache) + } + if !cfg.Cache.L1Enabled || cfg.Cache.L1TTL != 5*time.Second || cfg.Cache.DeleteScanCount != 1200 { + t.Fatalf("unexpected cache l1/delete config: %#v", cfg.Cache) + } +} + func TestLoadPrefersEnvSchedulerMetricsSettings(t *testing.T) { t.Setenv("SCHEDULER_HTTP_HOST", "0.0.0.0") t.Setenv("SCHEDULER_INTERNAL_METRICS_TOKEN", "env-metrics-token") @@ -281,6 +392,75 @@ generation: } } +func TestLoadAppliesGenerationRecoveryDefaults(t *testing.T) { + configPath := writeTestConfig(t, ` +generation: + article_timeout: 90s +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.Generation.QueueSize != 128 { + t.Fatalf("expected default queue size 128, got %d", cfg.Generation.QueueSize) + } + if cfg.Generation.WorkerConcurrency != 1 { + t.Fatalf("expected default worker concurrency 1, got %d", cfg.Generation.WorkerConcurrency) + } + if cfg.Generation.TaskLeaseTTL != 3*time.Minute+30*time.Second { + t.Fatalf("expected default task lease ttl article_timeout+2m, got %s", cfg.Generation.TaskLeaseTTL) + } + if cfg.Generation.TaskRecoveryInterval != time.Minute { + t.Fatalf("expected default task recovery interval 1m, got %s", cfg.Generation.TaskRecoveryInterval) + } + if cfg.Generation.TaskRecoveryTimeout != 30*time.Second { + t.Fatalf("expected default task recovery timeout 30s, got %s", cfg.Generation.TaskRecoveryTimeout) + } + if cfg.Generation.TaskRecoveryBatchSize != 100 { + t.Fatalf("expected default task recovery batch size 100, got %d", cfg.Generation.TaskRecoveryBatchSize) + } + if cfg.Generation.TaskQueuedStaleAfter != 2*time.Minute { + t.Fatalf("expected default queued stale after 2m, got %s", cfg.Generation.TaskQueuedStaleAfter) + } + if cfg.Generation.TaskMaxAttempts != 3 { + t.Fatalf("expected default task max attempts 3, got %d", cfg.Generation.TaskMaxAttempts) + } +} + +func TestLoadAppliesGenerationRecoveryOverrides(t *testing.T) { + configPath := writeTestConfig(t, ` +generation: + queue_size: 64 + worker_concurrency: 4 + article_timeout: 5m + task_lease_ttl: 7m + task_recovery_interval: 15s + task_recovery_timeout: 10s + task_recovery_batch_size: 12 + task_queued_stale_after: 45s + task_max_attempts: 5 +`) + + cfg, err := Load(configPath) + if err != nil { + t.Fatalf("Load() error = %v", err) + } + + if cfg.Generation.QueueSize != 64 || cfg.Generation.WorkerConcurrency != 4 || cfg.Generation.ArticleTimeout != 5*time.Minute { + t.Fatalf("unexpected generation base config: %#v", cfg.Generation) + } + if cfg.Generation.TaskLeaseTTL != 7*time.Minute || + cfg.Generation.TaskRecoveryInterval != 15*time.Second || + cfg.Generation.TaskRecoveryTimeout != 10*time.Second || + cfg.Generation.TaskRecoveryBatchSize != 12 || + cfg.Generation.TaskQueuedStaleAfter != 45*time.Second || + cfg.Generation.TaskMaxAttempts != 5 { + t.Fatalf("unexpected generation recovery config: %#v", cfg.Generation) + } +} + func TestStoreReloadsConfigAndKeepsPreviousOnInvalidYAML(t *testing.T) { configPath := writeTestConfig(t, ` jwt: diff --git a/server/internal/shared/config/reload.go b/server/internal/shared/config/reload.go index 6aa4016..52a3d40 100644 --- a/server/internal/shared/config/reload.go +++ b/server/internal/shared/config/reload.go @@ -85,7 +85,13 @@ func Diff(previous, current *Config) []FieldChange { addChange("generation.worker", false) } if previous.Generation.StreamEnabled != current.Generation.StreamEnabled || - previous.Generation.ArticleTimeout != current.Generation.ArticleTimeout { + previous.Generation.ArticleTimeout != current.Generation.ArticleTimeout || + previous.Generation.TaskLeaseTTL != current.Generation.TaskLeaseTTL || + previous.Generation.TaskRecoveryInterval != current.Generation.TaskRecoveryInterval || + previous.Generation.TaskRecoveryTimeout != current.Generation.TaskRecoveryTimeout || + previous.Generation.TaskRecoveryBatchSize != current.Generation.TaskRecoveryBatchSize || + previous.Generation.TaskQueuedStaleAfter != current.Generation.TaskQueuedStaleAfter || + previous.Generation.TaskMaxAttempts != current.Generation.TaskMaxAttempts { addChange("generation", true) } diff --git a/server/internal/shared/repository/redis/redis.go b/server/internal/shared/repository/redis/redis.go index eb744cf..e1d9e9f 100644 --- a/server/internal/shared/repository/redis/redis.go +++ b/server/internal/shared/repository/redis/redis.go @@ -20,8 +20,17 @@ func NewClient(ctx context.Context, cfg config.RedisConfig) (*goredis.Client, er func NewLazyClient(cfg config.RedisConfig) *goredis.Client { client := goredis.NewClient(&goredis.Options{ - Addr: cfg.Addr, - DB: cfg.DB, + Addr: cfg.Addr, + Password: cfg.Password, + DB: cfg.DB, + DialTimeout: cfg.DialTimeout, + ReadTimeout: cfg.ReadTimeout, + WriteTimeout: cfg.WriteTimeout, + PoolSize: cfg.PoolSize, + MinIdleConns: cfg.MinIdleConns, + MaxRetries: cfg.MaxRetries, + PoolTimeout: cfg.PoolTimeout, + TLSConfig: cfg.TLSConfig(), }) return client } diff --git a/server/internal/tenant/app/article_service.go b/server/internal/tenant/app/article_service.go index 0fd3dae..d57bc1d 100644 --- a/server/internal/tenant/app/article_service.go +++ b/server/internal/tenant/app/article_service.go @@ -136,7 +136,8 @@ type ArticleListResponse struct { func (s *ArticleService) List(ctx context.Context, params ArticleListParams) (*ArticleListResponse, error) { actor := auth.MustActor(ctx) - return sharedcache.LoadJSON(ctx, s.cache, &s.cacheGroup, articleListCacheKey(actor.TenantID, params), defaultCacheTTL(), func(loadCtx context.Context) (*ArticleListResponse, error) { + version := articleCacheVersion(ctx, s.cache, actor.TenantID) + return sharedcache.LoadJSON(ctx, s.cache, &s.cacheGroup, articleListCacheKey(actor.TenantID, version, params), defaultCacheTTL(), func(loadCtx context.Context) (*ArticleListResponse, error) { return s.loadArticles(loadCtx, actor.TenantID, params) }) } @@ -167,7 +168,8 @@ type ArticleDetailResponse struct { func (s *ArticleService) Detail(ctx context.Context, id int64) (*ArticleDetailResponse, error) { actor := auth.MustActor(ctx) - record, found, err := sharedcache.LoadJSONWithEmpty(ctx, s.cache, &s.cacheGroup, articleDetailCacheKey(actor.TenantID, id), defaultCacheTTL(), defaultCacheEmptyTTL(), func(loadCtx context.Context) (*ArticleDetailResponse, bool, error) { + version := articleCacheVersion(ctx, s.cache, actor.TenantID) + record, found, err := sharedcache.LoadJSONWithEmpty(ctx, s.cache, &s.cacheGroup, articleDetailCacheKey(actor.TenantID, id, version), defaultCacheTTL(), defaultCacheEmptyTTL(), func(loadCtx context.Context) (*ArticleDetailResponse, bool, error) { return s.loadArticleDetail(loadCtx, actor.TenantID, id) }) if err != nil { diff --git a/server/internal/tenant/app/cache_support.go b/server/internal/tenant/app/cache_support.go index 660d8da..9898926 100644 --- a/server/internal/tenant/app/cache_support.go +++ b/server/internal/tenant/app/cache_support.go @@ -6,14 +6,18 @@ import ( "encoding/hex" "encoding/json" "fmt" + "strings" "time" + "github.com/google/uuid" + sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache" ) const ( defaultReadCacheTTL = 5 * time.Minute defaultReadCacheEmptyTTL = 1 * time.Minute + articleCacheVersionTTL = 24 * time.Hour ) func defaultCacheTTL() time.Duration { @@ -203,18 +207,54 @@ func articleDetailCachePrefix(tenantID int64) string { return fmt.Sprintf("article:detail:%d:", tenantID) } -func articleDetailCacheKey(tenantID, articleID int64) string { - return fmt.Sprintf("%s%d", articleDetailCachePrefix(tenantID), articleID) +func articleDetailCacheArticlePrefix(tenantID, articleID int64) string { + return fmt.Sprintf("%s%d:", articleDetailCachePrefix(tenantID), articleID) } -func articleListCacheKey(tenantID int64, params interface{}) string { - return articleListCachePrefix(tenantID) + digestCacheKey(params) +func articleDetailCacheKey(tenantID, articleID int64, version string) string { + return fmt.Sprintf("%s%s", articleDetailCacheArticlePrefix(tenantID, articleID), normalizeArticleCacheVersion(version)) +} + +func articleListCacheKey(tenantID int64, version string, params interface{}) string { + return articleListCachePrefix(tenantID) + normalizeArticleCacheVersion(version) + ":" + digestCacheKey(params) +} + +func articleCacheVersionKey(tenantID int64) string { + return fmt.Sprintf("cache:version:article:%d", tenantID) +} + +func articleCacheVersion(ctx context.Context, c sharedcache.Cache, tenantID int64) string { + if c == nil || tenantID <= 0 { + return "0" + } + raw, err := c.Get(ctx, articleCacheVersionKey(tenantID)) + if err != nil || len(raw) == 0 { + return "0" + } + return normalizeArticleCacheVersion(string(raw)) +} + +func bumpArticleCacheVersion(ctx context.Context, c sharedcache.Cache, tenantID int64) { + if c == nil || tenantID <= 0 { + return + } + version := uuid.NewString() + _ = c.Set(ctx, articleCacheVersionKey(tenantID), []byte(version), articleCacheVersionTTL) +} + +func normalizeArticleCacheVersion(version string) string { + version = strings.TrimSpace(version) + if version == "" { + return "0" + } + return version } func invalidateArticleCaches(ctx context.Context, c sharedcache.Cache, tenantID int64, articleID *int64) { + bumpArticleCacheVersion(ctx, c, tenantID) deleteCachePrefix(ctx, c, articleListCachePrefix(tenantID)) if articleID != nil && *articleID > 0 { - deleteCacheKey(ctx, c, articleDetailCacheKey(tenantID, *articleID)) + deleteCachePrefix(ctx, c, articleDetailCacheArticlePrefix(tenantID, *articleID)) } else { deleteCachePrefix(ctx, c, articleDetailCachePrefix(tenantID)) } diff --git a/server/internal/tenant/app/cache_support_test.go b/server/internal/tenant/app/cache_support_test.go new file mode 100644 index 0000000..01b7e97 --- /dev/null +++ b/server/internal/tenant/app/cache_support_test.go @@ -0,0 +1,50 @@ +package app + +import ( + "context" + "errors" + "testing" + "time" + + sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache" +) + +func TestArticleCacheVersionSeparatesLateWritesAfterInvalidation(t *testing.T) { + t.Parallel() + + ctx := context.Background() + c := sharedcache.New("memory", nil, sharedcache.WithNamespace("test")) + tenantID := int64(42) + articleID := int64(7) + params := ArticleListParams{Page: 1, PageSize: 10} + + oldVersion := articleCacheVersion(ctx, c, tenantID) + oldListKey := articleListCacheKey(tenantID, oldVersion, params) + oldDetailKey := articleDetailCacheKey(tenantID, articleID, oldVersion) + if err := c.Set(ctx, oldListKey, []byte("generating-list"), time.Minute); err != nil { + t.Fatalf("set old list cache: %v", err) + } + if err := c.Set(ctx, oldDetailKey, []byte("generating-detail"), time.Minute); err != nil { + t.Fatalf("set old detail cache: %v", err) + } + + invalidateArticleCaches(ctx, c, tenantID, &articleID) + newVersion := articleCacheVersion(ctx, c, tenantID) + if newVersion == oldVersion { + t.Fatalf("expected article cache version to change, stayed %q", newVersion) + } + + if err := c.Set(ctx, oldListKey, []byte("late-generating-list"), time.Minute); err != nil { + t.Fatalf("late set old list cache: %v", err) + } + if err := c.Set(ctx, oldDetailKey, []byte("late-generating-detail"), time.Minute); err != nil { + t.Fatalf("late set old detail cache: %v", err) + } + + if _, err := c.Get(ctx, articleListCacheKey(tenantID, newVersion, params)); !errors.Is(err, sharedcache.ErrNotFound) { + t.Fatalf("expected new list key to miss stale late write, got err %v", err) + } + if _, err := c.Get(ctx, articleDetailCacheKey(tenantID, articleID, newVersion)); !errors.Is(err, sharedcache.ErrNotFound) { + t.Fatalf("expected new detail key to miss stale late write, got err %v", err) + } +} diff --git a/server/internal/tenant/app/generation_task_runtime.go b/server/internal/tenant/app/generation_task_runtime.go index 2a9e651..a9d348d 100644 --- a/server/internal/tenant/app/generation_task_runtime.go +++ b/server/internal/tenant/app/generation_task_runtime.go @@ -52,6 +52,9 @@ func HandleClaimedGenerationTaskFailure( defer cancel() errMsg := formatGenerationFailure("task_load", taskErr) + if taskErr != nil && strings.Contains(taskErr.Error(), "[worker_recovery]") { + errMsg = strings.TrimSpace(taskErr.Error()) + } now := time.Now() auditRepo := repository.NewAuditRepository(pool) diff --git a/server/internal/tenant/app/knowledge_service.go b/server/internal/tenant/app/knowledge_service.go index 5508551..3e17e1e 100644 --- a/server/internal/tenant/app/knowledge_service.go +++ b/server/internal/tenant/app/knowledge_service.go @@ -933,7 +933,7 @@ func (s *KnowledgeService) logKnowledgeResolve( zap.Int("candidate_count", len(candidates)), zap.Int("selected_count", len(selected)), zap.Int("precise_fact_count", len(knowledgeCtx.PreciseFacts)), - zap.Strings("precise_facts", summarizeKnowledgeFactsForLog(knowledgeCtx.PreciseFacts)), + // zap.Strings("precise_facts", summarizeKnowledgeFactsForLog(knowledgeCtx.PreciseFacts)), ) } diff --git a/server/internal/tenant/app/monitoring_metrics_prometheus.go b/server/internal/tenant/app/monitoring_metrics_prometheus.go index 5c30bb9..3df0e19 100644 --- a/server/internal/tenant/app/monitoring_metrics_prometheus.go +++ b/server/internal/tenant/app/monitoring_metrics_prometheus.go @@ -9,10 +9,12 @@ import ( promcollectors "github.com/prometheus/client_golang/prometheus/collectors" "github.com/prometheus/client_golang/prometheus/promhttp" "github.com/prometheus/common/expfmt" + + "github.com/geo-platform/tenant-api/internal/shared/cache" ) func MonitoringSchedulerMetricsPrometheusHandler() http.Handler { - return monitoringPrometheusHandler(true, monitoringDailyTaskMetricsCollector{}) + return monitoringPrometheusHandler(true, monitoringDailyTaskMetricsCollector{}, cache.MetricsCollector{}) } func MonitoringDailyTaskMetricsPrometheus() string { diff --git a/server/internal/tenant/repository/generated/audit.sql.go b/server/internal/tenant/repository/generated/audit.sql.go index 3e11b47..a2b88d1 100644 --- a/server/internal/tenant/repository/generated/audit.sql.go +++ b/server/internal/tenant/repository/generated/audit.sql.go @@ -125,7 +125,12 @@ func (q *Queries) CreateTaskRecord(ctx context.Context, arg CreateTaskRecordPara const updateGenerationTaskStatus = `-- name: UpdateGenerationTaskStatus :exec UPDATE generation_tasks SET status = $1, error_message = $2, - started_at = $3, completed_at = $4, updated_at = NOW() + started_at = $3, + completed_at = $4, + lease_token = CASE WHEN $1::text IN ('completed', 'failed') THEN NULL ELSE lease_token END, + lease_owner = CASE WHEN $1::text IN ('completed', 'failed') THEN NULL ELSE lease_owner END, + lease_expires_at = CASE WHEN $1::text IN ('completed', 'failed') THEN NULL ELSE lease_expires_at END, + updated_at = NOW() WHERE id = $5 AND tenant_id = $6 ` diff --git a/server/internal/tenant/repository/generated/models.go b/server/internal/tenant/repository/generated/models.go index 4b47543..9302e97 100644 --- a/server/internal/tenant/repository/generated/models.go +++ b/server/internal/tenant/repository/generated/models.go @@ -265,6 +265,11 @@ type GenerationTask struct { CreatedAt pgtype.Timestamptz `json:"created_at"` UpdatedAt pgtype.Timestamptz `json:"updated_at"` OperatorID pgtype.Int8 `json:"operator_id"` + LeaseToken pgtype.Text `json:"lease_token"` + LeaseOwner pgtype.Text `json:"lease_owner"` + LeaseExpiresAt pgtype.Timestamptz `json:"lease_expires_at"` + AttemptCount int32 `json:"attempt_count"` + LastHeartbeatAt pgtype.Timestamptz `json:"last_heartbeat_at"` } type ImageAsset struct { diff --git a/server/internal/tenant/repository/queries/audit.sql b/server/internal/tenant/repository/queries/audit.sql index 8e6d61f..9156dc0 100644 --- a/server/internal/tenant/repository/queries/audit.sql +++ b/server/internal/tenant/repository/queries/audit.sql @@ -40,5 +40,10 @@ WHERE tenant_id = sqlc.arg(tenant_id) -- name: UpdateGenerationTaskStatus :exec UPDATE generation_tasks SET status = sqlc.arg(status), error_message = sqlc.arg(error_message), - started_at = sqlc.arg(started_at), completed_at = sqlc.arg(completed_at), updated_at = NOW() + started_at = sqlc.arg(started_at), + completed_at = sqlc.arg(completed_at), + lease_token = CASE WHEN sqlc.arg(status)::text IN ('completed', 'failed') THEN NULL ELSE lease_token END, + lease_owner = CASE WHEN sqlc.arg(status)::text IN ('completed', 'failed') THEN NULL ELSE lease_owner END, + lease_expires_at = CASE WHEN sqlc.arg(status)::text IN ('completed', 'failed') THEN NULL ELSE lease_expires_at END, + updated_at = NOW() WHERE id = sqlc.arg(id) AND tenant_id = sqlc.arg(tenant_id); diff --git a/server/internal/worker/generate/article_generation_recovery.go b/server/internal/worker/generate/article_generation_recovery.go new file mode 100644 index 0000000..1545043 --- /dev/null +++ b/server/internal/worker/generate/article_generation_recovery.go @@ -0,0 +1,721 @@ +package generate + +import ( + "context" + "crypto/rand" + "database/sql" + "encoding/hex" + "encoding/json" + "fmt" + "os" + "strings" + "time" + + "github.com/jackc/pgx/v5" + "go.uber.org/zap" + + "github.com/geo-platform/tenant-api/internal/shared/config" + tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" +) + +const ( + defaultGenerationTaskRecoveryInterval = time.Minute + defaultGenerationTaskRecoveryTimeout = 30 * time.Second +) + +type generationTaskRecoveryConfig struct { + LeaseTTL time.Duration + RecoveryInterval time.Duration + RecoveryTimeout time.Duration + BatchSize int + QueuedStaleAfter time.Duration + MaxAttempts int +} + +type claimedGenerationTaskLease struct { + task generationTaskRecoveryRecord + token string +} + +type generationTaskRecoveryRecord struct { + ID int64 + TenantID int64 + OperatorID int64 + ArticleID int64 + QuotaReservationID int64 + TaskType string + InputParams map[string]interface{} + AttemptCount int +} + +func (r generationTaskRecoveryRecord) toClaimedTask() tenantapp.ClaimedGenerationTask { + return tenantapp.ClaimedGenerationTask{ + ID: r.ID, + TenantID: r.TenantID, + OperatorID: r.OperatorID, + ArticleID: r.ArticleID, + QuotaReservationID: r.QuotaReservationID, + TaskType: r.TaskType, + InputParams: r.InputParams, + } +} + +type recoveredGenerationTask struct { + ID int64 + TenantID int64 + ArticleID int64 + TaskType string + AttemptCount int +} + +func (w *ArticleGenerationWorker) startRecovery(ctx context.Context) { + if w == nil || w.pool == nil { + return + } + go w.runRecovery(ctx) +} + +func (w *ArticleGenerationWorker) runRecovery(ctx context.Context) { + w.runRecoveryOnce(context.Background()) + + for { + interval := w.runtimeRecoveryConfig().RecoveryInterval + timer := time.NewTimer(interval) + select { + case <-ctx.Done(): + timer.Stop() + return + case <-timer.C: + w.runRecoveryOnce(context.Background()) + } + } +} + +func (w *ArticleGenerationWorker) runRecoveryOnce(parent context.Context) { + cfg := w.runtimeRecoveryConfig() + ctx, cancel := context.WithTimeout(parent, cfg.RecoveryTimeout) + defer cancel() + + queuedRecovered, requeueErr := w.requeueStaleGenerationTasks(ctx, cfg) + if requeueErr != nil && w.logger != nil { + w.logger.Warn("article generation queued task recovery failed", zap.Error(requeueErr)) + } + + expired, expireErr := w.expireStaleRunningGenerationTasks(ctx, cfg) + if expireErr != nil && w.logger != nil { + w.logger.Warn("article generation running task recovery failed", zap.Error(expireErr)) + } + + if w.logger != nil && (queuedRecovered > 0 || expired > 0) { + w.logger.Info("article generation task recovery completed", + zap.Int("queued_recovered_task_count", queuedRecovered), + zap.Int("running_recovered_task_count", expired), + ) + } +} + +func (w *ArticleGenerationWorker) runtimeRecoveryConfig() generationTaskRecoveryConfig { + if w != nil && w.configProvider != nil { + if cfg := w.configProvider.Current(); cfg != nil { + return normalizeGenerationTaskRecoveryConfig(cfg.Generation) + } + } + if w != nil { + return normalizeGenerationTaskRecoveryConfig(w.generationCfg) + } + return normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{}) +} + +func normalizeGenerationTaskRecoveryConfig(cfg config.GenerationConfig) generationTaskRecoveryConfig { + normalized := cfg + config.NormalizeGenerationConfig(&normalized) + return generationTaskRecoveryConfig{ + LeaseTTL: normalized.TaskLeaseTTL, + RecoveryInterval: normalized.TaskRecoveryInterval, + RecoveryTimeout: normalized.TaskRecoveryTimeout, + BatchSize: normalized.TaskRecoveryBatchSize, + QueuedStaleAfter: normalized.TaskQueuedStaleAfter, + MaxAttempts: normalized.TaskMaxAttempts, + } +} + +func (w *ArticleGenerationWorker) claimGenerationTask(ctx context.Context, taskID int64, owner string, leaseTTL time.Duration) (*claimedGenerationTaskLease, bool, error) { + if leaseTTL <= 0 { + leaseTTL = normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{}).LeaseTTL + } + token, err := newGenerationTaskLeaseToken() + if err != nil { + return nil, false, err + } + leaseExpiresAt := time.Now().UTC().Add(leaseTTL) + + var ( + row generationTaskRecoveryRecord + operatorID sql.NullInt64 + articleID sql.NullInt64 + reservationID sql.NullInt64 + inputParamsJSON []byte + ) + err = w.pool.QueryRow(ctx, ` + UPDATE generation_tasks + SET status = 'running', + started_at = COALESCE(started_at, NOW()), + lease_token = $2, + lease_owner = $3, + lease_expires_at = $4, + last_heartbeat_at = NOW(), + attempt_count = attempt_count + 1, + updated_at = NOW() + WHERE id = $1 + AND status = 'queued' + RETURNING id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, attempt_count + `, taskID, token, owner, leaseExpiresAt).Scan( + &row.ID, + &row.TenantID, + &operatorID, + &articleID, + &reservationID, + &row.TaskType, + &inputParamsJSON, + &row.AttemptCount, + ) + if err != nil { + if err == pgx.ErrNoRows { + return nil, true, nil + } + return nil, false, err + } + + row.OperatorID = operatorID.Int64 + row.ArticleID = articleID.Int64 + row.QuotaReservationID = reservationID.Int64 + params, err := parseJSONMap(inputParamsJSON) + if err != nil { + row.InputParams = map[string]interface{}{} + return &claimedGenerationTaskLease{task: row, token: token}, false, fmt.Errorf("decode generation input params: %w", err) + } + row.InputParams = params + + return &claimedGenerationTaskLease{task: row, token: token}, false, nil +} + +func (w *ArticleGenerationWorker) renewGenerationTaskLease(ctx context.Context, taskID int64, token string, leaseTTL time.Duration) (bool, error) { + if taskID <= 0 || token == "" { + return false, nil + } + if leaseTTL <= 0 { + leaseTTL = normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{}).LeaseTTL + } + tag, err := w.pool.Exec(ctx, ` + UPDATE generation_tasks + SET lease_expires_at = $3, + last_heartbeat_at = NOW(), + updated_at = NOW() + WHERE id = $1 + AND lease_token = $2 + AND status = 'running' + `, taskID, token, time.Now().UTC().Add(leaseTTL)) + if err != nil { + return false, err + } + return tag.RowsAffected() > 0, nil +} + +func (w *ArticleGenerationWorker) releaseGenerationTaskLease(ctx context.Context, taskID int64, token string) { + if w == nil || w.pool == nil || taskID <= 0 || token == "" { + return + } + releaseCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + defer cancel() + + if _, err := w.pool.Exec(releaseCtx, ` + UPDATE generation_tasks + SET lease_token = NULL, + lease_owner = NULL, + lease_expires_at = NULL, + updated_at = NOW() + WHERE id = $1 + AND lease_token = $2 + AND status IN ('completed', 'failed') + `, taskID, token); err != nil && w.logger != nil { + w.logger.Warn("article generation task lease release failed", + zap.Error(err), + zap.Int64("task_id", taskID), + ) + } +} + +func (w *ArticleGenerationWorker) runGenerationLeaseHeartbeat(ctx context.Context, taskID int64, token string, leaseTTL time.Duration) { + if taskID <= 0 || token == "" || leaseTTL <= 0 { + return + } + + interval := leaseTTL / 3 + if interval < 10*time.Second { + interval = 10 * time.Second + } + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-ctx.Done(): + return + case <-ticker.C: + heartbeatCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second) + ok, err := w.renewGenerationTaskLease(heartbeatCtx, taskID, token, leaseTTL) + cancel() + if err != nil { + if w.logger != nil { + w.logger.Warn("article generation task lease heartbeat failed", + zap.Error(err), + zap.Int64("task_id", taskID), + ) + } + continue + } + if !ok { + if w.logger != nil { + w.logger.Warn("article generation task lease lost", + zap.Int64("task_id", taskID), + ) + } + return + } + } + } +} + +func (w *ArticleGenerationWorker) requeueStaleGenerationTasks(ctx context.Context, cfg generationTaskRecoveryConfig) (int, error) { + if w == nil || w.pool == nil || w.rabbitMQ == nil || cfg.BatchSize <= 0 { + return 0, nil + } + staleBefore := time.Now().UTC().Add(-cfg.QueuedStaleAfter) + + rows, err := w.pool.Query(ctx, ` + WITH candidates AS ( + SELECT id, tenant_id, article_id, task_type, attempt_count + FROM generation_tasks + WHERE status = 'queued' + AND article_id IS NOT NULL + AND updated_at < $1 + ORDER BY updated_at ASC, id ASC + FOR UPDATE SKIP LOCKED + LIMIT $2 + ), + touched AS ( + UPDATE generation_tasks gt + SET updated_at = NOW() + FROM candidates c + WHERE gt.id = c.id + RETURNING c.id, c.tenant_id, c.article_id, c.task_type, c.attempt_count + ) + SELECT id, tenant_id, article_id, task_type, attempt_count + FROM touched + `, staleBefore, cfg.BatchSize) + if err != nil { + return 0, err + } + defer rows.Close() + + items := make([]recoveredGenerationTask, 0, cfg.BatchSize) + for rows.Next() { + var item recoveredGenerationTask + if err := rows.Scan(&item.ID, &item.TenantID, &item.ArticleID, &item.TaskType, &item.AttemptCount); err != nil { + return 0, err + } + items = append(items, item) + } + if err := rows.Err(); err != nil { + return 0, err + } + + requeued := 0 + for _, item := range items { + if item.AttemptCount >= cfg.MaxAttempts { + ok, err := w.finalizeExpiredGenerationTask(ctx, generationTaskRecoveryRecord{ + ID: item.ID, + TenantID: item.TenantID, + ArticleID: item.ArticleID, + TaskType: item.TaskType, + AttemptCount: item.AttemptCount, + }) + if err != nil && w.logger != nil { + w.logger.Warn("article generation stale queued task finalization failed", + zap.Error(err), + zap.Int64("task_id", item.ID), + zap.Int64("article_id", item.ArticleID), + ) + } + if ok { + requeued++ + } + continue + } + if err := w.publishGenerationTask(ctx, item.ID, item.TaskType, item.TenantID, item.ArticleID); err != nil { + if w.logger != nil { + w.logger.Warn("article generation stale queued task republish failed", + zap.Error(err), + zap.Int64("task_id", item.ID), + zap.Int64("article_id", item.ArticleID), + ) + } + continue + } + requeued++ + } + + return requeued, nil +} + +func (w *ArticleGenerationWorker) expireStaleRunningGenerationTasks(ctx context.Context, cfg generationTaskRecoveryConfig) (int, error) { + if w == nil || w.pool == nil || cfg.BatchSize <= 0 { + return 0, nil + } + now := time.Now().UTC() + + rows, err := w.pool.Query(ctx, ` + SELECT id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, attempt_count + FROM generation_tasks + WHERE status = 'running' + AND article_id IS NOT NULL + AND ( + lease_expires_at IS NULL + OR lease_expires_at < $1 + ) + ORDER BY COALESCE(lease_expires_at, started_at, updated_at, created_at) ASC, id ASC + LIMIT $2 + `, now, cfg.BatchSize) + if err != nil { + return 0, err + } + defer rows.Close() + + recovered := make([]generationTaskRecoveryRecord, 0, cfg.BatchSize) + for rows.Next() { + var ( + item generationTaskRecoveryRecord + operatorID sql.NullInt64 + articleID sql.NullInt64 + reservationID sql.NullInt64 + inputParamsJSON []byte + ) + if err := rows.Scan( + &item.ID, + &item.TenantID, + &operatorID, + &articleID, + &reservationID, + &item.TaskType, + &inputParamsJSON, + &item.AttemptCount, + ); err != nil { + return 0, err + } + item.OperatorID = operatorID.Int64 + item.ArticleID = articleID.Int64 + item.QuotaReservationID = reservationID.Int64 + params, err := parseJSONMap(inputParamsJSON) + if err != nil { + params = map[string]interface{}{} + } + item.InputParams = params + recovered = append(recovered, item) + } + if err := rows.Err(); err != nil { + return 0, err + } + + handled := 0 + for _, item := range recovered { + if item.AttemptCount < cfg.MaxAttempts { + ok, err := w.requeueExpiredRunningGenerationTask(ctx, item) + if err != nil && w.logger != nil { + w.logger.Warn("article generation expired task requeue failed", + zap.Error(err), + zap.Int64("task_id", item.ID), + zap.Int64("article_id", item.ArticleID), + ) + } + if ok { + handled++ + } + continue + } + + ok, err := w.finalizeExpiredGenerationTask(ctx, item) + if err != nil && w.logger != nil { + w.logger.Warn("article generation expired task finalization failed", + zap.Error(err), + zap.Int64("task_id", item.ID), + zap.Int64("article_id", item.ArticleID), + ) + } + if ok { + handled++ + } + } + + return handled, nil +} + +func (w *ArticleGenerationWorker) requeueExpiredRunningGenerationTask(ctx context.Context, item generationTaskRecoveryRecord) (bool, error) { + if w == nil || w.rabbitMQ == nil { + return false, nil + } + tx, err := w.pool.Begin(ctx) + if err != nil { + return false, err + } + defer func() { + _ = tx.Rollback(ctx) + }() + + var ( + tenantID int64 + articleID int64 + taskType string + ) + if err := tx.QueryRow(ctx, ` + UPDATE generation_tasks + SET status = 'queued', + error_message = NULL, + lease_token = NULL, + lease_owner = NULL, + lease_expires_at = NULL, + updated_at = NOW() + WHERE id = $1 + AND status = 'running' + AND ( + lease_expires_at IS NULL + OR lease_expires_at < NOW() + ) + RETURNING tenant_id, article_id, task_type + `, item.ID).Scan(&tenantID, &articleID, &taskType); err != nil { + if err == pgx.ErrNoRows { + return false, nil + } + return false, err + } + + if err := tx.Commit(ctx); err != nil { + return false, err + } + if err := w.publishGenerationTask(ctx, item.ID, taskType, tenantID, articleID); err != nil { + return false, err + } + return true, nil +} + +func (w *ArticleGenerationWorker) finalizeExpiredGenerationTask(ctx context.Context, item generationTaskRecoveryRecord) (bool, error) { + errMsg := formatGenerationTaskRecoveryFailure() + + tx, err := w.pool.Begin(ctx) + if err != nil { + return false, err + } + defer func() { + _ = tx.Rollback(ctx) + }() + + var ( + claimedID int64 + operatorID sql.NullInt64 + articleID sql.NullInt64 + quotaReservationID sql.NullInt64 + taskType string + ) + if err := tx.QueryRow(ctx, ` + UPDATE generation_tasks + SET status = 'failed', + error_message = $2, + completed_at = COALESCE(completed_at, NOW()), + lease_token = NULL, + lease_owner = NULL, + lease_expires_at = NULL, + updated_at = NOW() + WHERE id = $1 + AND status IN ('running', 'queued') + AND ( + status = 'queued' + OR lease_expires_at IS NULL + OR lease_expires_at < NOW() + ) + RETURNING id, operator_id, article_id, quota_reservation_id, task_type + `, item.ID, errMsg).Scan(&claimedID, &operatorID, &articleID, "aReservationID, &taskType); err != nil { + if err == pgx.ErrNoRows { + return false, nil + } + return false, err + } + _ = claimedID + item.OperatorID = operatorID.Int64 + item.ArticleID = articleID.Int64 + item.QuotaReservationID = quotaReservationID.Int64 + item.TaskType = strings.TrimSpace(taskType) + + if item.ArticleID > 0 { + if _, err := tx.Exec(ctx, ` + UPDATE articles + SET generate_status = 'failed', + updated_at = NOW() + WHERE id = $1 + AND tenant_id = $2 + AND generate_status = 'generating' + AND deleted_at IS NULL + `, item.ArticleID, item.TenantID); err != nil { + return false, err + } + } + + if item.TaskType == kolGenerationTaskType { + if _, err := tx.Exec(ctx, ` + UPDATE task_records + SET status = 'failed', + error_message = $1, + updated_at = NOW() + WHERE tenant_id = $2 + AND task_type = $3 + AND resource_type = $4 + AND resource_id = $5 + `, errMsg, item.TenantID, "article_generation", "generation_task", item.ID); err != nil { + return false, err + } + if _, err := tx.Exec(ctx, ` + UPDATE kol_usage_logs + SET status = 'failed', + error_message = $1, + finished_at = COALESCE(finished_at, NOW()) + WHERE generation_task_id = $2 + AND tenant_id = $3 + AND status IN ('pending', 'running') + `, errMsg, item.ID, item.TenantID); err != nil { + return false, err + } + } + + if item.QuotaReservationID > 0 { + var reservationStatus string + if err := tx.QueryRow(ctx, ` + WITH refunded AS ( + UPDATE quota_reservations + SET status = 'refunded', + refunded_amount = reserved_amount, + updated_at = NOW() + WHERE id = $1 + AND tenant_id = $2 + AND status = 'pending' + RETURNING status + ) + SELECT status FROM refunded + `, item.QuotaReservationID, item.TenantID).Scan(&reservationStatus); err != nil && err != pgx.ErrNoRows { + return false, err + } + if reservationStatus == "refunded" { + if err := insertGenerationRecoveryRefundLedger(ctx, tx, item); err != nil { + return false, err + } + } + } + + if err := tx.Commit(ctx); err != nil { + return false, err + } + + cacheCtx, cancel := context.WithTimeout(context.Background(), 30*time.Second) + defer cancel() + if item.ArticleID > 0 { + tenantapp.InvalidateArticleCaches(cacheCtx, w.cache, item.TenantID, &item.ArticleID) + } + return true, nil +} + +func insertGenerationRecoveryRefundLedger(ctx context.Context, tx pgx.Tx, item generationTaskRecoveryRecord) error { + reason := "generation_refund" + referenceType := "generation_task" + var operatorID interface{} + if item.OperatorID > 0 { + operatorID = item.OperatorID + } + _, err := tx.Exec(ctx, ` + INSERT INTO tenant_quota_ledgers ( + tenant_id, + operator_id, + quota_type, + delta, + balance_after, + reason, + reference_type, + reference_id + ) + SELECT + $1, + $2, + 'article_generation', + 1, + COALESCE(( + SELECT balance_after + FROM tenant_quota_ledgers + WHERE tenant_id = $1 + AND quota_type = 'article_generation' + ORDER BY created_at DESC, id DESC + LIMIT 1 + ), 0) + 1, + $3, + $4, + $5 + WHERE NOT EXISTS ( + SELECT 1 + FROM tenant_quota_ledgers + WHERE tenant_id = $1 + AND quota_type = 'article_generation' + AND reason = $3 + AND reference_type = $4 + AND reference_id = $5 + ) + `, item.TenantID, operatorID, reason, referenceType, item.ID) + return err +} + +func formatGenerationTaskRecoveryFailure() string { + return "文章生成任务执行超时，worker 重启恢复后自动标记失败 [worker_recovery]" +} + +func (w *ArticleGenerationWorker) publishGenerationTask(ctx context.Context, taskID int64, taskType string, tenantID, articleID int64) error { + if w == nil || w.rabbitMQ == nil { + return nil + } + payload, err := json.Marshal(tenantapp.ArticleGenerationTaskEnvelope{ + TaskID: taskID, + TaskType: taskType, + TenantID: tenantID, + ArticleID: articleID, + }) + if err != nil { + return fmt.Errorf("marshal generation task envelope: %w", err) + } + if err := w.rabbitMQ.PublishGenerationTask(ctx, payload); err != nil { + return fmt.Errorf("publish generation task: %w", err) + } + return nil +} + +func newGenerationTaskLeaseToken() (string, error) { + var raw [16]byte + if _, err := rand.Read(raw[:]); err != nil { + return "", fmt.Errorf("create generation task lease token: %w", err) + } + return hex.EncodeToString(raw[:]), nil +} + +func generationWorkerOwner(consumerName string) string { + host, _ := os.Hostname() + owner := strings.TrimSpace(host + "/" + consumerName) + if owner == "/" { + return "worker-generate" + } + if len(owner) > 128 { + return owner[:128] + } + return owner +} diff --git a/server/internal/worker/generate/article_generation_recovery_test.go b/server/internal/worker/generate/article_generation_recovery_test.go new file mode 100644 index 0000000..6fae88c --- /dev/null +++ b/server/internal/worker/generate/article_generation_recovery_test.go @@ -0,0 +1,72 @@ +package generate + +import ( + "strings" + "testing" + "time" + + "github.com/geo-platform/tenant-api/internal/shared/config" +) + +func TestNormalizeGenerationTaskRecoveryConfigUsesGenerationDefaults(t *testing.T) { + cfg := normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{ + ArticleTimeout: 90 * time.Second, + }) + + if cfg.LeaseTTL != 3*time.Minute+30*time.Second { + t.Fatalf("expected lease ttl article_timeout+2m, got %s", cfg.LeaseTTL) + } + if cfg.RecoveryInterval != time.Minute { + t.Fatalf("expected recovery interval 1m, got %s", cfg.RecoveryInterval) + } + if cfg.RecoveryTimeout != 30*time.Second { + t.Fatalf("expected recovery timeout 30s, got %s", cfg.RecoveryTimeout) + } + if cfg.BatchSize != 100 { + t.Fatalf("expected batch size 100, got %d", cfg.BatchSize) + } + if cfg.QueuedStaleAfter != 2*time.Minute { + t.Fatalf("expected queued stale after 2m, got %s", cfg.QueuedStaleAfter) + } + if cfg.MaxAttempts != 3 { + t.Fatalf("expected max attempts 3, got %d", cfg.MaxAttempts) + } +} + +func TestNormalizeGenerationTaskRecoveryConfigKeepsOverrides(t *testing.T) { + cfg := normalizeGenerationTaskRecoveryConfig(config.GenerationConfig{ + ArticleTimeout: 5 * time.Minute, + TaskLeaseTTL: 7 * time.Minute, + TaskRecoveryInterval: 15 * time.Second, + TaskRecoveryTimeout: 10 * time.Second, + TaskRecoveryBatchSize: 12, + TaskQueuedStaleAfter: 45 * time.Second, + TaskMaxAttempts: 5, + }) + + if cfg.LeaseTTL != 7*time.Minute || + cfg.RecoveryInterval != 15*time.Second || + cfg.RecoveryTimeout != 10*time.Second || + cfg.BatchSize != 12 || + cfg.QueuedStaleAfter != 45*time.Second || + cfg.MaxAttempts != 5 { + t.Fatalf("unexpected recovery config: %#v", cfg) + } +} + +func TestGenerationWorkerOwnerIsBounded(t *testing.T) { + owner := generationWorkerOwner(strings.Repeat("x", 256)) + if len(owner) > 128 { + t.Fatalf("expected owner to be bounded to 128 chars, got %d", len(owner)) + } + if owner == "" { + t.Fatalf("expected non-empty owner") + } +} + +func TestGenerationTaskRecoveryFailureMessage(t *testing.T) { + message := formatGenerationTaskRecoveryFailure() + if !strings.Contains(message, "worker_recovery") { + t.Fatalf("expected recovery marker in failure message, got %q", message) + } +} diff --git a/server/internal/worker/generate/article_generation_worker.go b/server/internal/worker/generate/article_generation_worker.go index ba6ee81..94915d6 100644 --- a/server/internal/worker/generate/article_generation_worker.go +++ b/server/internal/worker/generate/article_generation_worker.go @@ -2,7 +2,6 @@ package generate import ( "context" - "database/sql" "encoding/json" "errors" "fmt" @@ -10,11 +9,11 @@ import ( "strings" "time" - "github.com/jackc/pgx/v5" "github.com/jackc/pgx/v5/pgxpool" amqp "github.com/rabbitmq/amqp091-go" "go.uber.org/zap" + sharedcache "github.com/geo-platform/tenant-api/internal/shared/cache" "github.com/geo-platform/tenant-api/internal/shared/config" "github.com/geo-platform/tenant-api/internal/shared/messaging/rabbitmq" tenantapp "github.com/geo-platform/tenant-api/internal/tenant/app" @@ -31,9 +30,11 @@ type ArticleGenerationWorker struct { promptRuleService *tenantapp.PromptRuleGenerationService imitationService *tenantapp.ArticleImitationService kolWorker *KolGenerationWorker + cache sharedcache.Cache logger *zap.Logger retryInterval time.Duration processTimeout time.Duration + generationCfg config.GenerationConfig configProvider config.Provider consumerPrefix string workerConcurrency int @@ -64,6 +65,7 @@ func NewArticleGenerationWorker( logger: logger, retryInterval: defaultArticleGenerationWorkerRetryInterval, processTimeout: articleGenerationProcessTimeout(cfg), + generationCfg: cfg, consumerPrefix: fmt.Sprintf("worker-generate-%d", os.Getpid()), workerConcurrency: workerConcurrency, } @@ -74,6 +76,11 @@ func (w *ArticleGenerationWorker) WithConfigProvider(provider config.Provider) * return w } +func (w *ArticleGenerationWorker) WithCache(c sharedcache.Cache) *ArticleGenerationWorker { + w.cache = c + return w +} + func (w *ArticleGenerationWorker) Start(ctx context.Context) { if w == nil || w.pool == nil || w.rabbitMQ == nil { return @@ -83,6 +90,7 @@ func (w *ArticleGenerationWorker) Start(ctx context.Context) { consumerName := fmt.Sprintf("%s-%d", w.consumerPrefix, i+1) go w.run(ctx, consumerName) } + w.startRecovery(ctx) } func (w *ArticleGenerationWorker) run(ctx context.Context, consumerName string) { @@ -135,31 +143,25 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver return } + recoveryCfg := w.runtimeRecoveryConfig() ctx, cancel := context.WithTimeout(parent, w.runtimeProcessTimeout()) defer cancel() - task, terminal, err := w.loadTask(ctx, envelope.TaskID) + lease, terminal, err := w.claimGenerationTask(ctx, envelope.TaskID, generationWorkerOwner(delivery.ConsumerTag), recoveryCfg.LeaseTTL) if err != nil { - if errors.Is(err, pgx.ErrNoRows) { - if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil { - w.logger.Warn("article generation ack failed for missing task", - zap.Error(ackErr), - zap.Int64("task_id", envelope.TaskID), - ) - } - return - } - if task.ID > 0 { + if lease != nil { + task := lease.task.toClaimedTask() w.handleTaskFailure(ctx, task, "task_load", err) + w.releaseGenerationTaskLease(context.Background(), task.ID, lease.token) if ackErr := delivery.Ack(false); ackErr != nil && w.logger != nil { - w.logger.Warn("article generation ack failed after task load failure", + w.logger.Warn("article generation ack failed after task claim failure", zap.Error(ackErr), zap.Int64("task_id", envelope.TaskID), ) } return } - requeue := !delivery.Redelivered + requeue := !delivery.Redelivered && !isPermanentGenerationTaskDecodeError(err) w.rejectDelivery(delivery, requeue, err, envelope.TaskID) return } @@ -172,6 +174,18 @@ func (w *ArticleGenerationWorker) handleDelivery(parent context.Context, deliver } return } + if lease == nil { + w.rejectDelivery(delivery, !delivery.Redelivered, fmt.Errorf("generation task claim returned empty lease"), envelope.TaskID) + return + } + + task := lease.task.toClaimedTask() + heartbeatCtx, stopHeartbeat := context.WithCancel(parent) + go w.runGenerationLeaseHeartbeat(heartbeatCtx, task.ID, lease.token, recoveryCfg.LeaseTTL) + defer func() { + stopHeartbeat() + w.releaseGenerationTaskLease(context.Background(), task.ID, lease.token) + }() task.TaskType = strings.TrimSpace(task.TaskType) switch { @@ -215,10 +229,13 @@ func (w *ArticleGenerationWorker) runtimeProcessTimeout() time.Duration { if w != nil && w.processTimeout > 0 { return w.processTimeout } - return articleGenerationProcessTimeout(config.GenerationConfig{}) + cfg := config.GenerationConfig{} + config.NormalizeGenerationConfig(&cfg) + return articleGenerationProcessTimeout(cfg) } func articleGenerationProcessTimeout(cfg config.GenerationConfig) time.Duration { + config.NormalizeGenerationConfig(&cfg) processTimeout := cfg.ArticleTimeout + 2*time.Minute if processTimeout <= 2*time.Minute { return 10 * time.Minute @@ -238,50 +255,11 @@ func (w *ArticleGenerationWorker) handleTaskFailure(ctx context.Context, task te } } -func (w *ArticleGenerationWorker) loadTask(ctx context.Context, taskID int64) (tenantapp.ClaimedGenerationTask, bool, error) { - var ( - task tenantapp.ClaimedGenerationTask - status string - operatorID sql.NullInt64 - articleID sql.NullInt64 - reservationID sql.NullInt64 - inputParamsJSON []byte - ) - - err := w.pool.QueryRow(ctx, ` - SELECT id, tenant_id, operator_id, article_id, quota_reservation_id, task_type, input_params_json, status - FROM generation_tasks - WHERE id = $1 - `, taskID).Scan( - &task.ID, - &task.TenantID, - &operatorID, - &articleID, - &reservationID, - &task.TaskType, - &inputParamsJSON, - &status, - ) - if err != nil { - return tenantapp.ClaimedGenerationTask{}, false, err +func isPermanentGenerationTaskDecodeError(err error) bool { + if err == nil { + return false } - - if status == "completed" || status == "failed" { - return tenantapp.ClaimedGenerationTask{}, true, nil - } - - task.OperatorID = operatorID.Int64 - task.ArticleID = articleID.Int64 - task.QuotaReservationID = reservationID.Int64 - - params, err := parseJSONMap(inputParamsJSON) - if err != nil { - task.InputParams = map[string]interface{}{} - return task, false, fmt.Errorf("decode generation input params: %w", err) - } - task.InputParams = params - - return task, false, nil + return strings.Contains(err.Error(), "decode generation input params") } func (w *ArticleGenerationWorker) rejectDelivery(delivery amqp.Delivery, requeue bool, err error, taskID int64) { diff --git a/server/migrations/20260503110000_add_generation_task_leases.down.sql b/server/migrations/20260503110000_add_generation_task_leases.down.sql new file mode 100644 index 0000000..06b679f --- /dev/null +++ b/server/migrations/20260503110000_add_generation_task_leases.down.sql @@ -0,0 +1,9 @@ +DROP INDEX IF EXISTS idx_generation_tasks_queued_stale; +DROP INDEX IF EXISTS idx_generation_tasks_running_lease; + +ALTER TABLE generation_tasks + DROP COLUMN IF EXISTS last_heartbeat_at, + DROP COLUMN IF EXISTS attempt_count, + DROP COLUMN IF EXISTS lease_expires_at, + DROP COLUMN IF EXISTS lease_owner, + DROP COLUMN IF EXISTS lease_token; diff --git a/server/migrations/20260503110000_add_generation_task_leases.up.sql b/server/migrations/20260503110000_add_generation_task_leases.up.sql new file mode 100644 index 0000000..462e0e2 --- /dev/null +++ b/server/migrations/20260503110000_add_generation_task_leases.up.sql @@ -0,0 +1,14 @@ +ALTER TABLE generation_tasks + ADD COLUMN lease_token VARCHAR(64), + ADD COLUMN lease_owner VARCHAR(128), + ADD COLUMN lease_expires_at TIMESTAMPTZ, + ADD COLUMN attempt_count INT NOT NULL DEFAULT 0, + ADD COLUMN last_heartbeat_at TIMESTAMPTZ; + +CREATE INDEX idx_generation_tasks_running_lease + ON generation_tasks(lease_expires_at, id) + WHERE status = 'running'; + +CREATE INDEX idx_generation_tasks_queued_stale + ON generation_tasks(updated_at, id) + WHERE status = 'queued';