dify-plugin-daemon/internal/cluster/plugin.go

package cluster

import (
	"errors"
	"strings"
	"sync/atomic"
	"time"

	"github.com/langgenius/dify-plugin-daemon/pkg/entities/plugin_entities"
	"github.com/langgenius/dify-plugin-daemon/pkg/utils/cache"
	"github.com/langgenius/dify-plugin-daemon/pkg/utils/log"
)

type pluginLifeTime struct {
	lifetime        plugin_entities.PluginLifetime
	lastScheduledAt time.Time
}

type pluginState struct {
	plugin_entities.PluginRuntimeState
	Identity string `json:"identity"`
}

func (l *pluginLifeTime) UpdateScheduledAt(t time.Time) {
	l.lastScheduledAt = t
}

// RegisterPlugin registers a plugin to the cluster, and start to be scheduled
func (c *Cluster) RegisterPlugin(lifetime plugin_entities.PluginLifetime) error {
	identity, err := lifetime.Identity()
	if err != nil {
		return errors.Join(err, errors.New("failed to get plugin identity"))
	}

	if c.showLog {
		log.Info("registering plugin", "identity", identity.String())
	}

	if c.plugins.Exists(identity.String()) {
		// idempotent: plugin already registered, just update its state
		if existing, ok := c.plugins.Load(identity.String()); ok {
			// update the lifetime reference
			existing.lifetime = lifetime
			// update plugin state immediately
			err = c.doPluginStateUpdate(existing)
			if err != nil {
				return errors.Join(err, errors.New("failed to update plugin state"))
			}
		} else {
			log.Warn("plugin does not exist", "identity", identity.String())
		}
		return nil
	}

	l := &pluginLifeTime{
		lifetime: lifetime,
	}

	c.plugins.Store(identity.String(), l)

	// do plugin state update immediately
	err = c.doPluginStateUpdate(l)
	if err != nil {
		return errors.Join(err, errors.New("failed to update plugin state"))
	}

	if c.showLog {
		log.Info("start to schedule plugin", "identity", identity)
	}

	return nil
}

func (c *Cluster) UnregisterPlugin(lifetime plugin_entities.PluginLifetime) error {
	identity, err := lifetime.Identity()
	if err != nil {
		return errors.Join(err, errors.New("failed to get plugin identity"))
	}

	if c.showLog {
		log.Info("unregistering plugin", "identity", identity.String())
	}

	hashedIdentity := plugin_entities.HashedIdentity(identity.String())
	const UNREGISTER_PLUGIN_RETRY_BASE_DELAY = 10 * time.Millisecond

	// remove plugin from cluster with retry logic
	maxRetries := 3
	for i := 0; i < maxRetries; i++ {
		err = c.removePluginState(c.id, hashedIdentity)
		if err == nil {
			break
		}

		// Log retry attempt
		if i < maxRetries-1 {
			log.Warn("failed to remove plugin state, retrying",
				"identity", identity.String(),
				"attempt", i+1,
				"max_retries", maxRetries,
				"error", err,
			)
			time.Sleep(time.Duration(i+1) * UNREGISTER_PLUGIN_RETRY_BASE_DELAY)
		}
	}

	if err != nil {
		return errors.Join(err, errors.New("failed to remove plugin state after retries"))
	}

	c.plugins.Delete(identity.String())

	return nil
}

const (
	PLUGIN_STATE_MAP_KEY = "plugin_state"
)

func (c *Cluster) getPluginStateKey(nodeId string, plugin_id string) string {
	return nodeId + ":" + plugin_id
}

func (c *Cluster) getScanPluginsByNodeKey(nodeId string) string {
	return nodeId + ":*"
}

func (c *Cluster) getScanPluginsByIdKey(plugin_id string) string {
	return "*:" + plugin_id
}

// SchedulePlugin schedules a plugin to the cluster
// it will walk through the plugin state map and update all the states
// as for the plugin has exited, normally, it will be removed automatically
// but once a plugin is not removed, it will be gc by the master node
func (c *Cluster) schedulePlugins() error {
	if c.showLog {
		log.Info("scheduling plugins", "count", c.plugins.Len())
	}

	c.notifyPluginSchedule()
	defer c.notifyPluginScheduleCompleted()

	c.plugins.Range(func(key string, value *pluginLifeTime) bool {
		if time.Since(value.lastScheduledAt) < c.pluginSchedulerInterval {
			return true
		}
		if c.showLog {
			log.Info("scheduling plugin", "key", key)
		}
		// do plugin state update
		err := c.doPluginStateUpdate(value)
		if err != nil {
			log.Error("failed to update plugin state", "error", err)
		}

		if c.showLog {
			log.Info("scheduled plugin", "key", key)
		}

		return true
	})

	if c.showLog {
		log.Info("scheduled plugins", "count", c.plugins.Len())
	}

	return nil
}

