Implement graceful shutdown for the server and add Shutdown method to InstanceManager

cmd/server/main.go

@@ -5,6 +5,8 @@ import (
 	llamactl "llamactl/pkg"
 	"net/http"
 	"os"
+	"os/signal"
+	"syscall"
 )
 
 // @title llamactl API
@@ -37,7 +39,34 @@ func main() {
 	// Setup the router with the handler
 	r := llamactl.SetupRouter(handler)
 
-	// Start the server with the router
-	fmt.Printf("Starting llamactl on port %d...\n", config.Server.Port)
-	http.ListenAndServe(fmt.Sprintf("%s:%d", config.Server.Host, config.Server.Port), r)
+	// Handle graceful shutdown
+	stop := make(chan os.Signal, 1)
+	signal.Notify(stop, os.Interrupt, syscall.SIGTERM)
+
+	server := http.Server{
+		Addr:    fmt.Sprintf("%s:%d", config.Server.Host, config.Server.Port),
+		Handler: r,
+	}
+
+	go func() {
+		fmt.Printf("Llamactl server listening on %s:%d\n", config.Server.Host, config.Server.Port)
+		if err := server.ListenAndServe(); err != nil && err != http.ErrServerClosed {
+			fmt.Printf("Error starting server: %v\n", err)
+		}
+	}()
+
+	// Wait for shutdown signal
+	<-stop
+	fmt.Println("Shutting down server...")
+
+	if err := server.Close(); err != nil {
+		fmt.Printf("Error shutting down server: %v\n", err)
+	} else {
+		fmt.Println("Server shut down gracefully.")
+	}
+
+	// Wait for all instances to stop
+	instanceManager.Shutdown()
+
+	fmt.Println("Exiting llamactl.")
 }

pkg/manager.go

@@ -16,6 +16,7 @@ type InstanceManager interface {
 	StopInstance(name string) (*Instance, error)
 	RestartInstance(name string) (*Instance, error)
 	GetInstanceLogs(name string) (string, error)
+	Shutdown()
 }
 
 type instanceManager struct {
@@ -247,3 +248,30 @@ func (im *instanceManager) getNextAvailablePort() (int, error) {
 
 	return 0, fmt.Errorf("no available ports in the specified range")
 }
+
+func (im *instanceManager) Shutdown() {
+	im.mu.Lock()
+	defer im.mu.Unlock()
+
+	var wg sync.WaitGroup
+	wg.Add(len(im.instances))
+
+	for name, instance := range im.instances {
+		if !instance.Running {
+			wg.Done() // If instance is not running, just mark it as done
+			continue
+		}
+
+		go func(name string, instance *Instance) {
+			defer wg.Done()
+			fmt.Printf("Stopping instance %s...\n", name)
+			// Attempt to stop the instance gracefully
+			if err := instance.Stop(); err != nil {
+				fmt.Printf("Error stopping instance %s: %v\n", name, err)
+			}
+		}(name, instance)
+	}
+
+	wg.Wait()
+	fmt.Println("All instances stopped.")
+}