stz plugin

Author: Sayan-

server/internal/plugins/manager.go

@@ -13,6 +13,7 @@ import (
 	"github.com/m1k1o/neko/server/internal/config"
 	"github.com/m1k1o/neko/server/internal/plugins/chat"
 	"github.com/m1k1o/neko/server/internal/plugins/filetransfer"
+	"github.com/m1k1o/neko/server/internal/plugins/scaletozero"
 	"github.com/m1k1o/neko/server/pkg/types"
 )
 
@@ -47,6 +48,7 @@ func New(config *config.Plugins) *ManagerCtx {
 	// add built-in plugins
 	manager.plugins.addPlugin(filetransfer.NewPlugin())
 	manager.plugins.addPlugin(chat.NewPlugin())
+	manager.plugins.addPlugin(scaletozero.NewPlugin())
 
 	return manager
 }

server/internal/plugins/scaletozero/config.go

@@ -0,0 +1,23 @@
+package scaletozero
+
+import (
+	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
+)
+
+type Config struct {
+	Enabled bool
+}
+
+func (Config) Init(cmd *cobra.Command) error {
+	cmd.PersistentFlags().Bool("scaletozero.enabled", false, "enable scale-to-zero")
+	if err := viper.BindPFlag("scaletozero.enabled", cmd.PersistentFlags().Lookup("scaletozero.enabled")); err != nil {
+		return err
+	}
+
+	return nil
+}
+
+func (c *Config) Set() {
+	c.Enabled = viper.GetBool("scaletozero.enabled")
+}

server/internal/plugins/scaletozero/manager.go

@@ -0,0 +1,78 @@
+package scaletozero
+
+import (
+	"context"
+	"sync"
+
+	"github.com/m1k1o/neko/server/pkg/types"
+	"github.com/onkernel/kernel-images/server/lib/scaletozero"
+	"github.com/rs/zerolog"
+	"github.com/rs/zerolog/log"
+)
+
+func NewManager(
+	sessions types.SessionManager,
+	config *Config,
+) *Manager {
+	logger := log.With().Str("module", "scaletozero").Logger()
+
+	return &Manager{
+		logger:   logger,
+		config:   config,
+		sessions: sessions,
+		ctrl:     scaletozero.NewUnikraftCloudController(),
+	}
+}
+
+type Manager struct {
+	logger   zerolog.Logger
+	config   *Config
+	sessions types.SessionManager
+	ctrl     scaletozero.Controller
+	mu       sync.Mutex
+	shutdown bool
+	pending  int
+}
+
+func (m *Manager) Start() error {
+	if !m.config.Enabled {
+		return nil
+	}
+	m.logger.Info().Msg("scale-to-zero plugin enabled")
+
+	m.sessions.OnConnected(func(session types.Session) {
+		m.mu.Lock()
+		defer m.mu.Unlock()
+		if m.shutdown {
+			return
+		}
+
+		m.pending++
+		m.ctrl.Disable(context.Background())
+	})
+
+	m.sessions.OnDisconnected(func(session types.Session) {
+		m.mu.Lock()
+		defer m.mu.Unlock()
+		if m.shutdown {
+			return
+		}
+		m.pending--
+		m.ctrl.Enable(context.Background())
+	})
+
+	m.logger.Info().Msg("scale-to-zero hooks enabled")
+	return nil
+}
+
+func (m *Manager) Shutdown() error {
+	m.mu.Lock()
+	defer m.mu.Unlock()
+	m.shutdown = true
+
+	for i := 0; i < m.pending; i++ {
+		m.ctrl.Enable(context.Background())
+	}
+
+	return nil
+}

server/internal/plugins/scaletozero/plugin.go

@@ -0,0 +1,33 @@
+package scaletozero
+
+import (
+	"github.com/m1k1o/neko/server/pkg/types"
+)
+
+type Plugin struct {
+	config  *Config
+	manager *Manager
+}
+
+func NewPlugin() *Plugin {
+	return &Plugin{
+		config: &Config{},
+	}
+}
+
+func (p *Plugin) Name() string {
+	return PluginName
+}
+
+func (p *Plugin) Config() types.PluginConfig {
+	return p.config
+}
+
+func (p *Plugin) Start(m types.PluginManagers) error {
+	p.manager = NewManager(m.SessionManager, p.config)
+	return p.manager.Start()
+}
+
+func (p *Plugin) Shutdown() error {
+	return p.manager.Shutdown()
+}

server/internal/plugins/scaletozero/types.go

@@ -0,0 +1,3 @@
+package scaletozero
+
+const PluginName = "scaletozero"