// doPluginUpdate updates the plugin state and schedule the plugin
func (c *Cluster) doPluginStateUpdate(lifetime *pluginLifeTime) error {
	state := lifetime.lifetime.RuntimeState()
	identity, err := lifetime.lifetime.Identity()
	if err != nil {
		return err
	}

	hashedIdentity := plugin_entities.HashedIdentity(identity.String())

	scheduleState := &pluginState{
		Identity:           identity.String(),
		PluginRuntimeState: state,
	}

	stateKey := c.getPluginStateKey(c.id, hashedIdentity)

	// check if the plugin has been removed
	if !c.plugins.Exists(identity.String()) {
		if c.showLog {
			log.Info("removing plugin state due no longer exists", "identity", identity.String())
		}
		// remove state
		err = c.removePluginState(c.id, hashedIdentity)
		if err != nil {
			return err
		}
	} else {
		// update plugin state
		scheduleState.ScheduledAt = &[]time.Time{time.Now()}[0]
		err = cache.SetMapOneField(PLUGIN_STATE_MAP_KEY, stateKey, scheduleState)
		if err != nil {
			return err
		}
		lifetime.UpdateScheduledAt(*scheduleState.ScheduledAt)
		if c.showLog {
			log.Info("updated plugin state",
				"identity", identity.String(),
				"state_key", stateKey,
				"scheduled_at", scheduleState.ScheduledAt,
				"status", state.Status,
			)
		}
	}

	lifetime.lastScheduledAt = time.Now()

	return nil
}

func (c *Cluster) removePluginState(nodeId string, hashed_identity string) error {
	if c.showLog {
		log.Info("removing plugin state", "hashed_identity", hashed_identity)
	}
	err := cache.DelMapField(PLUGIN_STATE_MAP_KEY, c.getPluginStateKey(nodeId, hashed_identity))
	if err != nil {
		return err
	}

	if c.showLog {
		log.Info("plugin has been removed from node", "hashed_identity", hashed_identity, "node_id", c.id)
	}

	return nil
}

// forceGCNodePlugins will force garbage collect all the plugins on the node
func (c *Cluster) forceGCNodePlugins(nodeId string) error {
	return cache.ScanMapAsync(
		PLUGIN_STATE_MAP_KEY,
		c.getScanPluginsByNodeKey(nodeId),
		func(m map[string]pluginState) error {
			for _, plugin_state := range m {
				if err := c.forceGCNodePlugin(nodeId, plugin_state.Identity); err != nil {
					return err
				}
			}
			return nil
		},
	)
}

// forceGCNodePlugin will force garbage collect the plugin on the node
func (c *Cluster) forceGCNodePlugin(nodeId string, plugin_id string) error {
	if nodeId == c.id {
		c.pluginLock.Lock()
		c.plugins.Delete(plugin_id)
		c.pluginLock.Unlock()
	}

	if err := c.removePluginState(nodeId, plugin_entities.HashedIdentity(plugin_id)); err != nil {
		return err
	}

	return nil
}

// forceGCPluginByNodePluginJoin will force garbage collect the plugin by node_plugin_join
func (c *Cluster) forceGCPluginByNodePluginJoin(node_plugin_join string) error {
	return cache.DelMapField(PLUGIN_STATE_MAP_KEY, node_plugin_join)
}

func (c *Cluster) isPluginActive(state *pluginState) bool {
	if state == nil {
		return false
	}
	if state.ScheduledAt == nil {
		return false
	}
	if time.Since(*state.ScheduledAt) > c.pluginDeactivatedTimeout {
		return false
	}
	return true
}

// isPluginStateValid checks if a plugin runtime state is still valid
// A state is considered valid if it has been scheduled within the deactivated timeout
func (c *Cluster) isPluginStateValid(state *plugin_entities.PluginRuntimeState) bool {
	if state == nil {
		return false
	}
	if state.ScheduledAt == nil {
		return false
	}
	// Consider state invalid if it hasn't been updated in more than half the deactivated timeout
	// This allows us to clean up stale states faster
	if time.Since(*state.ScheduledAt) > c.pluginDeactivatedTimeout/2 {
		return false
	}
	return true
}

func (c *Cluster) splitNodePluginJoin(node_plugin_join string) (nodeId string, plugin_hashed_id string, err error) {
	split := strings.Split(node_plugin_join, ":")
	if len(split) != 2 {
		return "", "", errors.New("invalid node_plugin_join")
	}

	return split[0], split[1], nil
}

// autoGCPlugins will automatically garbage collect the plugins that are no longer active
func (c *Cluster) autoGCPlugins() error {
	// skip if already in auto gc
	if atomic.LoadInt32(&c.isInAutoGcPlugins) == 1 {
		return nil
	}
	defer atomic.StoreInt32(&c.isInAutoGcPlugins, 0)

	return cache.ScanMapAsync(
		PLUGIN_STATE_MAP_KEY,
		"*",
		func(m map[string]pluginState) error {
			for node_plugin_join, plugin_state := range m {
				if !c.isPluginActive(&plugin_state) {
					nodeId, _, err := c.splitNodePluginJoin(node_plugin_join)
					if err != nil {
						return err
					}

					if plugin_state.ScheduledAt != nil {
						timeSinceScheduled := time.Since(*plugin_state.ScheduledAt)
						log.Warn("auto gc inactive plugin",
							"plugin", plugin_state.Identity,
							"node_id", nodeId,
							"scheduled_at", plugin_state.ScheduledAt,
							"time_since_scheduled", timeSinceScheduled.String(),
							"timeout", c.pluginDeactivatedTimeout.String(),
							"status", plugin_state.Status,
						)
					}

					// force gc the plugin
					if err := c.forceGCNodePlugin(nodeId, plugin_state.Identity); err != nil {
						return err
					}

					// one more time to force gc the plugin, there is a possibility
					// that the hash value of plugin's identity is not the same as the node_plugin_join
					// so we need to force gc the plugin by node_plugin_join again
					if err := c.forceGCPluginByNodePluginJoin(node_plugin_join); err != nil {
						return err
					}
				}
			}
			return nil
		},
	)
}