diff --git a/internal/apps/admin/system_config/routers.go b/internal/apps/admin/system_config/routers.go index 8fc96b05..7f6770b1 100644 --- a/internal/apps/admin/system_config/routers.go +++ b/internal/apps/admin/system_config/routers.go @@ -248,6 +248,7 @@ func UpdateSystemConfig(c *gin.Context) { if key == model.ConfigKeyStorageConfig { storage.ResetCache() + storage.PublishCacheInvalidation(c.Request.Context()) } c.JSON(http.StatusOK, util.OKNil()) diff --git a/internal/apps/upload/storage_migration_task.go b/internal/apps/upload/storage_migration_task.go index edb787c5..2b1ba6f3 100644 --- a/internal/apps/upload/storage_migration_task.go +++ b/internal/apps/upload/storage_migration_task.go @@ -14,6 +14,7 @@ import ( "os" "strings" "sync/atomic" + "time" "github.com/Rain-kl/Wavelet/internal/db" "github.com/Rain-kl/Wavelet/internal/model" @@ -67,6 +68,18 @@ func (h *MigrationHandler) ValidatePayload(payload []byte) ([]byte, error) { // Execute migrates all unique active-storage objects to the pending backend. func (h *MigrationHandler) Execute(ctx context.Context, payload []byte) (*task.TaskResult, error) { + if db.Redis != nil { + lockKey := db.PrefixedKey("lock:storage:migrate") + ok, err := db.Redis.SetNX(ctx, lockKey, "locked", time.Hour).Result() + if err != nil { + return nil, fmt.Errorf("acquire migration lock: %w", err) + } + if !ok { + return nil, errors.New("另一个存储迁移任务正在运行中") + } + defer db.Redis.Del(ctx, lockKey) + } + active, err := storage.LoadConfig(ctx) if err != nil { return nil, fmt.Errorf("load active storage config: %w", err) diff --git a/internal/apps/upload/storage_migration_task_test.go b/internal/apps/upload/storage_migration_task_test.go index 2f6a428b..86d4177f 100644 --- a/internal/apps/upload/storage_migration_task_test.go +++ b/internal/apps/upload/storage_migration_task_test.go @@ -14,10 +14,14 @@ import ( "path/filepath" "strings" "testing" + "time" + "github.com/Rain-kl/Wavelet/internal/db" "github.com/Rain-kl/Wavelet/internal/model" "github.com/Rain-kl/Wavelet/internal/storage" "github.com/Rain-kl/Wavelet/internal/testhelper" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" ) func TestMigrationHandlerExecute(t *testing.T) { @@ -220,3 +224,62 @@ func TestMigrationHandlerExecuteWithHashValidation(t *testing.T) { t.Errorf("StorageDriver = %q, want %q", migrated.StorageDriver, storage.DriverS3) } } + +func TestMigrationHandlerExecuteWithRedisLock(t *testing.T) { + _, _, cleanup := testhelper.SetupTestEnvironment(t) + defer cleanup() + + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("Failed to run miniredis: %v", err) + } + defer mr.Close() + + rdb := redis.NewClient(&redis.Options{ + Addr: mr.Addr(), + }) + defer rdb.Close() + + oldRedis := db.Redis + db.Redis = rdb + defer func() { + db.Redis = oldRedis + }() + + ctx := context.Background() + + // Acquire lock manually + lockKey := db.PrefixedKey("lock:storage:migrate") + if err := rdb.Set(ctx, lockKey, "locked", time.Hour).Err(); err != nil { + t.Fatalf("Failed to set manual lock in Redis: %v", err) + } + + active := storage.DefaultConfig() + if err := storage.SaveActiveConfig(ctx, active); err != nil { + t.Fatalf("SaveActiveConfig() returned error: %v", err) + } + + payload, err := json.Marshal(storageMigrationPayload{Target: active}) + if err != nil { + t.Fatalf("Marshal payload failed: %v", err) + } + + // Execution should fail because lock is already acquired + _, err = (&MigrationHandler{}).Execute(ctx, payload) + if err == nil { + t.Fatal("Execute() succeeded when lock was held, want error") + } + if !strings.Contains(err.Error(), "另一个存储迁移任务正在运行中") { + t.Errorf("expected lock warning, got: %v", err) + } + + // Release lock and run again, should succeed + if err := rdb.Del(ctx, lockKey).Err(); err != nil { + t.Fatalf("Failed to delete lock: %v", err) + } + + _, err = (&MigrationHandler{}).Execute(ctx, payload) + if err != nil { + t.Fatalf("Execute() failed after lock released: %v", err) + } +} diff --git a/internal/storage/config.go b/internal/storage/config.go index 6dfcd1a8..039d15ef 100644 --- a/internal/storage/config.go +++ b/internal/storage/config.go @@ -88,6 +88,8 @@ func DefaultConfig() Config { // LoadConfig loads the active storage configuration. func LoadConfig(ctx context.Context) (Config, error) { + pubSubOnce.Do(startPubSubListener) + cacheMutex.RLock() isCacheValid := time.Since(lastChecked) < 5*time.Second && activeConfigJSON != "" configJSON := activeConfigJSON @@ -184,6 +186,7 @@ func saveSystemConfig(ctx context.Context, key string, value any, description st }) if err == nil && key == model.ConfigKeyStorageConfig { ResetCache() + PublishCacheInvalidation(ctx) } return err } diff --git a/internal/storage/storage.go b/internal/storage/storage.go index 05d44cfc..7965acb2 100644 --- a/internal/storage/storage.go +++ b/internal/storage/storage.go @@ -13,6 +13,7 @@ import ( "sync" "time" + "github.com/Rain-kl/Wavelet/internal/db" "github.com/Rain-kl/Wavelet/internal/model" "gorm.io/gorm" ) @@ -51,6 +52,10 @@ var ( cacheMutex sync.RWMutex ) +const configInvalidationChannel = "storage:config_invalidation" + +var pubSubOnce sync.Once + // ResetCache clears the local cache for storage configuration and client singletons. func ResetCache() { cacheMutex.Lock() @@ -61,12 +66,39 @@ func ResetCache() { lastChecked = time.Time{} } +// PublishCacheInvalidation broadcasts cache eviction to all nodes in the cluster via Redis. +func PublishCacheInvalidation(ctx context.Context) { + if db.Redis != nil { + _ = db.Redis.Publish(ctx, configInvalidationChannel, "reset").Err() + } +} + +// startPubSubListener starts the background subscriber for cache invalidations. +func startPubSubListener() { + if db.Redis == nil { + return + } + go func() { + pubsub := db.Redis.Subscribe(context.Background(), configInvalidationChannel) + defer func() { + _ = pubsub.Close() + }() + + ch := pubsub.Channel() + for range ch { + ResetCache() + } + }() +} + // Active returns the configured active driver and backend, using an in-memory cache with 5s TTL. func Active(ctx context.Context) (Driver, Backend, error) { if IsEnabledFunc() && mockBackend != nil { return DriverS3, mockBackend, nil } + pubSubOnce.Do(startPubSubListener) + cacheMutex.RLock() isCacheValid := time.Since(lastChecked) < 5*time.Second && activeBackend != nil if isCacheValid { diff --git a/internal/storage/storage_test.go b/internal/storage/storage_test.go index 760a27e9..534fd67c 100644 --- a/internal/storage/storage_test.go +++ b/internal/storage/storage_test.go @@ -4,11 +4,16 @@ package storage import ( - "context" - "encoding/json" - "io" - "testing" - "time" + "context" + "encoding/json" + "io" + "sync" + "testing" + "time" + + "github.com/Rain-kl/Wavelet/internal/db" + "github.com/alicebob/miniredis/v2" + "github.com/redis/go-redis/v9" ) func TestStorageCache(t *testing.T) { @@ -80,3 +85,68 @@ func TestStorageCache(t *testing.T) { t.Fatal("ResetCache did not clear cache variables after setting them") } } + +func TestStorageCachePubSub(t *testing.T) { + // 1. Start miniredis + mr, err := miniredis.Run() + if err != nil { + t.Fatalf("Failed to run miniredis: %v", err) + } + defer mr.Close() + + // 2. Initialize Redis client + rdb := redis.NewClient(&redis.Options{ + Addr: mr.Addr(), + }) + defer rdb.Close() + + // 3. Set db.Redis to our client + oldRedis := db.Redis + db.Redis = rdb + defer func() { + db.Redis = oldRedis + }() + + // Reset cache and set some cached config + ResetCache() + cacheMutex.Lock() + activeConfigJSON = "some_config" + lastChecked = time.Now() + cacheMutex.Unlock() + + // 4. Force trigger lazy initialization of subscription + // Reset the once guard so it runs the listener + pubSubOnce = sync.Once{} + ctx := context.Background() + + // Create mock backend for Active call + mockBnd := &functionBackend{ + put: func(context.Context, string, io.Reader, int64, string) error { return nil }, + get: func(context.Context, string) (*Object, error) { return nil, nil }, + delete: func(context.Context, string) error { return nil }, + } + cacheMutex.Lock() + activeBackend = mockBnd + activeDriver = DriverLocal + cacheMutex.Unlock() + + _, _, _ = Active(ctx) // This calls startPubSubListener() + + // Allow some time for subscriber connection + time.Sleep(100 * time.Millisecond) + + // 5. Publish cache invalidation + PublishCacheInvalidation(ctx) + + // Allow message propagation + time.Sleep(100 * time.Millisecond) + + // 6. Verify cache was cleared + cacheMutex.RLock() + configJSON := activeConfigJSON + cacheMutex.RUnlock() + + if configJSON != "" { + t.Error("Memory cache was not cleared after Redis Pub/Sub broadcast") + } +}