diff --git a/docs/claude/cache.md b/docs/claude/cache.md index 623f045ef..807ef66aa 100644 --- a/docs/claude/cache.md +++ b/docs/claude/cache.md @@ -70,6 +70,61 @@ if acquired { } ``` +## Owner-Aware Distributed Locks (OwnedLock) + +The `AcquireOwnedLock` function provides token-based lock ownership with advanced features: + +```go +// Acquire lock with timeout +lock, err := cache.AcquireOwnedLock(key, expire, tryLockTimeout) +if err == cache.ErrLockTimeout { + // Lock acquisition timed out + return err +} +if err != nil { + return err +} +defer lock.Unlock() + +// Critical section +``` + +The OwnedLock type supports: + +**Renewal**: Extend lock expiration while holding it +```go +if err := lock.Renew(time.Minute); err == cache.ErrLockNotOwned { + // Lock was lost +} +``` + +**Keep-Alive**: Automatic periodic renewal +```go +ctx, cancel := context.WithCancel(context.Background()) +defer cancel() + +errCh := lock.KeepAlive(ctx, time.Minute) +// Lock is renewed every expire/3 interval +// Cancel context to stop keep-alive +if err := <-errCh; err != nil { + // Keep-alive failed (lock lost or error) +} +``` + +**Safe Unlock**: Only the lock owner can unlock +```go +if err := lock.Unlock(); err == cache.ErrLockNotOwned { + // Lock expired or was never owned +} +``` + +The OwnedLock mechanism uses unique tokens to ensure only the lock owner can renew or release the lock, preventing accidental releases by other processes. + +**Error Constants**: +- `ErrLockTimeout`: Lock acquisition timed out +- `ErrLockNotOwned`: Returned when attempting to renew or unlock a lock that is not owned +- `ErrInvalidLockExpiration`: Returned when lock expiration duration is not positive + ## Auto-type Operations Helper functions in `redis_auto_type.go`: diff --git a/internal/core/control_panel/launcher_local.go b/internal/core/control_panel/launcher_local.go index 128357545..869e7063a 100644 --- a/internal/core/control_panel/launcher_local.go +++ b/internal/core/control_panel/launcher_local.go @@ -76,26 +76,48 @@ func (c *ControlPanel) LaunchLocalPlugin( expire := 15 * time.Minute tryTimeout := 2 * time.Minute log.Info("acquiring distributed init lock", "plugin", pluginUniqueIdentifier.String(), "expire", expire.String()) - if err := cache.Lock(lockKey, expire, tryTimeout); err != nil { - // failed to acquire the lock within timeout - err = errors.Join(err, fmt.Errorf("failed to acquire distributed env-init lock")) - c.WalkNotifiers(func(notifier ControlPanelNotifier) { - notifier.OnLocalRuntimeStartFailed(pluginUniqueIdentifier, err) - }) - // release semaphore and local lock - releaseLockAndSemaphore() - return nil, nil, err + lock, err := cache.AcquireOwnedLock(lockKey, expire, tryTimeout) + if err != nil { + log.Warn( + "failed to acquire distributed init lock; continuing with idempotent initialization", + "plugin", pluginUniqueIdentifier.String(), + "error", err.Error(), + ) } - defer func() { - if unlockErr := cache.Unlock(lockKey); unlockErr != nil { - log.Warn("failed to release distributed init lock", "plugin", pluginUniqueIdentifier.String(), "error", unlockErr.Error()) + + var stopRenew context.CancelFunc + var renewResult <-chan error + if lock != nil { + var renewCtx context.Context + renewCtx, stopRenew = context.WithCancel(context.Background()) + renewResult = lock.KeepAlive(renewCtx, expire) + } + + initErr := runtime.InitEnvironment(decoder) + if lock != nil { + stopRenew() + if renewErr := <-renewResult; renewErr != nil { + log.Warn( + "lost distributed init lock; runtime initialization remains authoritative", + "plugin", pluginUniqueIdentifier.String(), + "error", renewErr.Error(), + ) + } + if unlockErr := lock.Unlock(); unlockErr != nil { + if !errors.Is(unlockErr, cache.ErrLockNotOwned) { + log.Warn( + "failed to release distributed init lock", + "plugin", pluginUniqueIdentifier.String(), + "error", unlockErr.Error(), + ) + } } else { log.Info("released distributed init lock", "plugin", pluginUniqueIdentifier.String()) } - }() + } - if err := runtime.InitEnvironment(decoder); err != nil { - err = errors.Join(err, fmt.Errorf("failed to init environment")) + if initErr != nil { + err = errors.Join(initErr, fmt.Errorf("failed to init environment")) // notify new runtime launch failed c.WalkNotifiers(func(notifier ControlPanelNotifier) { notifier.OnLocalRuntimeStartFailed(pluginUniqueIdentifier, err) diff --git a/internal/core/plugin_manager/installer.go b/internal/core/plugin_manager/installer.go index d49b31e88..993219b94 100644 --- a/internal/core/plugin_manager/installer.go +++ b/internal/core/plugin_manager/installer.go @@ -18,6 +18,7 @@ import ( "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" "github.com/langgenius/dify-plugin-daemon/pkg/utils/routine" "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" + "gorm.io/gorm/clause" ) var ( @@ -27,6 +28,13 @@ var ( func (p *PluginManager) Install( ctx context.Context, pluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier, +) (*stream.Stream[installation_entities.PluginInstallResponse], error) { + return p.EnsureRuntime(ctx, pluginUniqueIdentifier) +} + +func (p *PluginManager) EnsureRuntime( + ctx context.Context, + pluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier, ) (*stream.Stream[installation_entities.PluginInstallResponse], error) { if p.config.Platform == app.PLATFORM_LOCAL { return p.installLocal(ctx, pluginUniqueIdentifier) @@ -170,6 +178,33 @@ func (p *PluginManager) updateServerlessRuntimeModel( return db.Update(&serverlessModel) } +func (p *PluginManager) persistServerlessRuntime( + pluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier, + functionURL string, + functionName string, +) error { + serverlessModel := &models.ServerlessRuntime{ + Checksum: pluginUniqueIdentifier.Checksum(), + Type: models.SERVERLESS_RUNTIME_TYPE_SERVERLESS, + FunctionURL: functionURL, + FunctionName: functionName, + PluginUniqueIdentifier: pluginUniqueIdentifier.String(), + } + if err := db.DifyPluginDB.Clauses(clause.OnConflict{ + Columns: []clause.Column{{Name: "plugin_unique_identifier"}}, + DoUpdates: clause.AssignmentColumns([]string{ + "checksum", + "type", + "function_url", + "function_name", + }), + }).Create(serverlessModel).Error; err != nil { + return err + } + + return p.ClearServerlessRuntimeCache(pluginUniqueIdentifier) +} + // whenever a plugin was installed successfully, a record will be inserted into `models.ServerlessRuntime` func (p *PluginManager) installServerless( ctx context.Context, @@ -209,37 +244,11 @@ func (p *PluginManager) installServerless( return } - // check if the plugin is already installed - // NOTE: models.ServerlessRuntime is a tenant-isolated model - // it hands only engine-level persist data like which serverless runtime is installed - // that's why we placed it here, not in service layer. - // - // service layer takes care of tenant-level persist data like "which tenant installed which plugin" - _, err := db.GetOne[models.ServerlessRuntime]( - db.Equal("plugin_unique_identifier", pluginUniqueIdentifier.String()), - db.Equal("type", string(models.SERVERLESS_RUNTIME_TYPE_SERVERLESS)), - ) - if err == db.ErrDatabaseNotFound { - // create a new serverless runtime - serverlessModel := &models.ServerlessRuntime{ - Checksum: pluginUniqueIdentifier.Checksum(), - Type: models.SERVERLESS_RUNTIME_TYPE_SERVERLESS, - FunctionURL: functionUrl, - FunctionName: functionName, - PluginUniqueIdentifier: pluginUniqueIdentifier.String(), - } - err = db.Create(serverlessModel) - if err != nil { - responseStream.Write(installation_entities.PluginInstallResponse{ - Event: installation_entities.PluginInstallEventError, - Data: "failed to create serverless runtime", - }) - return - } - } else if err != nil { + if err := p.persistServerlessRuntime(pluginUniqueIdentifier, functionUrl, functionName); err != nil { + log.Error("failed to persist serverless runtime", "error", err) responseStream.Write(installation_entities.PluginInstallResponse{ Event: installation_entities.PluginInstallEventError, - Data: "failed to check if the plugin is already installed", + Data: "failed to persist serverless runtime", }) return } @@ -282,21 +291,21 @@ func (p *PluginManager) installLocal( routinepkg.RoutineLabelKeyModule: "plugin_manager", routinepkg.RoutineLabelKeyMethod: "installLocal", }, func() { - // firstly, install the plugin, then launch it, delete it if process fails + // First publish the plugin package, then launch it. A failed attempt must not + // delete the shared package because another caller may already depend on it. var success bool = false + var stopRuntime bool = false var runtime *local_runtime.LocalPluginRuntime var ch <-chan error defer responseStream.Close() + defer p.controlPanel.EnableLocalPluginAutoLaunch(pluginUniqueIdentifier) defer func() { - if !success { - p.controlPanel.RemoveLocalPlugin(pluginUniqueIdentifier) - - // release the lock, avoid a potential race condition - // which causes plugins never to be scheduled automatically - p.controlPanel.EnableLocalPluginAutoLaunch(pluginUniqueIdentifier) - - // forcefully stop runtime, prevent continuous scheduling + if !success && stopRuntime { + // Stop only the runtime constructed by this attempt. The canonical + // package remains available for concurrent callers and retries. A + // timeout does not stop the runtime because it may become ready after + // the caller has left and be observed by another ensure attempt. if runtime != nil { runtime.Stop(false) } @@ -340,8 +349,10 @@ func (p *PluginManager) installLocal( } ticker := time.NewTicker(5 * time.Second) + defer ticker.Stop() timeout := time.Duration(p.config.PythonEnvInitTimeout) * time.Second timer := time.NewTimer(timeout) + defer timer.Stop() for { select { @@ -363,6 +374,7 @@ func (p *PluginManager) installLocal( }) case err := <-ch: if err != nil { + stopRuntime = true responseStream.Write(installation_entities.PluginInstallResponse{ Event: installation_entities.PluginInstallEventError, Data: fmt.Sprintf("failed to launch plugin: %s", err.Error()), diff --git a/internal/core/plugin_manager/local_launch_test.go b/internal/core/plugin_manager/local_launch_test.go index 63be51181..bb6195d75 100644 --- a/internal/core/plugin_manager/local_launch_test.go +++ b/internal/core/plugin_manager/local_launch_test.go @@ -1 +1,87 @@ package plugin_manager + +import ( + "context" + "os" + "path/filepath" + "testing" + + miniredis "github.com/alicebob/miniredis/v2" + cloudoss "github.com/langgenius/dify-cloud-kit/oss" + "github.com/langgenius/dify-cloud-kit/oss/factory" + "github.com/langgenius/dify-plugin-daemon/internal/types/app" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/installation_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/plugin_packager/decoder" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/routine" + "github.com/stretchr/testify/require" +) + +func TestEnsureLocalRuntimeFailureKeepsInstalledPackage(t *testing.T) { + eventData, installed := runFailingLocalRuntimeInstall(t, false) + + require.NotEmpty(t, eventData) + require.True(t, installed) +} + +func TestEnsureLocalRuntimeContinuesWhenRedisLockIsUnavailable(t *testing.T) { + eventData, installed := runFailingLocalRuntimeInstall(t, true) + + require.Contains(t, eventData, "missing-uv") + require.NotContains(t, eventData, "failed to acquire distributed env-init lock") + require.True(t, installed) +} + +func runFailingLocalRuntimeInstall(t *testing.T, closeRedisBeforeInstall bool) (string, bool) { + t.Helper() + routine.InitPool(4) + + redisServer := miniredis.RunT(t) + require.NoError(t, cache.InitRedisClient(redisServer.Addr(), cache.RedisCredentials{}, false, 0, nil)) + t.Cleanup(func() { + _ = cache.Close() + }) + + storageDir := t.TempDir() + storage, err := factory.Load("local", cloudoss.OSSArgs{ + Local: &cloudoss.Local{Path: storageDir}, + }) + require.NoError(t, err) + + manager := InitGlobalManager(storage, &app.Config{ + Platform: app.PLATFORM_LOCAL, + PluginMediaCachePath: "assets", + PluginMediaCacheSize: 4, + PluginAssetCacheSize: 4, + PluginInstalledPath: "installed", + PluginPackageCachePath: "packages", + PluginLocalLaunchingConcurrent: 1, + PluginWorkingPath: filepath.Join(storageDir, "working"), + PythonEnvInitTimeout: 1, + UvPath: filepath.Join(storageDir, "missing-uv"), + }) + + packageBytes, err := os.ReadFile("testdata/openai.difypkg") + require.NoError(t, err) + packageDecoder, err := decoder.NewZipPluginDecoder(packageBytes) + require.NoError(t, err) + identifier, err := packageDecoder.UniqueIdentity() + require.NoError(t, err) + require.NoError(t, manager.packageBucket.Save(identifier.String(), packageBytes)) + if closeRedisBeforeInstall { + require.NoError(t, cache.Close()) + } + + response, err := manager.EnsureRuntime(context.Background(), identifier) + require.NoError(t, err) + eventData := "" + require.NoError(t, response.Process(func(event installation_entities.PluginInstallResponse) { + if event.Event == installation_entities.PluginInstallEventError { + eventData = event.Data + } + })) + + exists, err := manager.installedBucket.Exists(identifier) + require.NoError(t, err) + return eventData, exists +} diff --git a/internal/core/plugin_manager/serverless_install_test.go b/internal/core/plugin_manager/serverless_install_test.go new file mode 100644 index 000000000..b000ef6af --- /dev/null +++ b/internal/core/plugin_manager/serverless_install_test.go @@ -0,0 +1,86 @@ +package plugin_manager + +import ( + "strings" + "sync" + "testing" + "time" + + miniredis "github.com/alicebob/miniredis/v2" + "github.com/langgenius/dify-plugin-daemon/internal/db" + "github.com/langgenius/dify-plugin-daemon/internal/types/models" + "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/stretchr/testify/require" + "gorm.io/driver/sqlite" + "gorm.io/gorm" +) + +func TestPersistServerlessRuntimeConvergesUnderConcurrency(t *testing.T) { + redisServer := miniredis.RunT(t) + require.NoError(t, cache.InitRedisClient(redisServer.Addr(), cache.RedisCredentials{}, false, 0, nil)) + t.Cleanup(func() { + _ = cache.Close() + }) + + dsn := "file:" + strings.ReplaceAll(t.Name(), "/", "_") + "?mode=memory&cache=shared&_busy_timeout=5000" + gormDB, err := gorm.Open(sqlite.Open(dsn), &gorm.Config{}) + require.NoError(t, err) + require.NoError(t, gormDB.AutoMigrate(&models.ServerlessRuntime{})) + sqlDB, err := gormDB.DB() + require.NoError(t, err) + sqlDB.SetMaxOpenConns(1) + db.DifyPluginDB = gormDB + t.Cleanup(func() { + _ = sqlDB.Close() + }) + + identifier := plugin_entities.PluginUniqueIdentifier( + "langgenius/serverless_concurrency:1.0.0@0123456789abcdef0123456789abcdef", + ) + manager := &PluginManager{} + require.NoError(t, cache.Store( + manager.getServerlessRuntimeCacheKey(identifier), + models.ServerlessRuntime{FunctionName: "stale"}, + time.Minute, + )) + + const workers = 8 + start := make(chan struct{}) + var ready sync.WaitGroup + var finished sync.WaitGroup + ready.Add(workers) + finished.Add(workers) + errs := make(chan error, workers) + for range workers { + go func() { + defer finished.Done() + ready.Done() + <-start + errs <- manager.persistServerlessRuntime( + identifier, + "https://runtime.example.test", + "serverless-concurrency", + ) + }() + } + ready.Wait() + close(start) + finished.Wait() + close(errs) + + for err := range errs { + require.NoError(t, err) + } + + runtimes, err := db.GetAll[models.ServerlessRuntime]( + db.Equal("plugin_unique_identifier", identifier.String()), + ) + require.NoError(t, err) + require.Len(t, runtimes, 1) + require.Equal(t, "https://runtime.example.test", runtimes[0].FunctionURL) + require.Equal(t, "serverless-concurrency", runtimes[0].FunctionName) + + _, err = cache.Get[models.ServerlessRuntime](manager.getServerlessRuntimeCacheKey(identifier)) + require.ErrorIs(t, err, cache.ErrNotFound) +} diff --git a/internal/core/serverless_connector/launch.go b/internal/core/serverless_connector/launch.go index 489e55c0d..914fc47de 100644 --- a/internal/core/serverless_connector/launch.go +++ b/internal/core/serverless_connector/launch.go @@ -3,11 +3,14 @@ package serverless import ( "bytes" "context" + "errors" + "sync" "time" "github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities" "github.com/langgenius/dify-plugin-daemon/pkg/plugin_packager/decoder" "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" "github.com/langgenius/dify-plugin-daemon/pkg/utils/stream" ) @@ -31,16 +34,45 @@ func LaunchPlugin( } // check if the plugin has already been initialized - if err := cache.Lock( + lock, err := cache.AcquireOwnedLock( SERVERLESS_LAUNCH_LOCK_PREFIX+checksum, time.Duration(timeout)*time.Second, time.Duration(timeout)*time.Second, - ); err != nil { - return nil, err + ) + if err != nil { + log.Warn( + "failed to acquire serverless launch lock; continuing with idempotent launch", + "checksum", checksum, + "error", err.Error(), + ) + } + + var stopRenew context.CancelFunc + var renewResult <-chan error + if lock != nil { + var renewCtx context.Context + renewCtx, stopRenew = context.WithCancel(context.Background()) + renewResult = lock.KeepAlive(renewCtx, time.Duration(timeout)*time.Second) } + var releaseOnce sync.Once unlock := func(e error) error { - cache.Unlock(SERVERLESS_LAUNCH_LOCK_PREFIX + checksum) + releaseOnce.Do(func() { + if lock == nil { + return + } + stopRenew() + if renewErr := <-renewResult; renewErr != nil { + log.Warn( + "lost serverless launch lock; runtime launch remains authoritative", + "checksum", checksum, + "error", renewErr.Error(), + ) + } + if unlockErr := lock.Unlock(); unlockErr != nil && !errors.Is(unlockErr, cache.ErrLockNotOwned) { + log.Warn("failed to release serverless launch lock", "checksum", checksum, "error", unlockErr.Error()) + } + }) return e } diff --git a/internal/core/serverless_connector/launch_test.go b/internal/core/serverless_connector/launch_test.go new file mode 100644 index 000000000..6b056e197 --- /dev/null +++ b/internal/core/serverless_connector/launch_test.go @@ -0,0 +1,76 @@ +package serverless + +import ( + "context" + "net/http" + "net/http/httptest" + "net/url" + "os" + "testing" + + miniredis "github.com/alicebob/miniredis/v2" + "github.com/langgenius/dify-plugin-daemon/pkg/plugin_packager/decoder" + "github.com/langgenius/dify-plugin-daemon/pkg/utils/cache" + "github.com/stretchr/testify/require" +) + +func TestLaunchPluginContinuesWhenRedisLockIsUnavailable(t *testing.T) { + server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { + require.Equal(t, "/v1/runner/instances", r.URL.Path) + w.Header().Set("Content-Type", "application/json") + _, err := w.Write([]byte(`{ + "error":"", + "Items":[{ + "ID":"runtime-id", + "Name":"runtime-name", + "Endpoint":"https://runtime.example.test", + "ResourceName":"runtime-resource", + "Status":{"State":"running"} + }] + }`)) + require.NoError(t, err) + })) + t.Cleanup(server.Close) + + previousBaseURL := baseurl + previousClient := client + parsedURL, err := url.Parse(server.URL) + require.NoError(t, err) + baseurl = parsedURL + client = server.Client() + t.Cleanup(func() { + baseurl = previousBaseURL + client = previousClient + }) + + redisServer := miniredis.RunT(t) + require.NoError(t, cache.InitRedisClient(redisServer.Addr(), cache.RedisCredentials{}, false, 0, nil)) + require.NoError(t, cache.Close()) + + packageBytes, err := os.ReadFile("../plugin_manager/testdata/openai.difypkg") + require.NoError(t, err) + packageDecoder, err := decoder.NewZipPluginDecoder(packageBytes) + require.NoError(t, err) + identifier, err := packageDecoder.UniqueIdentity() + require.NoError(t, err) + + response, err := LaunchPlugin( + context.Background(), + identifier, + packageBytes, + packageDecoder, + 1, + false, + ) + require.NoError(t, err) + + events := make([]LaunchFunctionResponse, 0, 3) + require.NoError(t, response.Process(func(event LaunchFunctionResponse) { + events = append(events, event) + })) + require.Equal(t, []LaunchFunctionResponse{ + {Event: FunctionUrl, Message: "https://runtime.example.test"}, + {Event: Function, Message: "runtime-name"}, + {Event: Done, Message: ""}, + }, events) +} diff --git a/internal/service/install_plugin.go b/internal/service/install_plugin.go index e4df027a6..005bb0d01 100644 --- a/internal/service/install_plugin.go +++ b/internal/service/install_plugin.go @@ -42,6 +42,27 @@ func InstallMultiplePluginsToTenant( metas []map[string]any, ) *entities.Response { runtimeType := config.Platform.ToPluginRuntimeType() + needsRuntimeInstall := make([]bool, len(pluginUniqueIdentifiers)) + allInstalled := true + + for i, pluginUniqueIdentifier := range pluginUniqueIdentifiers { + installed, err := isPluginInstalledForTenant(tenantId, pluginUniqueIdentifier) + if err != nil { + return exception.InternalServerError(err).ToResponse() + } + needsRuntimeInstall[i] = !installed + if !installed { + allInstalled = false + } + } + + if allInstalled { + return entities.NewSuccessResponse(&InstallPluginResponse{ + AllInstalled: true, + TaskID: "", + }) + } + manager := plugin_manager.Manager() if manager == nil { return exception.InternalServerError(errors.New("plugin manager is not initialized")).ToResponse() @@ -51,7 +72,6 @@ func InstallMultiplePluginsToTenant( // and runs in a single goroutine after the task is created jobs := make([]tasks.PluginInstallJob, 0, len(pluginUniqueIdentifiers)) declarations := make([]*plugin_entities.PluginDeclaration, 0, len(pluginUniqueIdentifiers)) - allInstalled := true for i, pluginUniqueIdentifier := range pluginUniqueIdentifiers { declaration, err := helper.CombinedGetPluginDeclaration( @@ -62,23 +82,11 @@ func InstallMultiplePluginsToTenant( return exception.InternalServerError(errors.Join(err, errors.New("failed to get plugin declaration"))).ToResponse() } - _, err = db.GetOne[models.Plugin]( - db.Equal("plugin_unique_identifier", pluginUniqueIdentifier.String()), - ) - - needsRuntimeInstall := false - if err == db.ErrDatabaseNotFound { - needsRuntimeInstall = true - allInstalled = false - } else if err != nil { - return exception.InternalServerError(err).ToResponse() - } - job := tasks.PluginInstallJob{ Identifier: pluginUniqueIdentifier, Declaration: declaration, Meta: metas[i], - NeedsRuntimeInstall: needsRuntimeInstall, + NeedsRuntimeInstall: needsRuntimeInstall[i], } jobs = append(jobs, job) @@ -87,26 +95,6 @@ func InstallMultiplePluginsToTenant( tenants := []string{tenantId} - // all plugins are installed, no need to create tasks - // just add DB record and return - if allInstalled { - for i := range jobs { - if err := tasks.SaveInstallationForTenantsToDB( - tenants, - jobs[i], - runtimeType, - source, - ); err != nil { - return exception.InternalServerError(errors.Join(err, errors.New("failed on plugin installation"))).ToResponse() - } - } - - return entities.NewSuccessResponse(&InstallPluginResponse{ - AllInstalled: true, - TaskID: "", - }) - } - // create tasks for each plugin statuses := buildTaskStatuses(pluginUniqueIdentifiers, declarations, source) taskRegistry, err := createInstallTasks(tenants, statuses) @@ -145,6 +133,23 @@ func InstallMultiplePluginsToTenant( }) } +func isPluginInstalledForTenant( + tenantID string, + pluginUniqueIdentifier plugin_entities.PluginUniqueIdentifier, +) (bool, error) { + _, err := db.GetOne[models.PluginInstallation]( + db.Equal("tenant_id", tenantID), + db.Equal("plugin_id", pluginUniqueIdentifier.PluginID()), + ) + if err == nil { + return true, nil + } + if errors.Is(err, db.ErrDatabaseNotFound) { + return false, nil + } + return false, err +} + /* * Reinstall a plugin from a given identifier, no tenant_id is needed */ diff --git a/internal/service/plugin_decoder_test.go b/internal/service/plugin_decoder_test.go index 867de43b1..816551832 100644 --- a/internal/service/plugin_decoder_test.go +++ b/internal/service/plugin_decoder_test.go @@ -60,6 +60,7 @@ func setupUploadTestEnv(t *testing.T, forceVerify bool) uploadTestEnv { &models.Plugin{}, &models.PluginDeclaration{}, &models.PluginInstallation{}, + &models.InstallTask{}, )) db.DifyPluginDB = gormDB t.Cleanup(func() { @@ -158,6 +159,12 @@ func TestUploadPluginPkgPersistsValidPackageAndSupportsInstallLookup(t *testing. PluginID: identifier.PluginID(), InstallType: plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL, })) + require.NoError(t, db.Create(&models.PluginInstallation{ + TenantID: "00000000-0000-0000-0000-000000000001", + PluginID: identifier.PluginID(), + PluginUniqueIdentifier: identifier.String(), + RuntimeType: string(plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL), + })) installResp := InstallMultiplePluginsToTenant( context.Background(), env.config, @@ -172,6 +179,39 @@ func TestUploadPluginPkgPersistsValidPackageAndSupportsInstallLookup(t *testing. require.Empty(t, installData.TaskID) } +func TestTenantInstallStateDoesNotReuseGlobalPluginState(t *testing.T) { + setupUploadTestEnv(t, false) + pkgBytes := unsignedPluginPackage(t, "tenantinstallstate") + identifier := packageIdentifier(t, pkgBytes) + + require.NoError(t, db.Create(&models.Plugin{ + PluginUniqueIdentifier: identifier.String(), + PluginID: identifier.PluginID(), + InstallType: plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL, + })) + + installed, err := isPluginInstalledForTenant( + "00000000-0000-0000-0000-000000000001", + identifier, + ) + require.NoError(t, err) + require.False(t, installed) + + require.NoError(t, db.Create(&models.PluginInstallation{ + TenantID: "00000000-0000-0000-0000-000000000001", + PluginID: identifier.PluginID(), + PluginUniqueIdentifier: identifier.String(), + RuntimeType: string(plugin_entities.PLUGIN_RUNTIME_TYPE_LOCAL), + })) + + installed, err = isPluginInstalledForTenant( + "00000000-0000-0000-0000-000000000001", + identifier, + ) + require.NoError(t, err) + require.True(t, installed) +} + func TestUploadPluginPkgStillAllowsUnsignedPackageWhenVerificationNotRequired(t *testing.T) { env := setupUploadTestEnv(t, false) pkgBytes := unsignedPluginPackage(t, "unsignedallowed") diff --git a/internal/tasks/install_plugin.go b/internal/tasks/install_plugin.go index d25575c90..e4f1ef32f 100644 --- a/internal/tasks/install_plugin.go +++ b/internal/tasks/install_plugin.go @@ -67,7 +67,7 @@ func ProcessInstallJob( SetTaskStatusForOnePlugin(taskIDs, job.Identifier, models.InstallTaskStatusRunning, "starting") // start installation process - installationStream, err := manager.Install(ctx, job.Identifier) + installationStream, err := manager.EnsureRuntime(ctx, job.Identifier) if err != nil { status = "failed" SetTaskStatusForOnePlugin(taskIDs, job.Identifier, models.InstallTaskStatusFailed, fmt.Sprintf("failed to start installation: %v", err)) @@ -121,7 +121,7 @@ func ProcessUpgradeJob( SetTaskStatusForOnePlugin(taskIDs, job.NewIdentifier, models.InstallTaskStatusRunning, "starting") // start installation process - installationStream, err := manager.Install(ctx, job.NewIdentifier) + installationStream, err := manager.EnsureRuntime(ctx, job.NewIdentifier) if err != nil { SetTaskStatusForOnePlugin(taskIDs, job.NewIdentifier, models.InstallTaskStatusFailed, fmt.Sprintf("failed to start installation: %v", err)) return diff --git a/pkg/utils/cache/redis.go b/pkg/utils/cache/redis.go index dcb21812a..85fbc44c0 100644 --- a/pkg/utils/cache/redis.go +++ b/pkg/utils/cache/redis.go @@ -8,6 +8,7 @@ import ( "sync" "time" + "github.com/google/uuid" "github.com/langgenius/dify-plugin-daemon/pkg/utils/log" "github.com/langgenius/dify-plugin-daemon/pkg/utils/parser" "github.com/redis/go-redis/extra/redisotel/v9" @@ -523,9 +524,145 @@ func SetNX[T any](key string, value T, expire time.Duration, context ...redis.Cm } var ( - ErrLockTimeout = errors.New("lock timeout") + ErrLockTimeout = errors.New("lock timeout") + ErrLockNotOwned = errors.New("lock is not owned") + ErrInvalidLockExpiration = errors.New("lock expiration must be positive") ) +var unlockOwnedLockScript = redis.NewScript(` +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("del", KEYS[1]) +end +return 0 +`) + +var renewOwnedLockScript = redis.NewScript(` +if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("pexpire", KEYS[1], ARGV[2]) +end +return 0 +`) + +type OwnedLock struct { + key string + token string + cmdable redis.Cmdable +} + +func AcquireOwnedLock( + key string, + expire time.Duration, + tryLockTimeout time.Duration, + commands ...redis.Cmdable, +) (*OwnedLock, error) { + if client == nil { + return nil, ErrDBNotInit + } + if expire <= 0 { + return nil, ErrInvalidLockExpiration + } + + cmdable := getCmdable(commands...) + token := uuid.NewString() + deadline := time.NewTimer(tryLockTimeout) + defer deadline.Stop() + + const retryInterval = 20 * time.Millisecond + ticker := time.NewTicker(retryInterval) + defer ticker.Stop() + + for { + acquired, err := cmdable.SetNX(ctx, serialKey(key), token, expire).Result() + if err != nil { + return nil, err + } + if acquired { + return &OwnedLock{ + key: key, + token: token, + cmdable: cmdable, + }, nil + } + + if tryLockTimeout <= 0 { + return nil, ErrLockTimeout + } + + select { + case <-deadline.C: + return nil, ErrLockTimeout + case <-ticker.C: + } + } +} + +func (l *OwnedLock) Renew(expire time.Duration) error { + if expire <= 0 { + return ErrInvalidLockExpiration + } + + renewed, err := renewOwnedLockScript.Run( + ctx, + l.cmdable, + []string{serialKey(l.key)}, + l.token, + expire.Milliseconds(), + ).Int64() + if err != nil { + return err + } + if renewed == 0 { + return ErrLockNotOwned + } + return nil +} + +func (l *OwnedLock) KeepAlive(keepAliveCtx context.Context, expire time.Duration) <-chan error { + result := make(chan error, 1) + if expire <= 0 { + result <- ErrInvalidLockExpiration + close(result) + return result + } + go func() { + defer close(result) + + interval := expire / 3 + ticker := time.NewTicker(interval) + defer ticker.Stop() + + for { + select { + case <-keepAliveCtx.Done(): + result <- nil + return + case <-ticker.C: + if err := l.Renew(expire); err != nil { + result <- err + return + } + } + } + }() + return result +} + +func (l *OwnedLock) Unlock() error { + released, err := unlockOwnedLockScript.Run( + ctx, + l.cmdable, + []string{serialKey(l.key)}, + l.token, + ).Int64() + if err != nil { + return err + } + if released == 0 { + return ErrLockNotOwned + } + return nil +} + var ( distributedLocks = sync.Map{} ) diff --git a/pkg/utils/cache/redis_owned_lock_test.go b/pkg/utils/cache/redis_owned_lock_test.go new file mode 100644 index 000000000..21dcd236a --- /dev/null +++ b/pkg/utils/cache/redis_owned_lock_test.go @@ -0,0 +1,71 @@ +package cache + +import ( + "errors" + "testing" + "time" + + miniredis "github.com/alicebob/miniredis/v2" + "github.com/stretchr/testify/require" +) + +func setupOwnedLockTest(t *testing.T) *miniredis.Miniredis { + t.Helper() + + server := miniredis.RunT(t) + require.NoError(t, InitRedisClient(server.Addr(), RedisCredentials{}, false, 0, nil)) + t.Cleanup(func() { + _ = Close() + }) + return server +} + +func TestOwnedLockRejectsStaleOwnerUnlock(t *testing.T) { + server := setupOwnedLockTest(t) + + first, err := AcquireOwnedLock("owned-lock", time.Second, time.Second) + require.NoError(t, err) + + server.FastForward(time.Second + time.Millisecond) + + second, err := AcquireOwnedLock("owned-lock", time.Second, time.Second) + require.NoError(t, err) + require.ErrorIs(t, first.Unlock(), ErrLockNotOwned) + + value, err := server.Get(serialKey("owned-lock")) + require.NoError(t, err) + require.Equal(t, second.token, value) + require.NoError(t, second.Unlock()) +} + +func TestOwnedLockOnlyRenewsCurrentOwner(t *testing.T) { + server := setupOwnedLockTest(t) + + first, err := AcquireOwnedLock("renewable-lock", time.Second, time.Second) + require.NoError(t, err) + server.FastForward(time.Second + time.Millisecond) + + second, err := AcquireOwnedLock("renewable-lock", time.Second, time.Second) + require.NoError(t, err) + require.ErrorIs(t, first.Renew(2*time.Second), ErrLockNotOwned) + require.NoError(t, second.Renew(2*time.Second)) + + server.FastForward(time.Second + time.Millisecond) + value, err := server.Get(serialKey("renewable-lock")) + require.NoError(t, err) + require.Equal(t, second.token, value) + require.NoError(t, second.Unlock()) +} + +func TestOwnedLockValidatesExpiration(t *testing.T) { + setupOwnedLockTest(t) + + _, err := AcquireOwnedLock("invalid-lock", 0, time.Second) + require.ErrorIs(t, err, ErrInvalidLockExpiration) + + lock, err := AcquireOwnedLock("valid-lock", time.Second, time.Second) + require.NoError(t, err) + require.True(t, errors.Is(lock.Renew(0), ErrInvalidLockExpiration)) + require.ErrorIs(t, <-lock.KeepAlive(t.Context(), 0), ErrInvalidLockExpiration) + require.NoError(t, lock.Unlock()) +}