feat(cache,worker): overhaul cache layer and add generation task lease recovery
- shared/cache: add Options/L1/async/metrics/prefix decorators, multi-key ops, Redis pool tuning, and JSON readthrough metrics - worker-generate: claim tasks via DB lease + heartbeat, requeue stale queued tasks, expire dead leases with refund/cache invalidation - tenant: version article cache keys so worker recovery invalidations propagate cleanly - shared/config: expand Redis (pool/timeouts/TLS) and Generation (lease/recovery) configs with defaults Co-Authored-By: Claude Opus 4.7 (1M context) <noreply@anthropic.com>
This commit is contained in:
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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"
|
||||
}
|
||||
|
||||
+178
@@ -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)
|
||||
}
|
||||
+9
@@ -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
|
||||
}
|
||||
|
||||
+31
-3
@@ -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)
|
||||
}
|
||||
|
||||
Vendored
+117
@@ -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:")
|
||||
}
|
||||
+37
@@ -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
|
||||
|
||||
+44
@@ -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)
|
||||
}
|
||||
}
|
||||
+263
@@ -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),
|
||||
)
|
||||
}
|
||||
+159
@@ -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)
|
||||
}
|
||||
+91
@@ -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)
|
||||
}
|
||||
+39
-3
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
+60
-5
@@ -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
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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 == "" {
|
||||
|
||||
@@ -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:
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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))
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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)
|
||||
|
||||
@@ -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)),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
`
|
||||
|
||||
|
||||
@@ -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 {
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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)
|
||||
}
|
||||
}
|
||||
@@ -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) {
|
||||
|
||||
Reference in New Issue
Block a user