diff --git a/cmd/server/main.go b/cmd/server/main.go index dc58eaa..b0bc1ab 100644 --- a/cmd/server/main.go +++ b/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.") } diff --git a/pkg/manager.go b/pkg/manager.go index b7808a9..d2f0757 100644 --- a/pkg/manager.go +++ b/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.") +}