refactor

cmd/server/main.go

@@ -28,7 +28,7 @@ func main() {
 		log.Panic("Invalid configuration: %s", err.Error())
 	}
 
-	server.Run(&config)
+	(&server.App{}).Run(&config)
 }
 
 func setDefault(config *app.Config) {

internal/cluster/init.go

@@ -1,25 +1,25 @@
 package cluster
 
-import "github.com/langgenius/dify-plugin-daemon/internal/types/app"
+import (
+	"sync"
+
+	"github.com/langgenius/dify-plugin-daemon/internal/types/app"
+)
 
 type Cluster struct {
 	port uint16
-}
 
-var (
-	cluster *Cluster
-)
+	plugins     map[string]*PluginLifeTime
+	plugin_lock sync.Mutex
+}
 
-func Launch(config *app.Config) {
-	cluster = &Cluster{
-		port: uint16(config.ServerPort),
+func NewCluster(config *app.Config) *Cluster {
+	return &Cluster{
+		port:    uint16(config.ServerPort),
+		plugins: make(map[string]*PluginLifeTime),
 	}
-
-	go func() {
-		cluster.clusterLifetime()
-	}()
 }
 
-func GetCluster() *Cluster {
-	return cluster
+func (c *Cluster) Launch(config *app.Config) {
+	go c.clusterLifetime()
 }

internal/cluster/state.go

@@ -1 +1,40 @@
 package cluster
+
+import (
+	"github.com/langgenius/dify-plugin-daemon/internal/types/entities"
+	"github.com/langgenius/dify-plugin-daemon/internal/utils/log"
+)
+
+type PluginLifeTime struct {
+	lifetime entities.PluginRuntimeTimeLifeInterface
+}
+
+// RegisterPlugin registers a plugin to the cluster, and start to be scheduled
+func (c *Cluster) RegisterPlugin(lifetime entities.PluginRuntimeTimeLifeInterface) error {
+	identity, err := lifetime.Identity()
+	if err != nil {
+		return err
+	}
+
+	lifetime.OnStop(func() {
+		c.plugin_lock.Lock()
+		delete(c.plugins, identity)
+		c.plugin_lock.Unlock()
+	})
+
+	c.plugin_lock.Lock()
+	if !lifetime.Stopped() {
+		c.plugins[identity] = &PluginLifeTime{
+			lifetime: lifetime,
+		}
+	}
+	c.plugin_lock.Unlock()
+
+	log.Info("start to schedule plugin %s", identity)
+
+	return nil
+}
+
+func (c *Cluster) SchedulePlugin(lifetime entities.PluginRuntimeTimeLifeInterface) error {
+	return nil
+}

internal/core/plugin_daemon/model_service.go

@@ -20,10 +20,8 @@ func genericInvokePlugin[Req any, Rsp any](
 	response_buffer_size int,
 	typ backwards_invocation.PluginAccessType,
 	action backwards_invocation.PluginAccessAction,
-) (
-	*stream.StreamResponse[Rsp], error,
-) {
-	runtime := plugin_manager.Get(session.PluginIdentity())
+) (*stream.StreamResponse[Rsp], error) {
+	runtime := plugin_manager.GetGlobalPluginManager().Get(session.PluginIdentity())
 	if runtime == nil {
 		return nil, errors.New("plugin not found")
 	}

internal/core/plugin_manager/init.go

@@ -6,7 +6,9 @@ import (
 	"github.com/langgenius/dify-plugin-daemon/internal/types/entities"
 )
 
-var m sync.Map
+var (
+	m sync.Map
+)
 
 func checkPluginExist(identity string) (entities.PluginRuntimeInterface, bool) {
 	if v, ok := m.Load(identity); ok {

internal/core/plugin_manager/lifetime.go

@@ -8,22 +8,22 @@ import (
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/log"
 )
 
-func lifetime(config *app.Config, r entities.PluginRuntimeInterface) {
+func (p *PluginManager) lifetime(config *app.Config, r entities.PluginRuntimeInterface) {
 	start_failed_times := 0
 	configuration := r.Configuration()
 
 	log.Info("new plugin logged in: %s", configuration.Identity())
 	defer log.Info("plugin %s has exited", configuration.Identity())
 
-	// store plugin runtime
-	m.Store(configuration.Identity(), r)
-	defer m.Delete(configuration.Identity())
-
-	// update lifetime state for this pod
-	addLifetimeState(r)
+	// add plugin to cluster
+	err := p.cluster.RegisterPlugin(r)
+	if err != nil {
+		log.Error("add plugin to cluster failed: %s", err.Error())
+		return
+	}
 
 	// remove lifetime state after plugin if it has been stopped
-	defer deleteLifetimeState(r)
+	defer r.TriggerStop()
 
 	for !r.Stopped() {
 		if err := r.InitEnvironment(); err != nil {

internal/core/plugin_manager/manager.go

@@ -3,14 +3,33 @@ package plugin_manager
 import (
 	"fmt"
 
+	"github.com/langgenius/dify-plugin-daemon/internal/cluster"
 	"github.com/langgenius/dify-plugin-daemon/internal/core/dify_invocation"
 	"github.com/langgenius/dify-plugin-daemon/internal/types/app"
 	"github.com/langgenius/dify-plugin-daemon/internal/types/entities"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/cache"
 	"github.com/langgenius/dify-plugin-daemon/internal/utils/log"
 )
 
-func List() []entities.PluginRuntimeInterface {
+type PluginManager struct {
+	cluster *cluster.Cluster
+}
+
+var (
+	manager *PluginManager
+)
+
+func InitGlobalPluginManager(cluster *cluster.Cluster) {
+	manager = &PluginManager{
+		cluster: cluster,
+	}
+}
+
+func GetGlobalPluginManager() *PluginManager {
+	return manager
+}
+
+func (p *PluginManager) List() []entities.PluginRuntimeInterface {
 	var runtimes []entities.PluginRuntimeInterface
 	m.Range(func(key, value interface{}) bool {
 		if v, ok := value.(entities.PluginRuntimeInterface); ok {
@@ -21,7 +40,7 @@ func List() []entities.PluginRuntimeInterface {
 	return runtimes
 }
 
-func Get(identity string) entities.PluginRuntimeInterface {
+func (p *PluginManager) Get(identity string) entities.PluginRuntimeInterface {
 	if v, ok := m.Load(identity); ok {
 		if r, ok := v.(entities.PluginRuntimeInterface); ok {
 			return r
@@ -30,15 +49,15 @@ func Get(identity string) entities.PluginRuntimeInterface {
 	return nil
 }
 
-func Put(path string, binary []byte) {
+func (p *PluginManager) Put(path string, binary []byte) {
 	//TODO: put binary into
 }
 
-func Delete(identity string) {
+func (p *PluginManager) Delete(identity string) {
 	//TODO: delete binary from
 }
 
-func Init(configuration *app.Config) {
+func (p *PluginManager) Init(configuration *app.Config) {
 	// TODO: init plugin manager
 	log.Info("start plugin manager daemon...")
 
@@ -57,8 +76,5 @@ func Init(configuration *app.Config) {
 	}
 
 	// start plugin watcher
-	startWatcher(configuration)
-
-	// start plugin lifetime manager
-	startLifeTimeManager(configuration)
+	p.startWatcher(configuration)
 }