feat: config hot-reload

Author: sammcj

llama-swap.go

@@ -4,9 +4,11 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"net/http"
 	"os"
 	"os/signal"
 	"path/filepath"
+	"sync"
 	"syscall"
 	"time"
 
@@ -15,9 +17,14 @@ import (
 	"github.com/mostlygeek/llama-swap/proxy"
 )
 
-var version string = "0"
-var commit string = "abcd1234"
-var date = "unknown"
+var (
+	version string = "0"
+	commit  string = "abcd1234"
+	date    string = "unknown"
+
+	// Global mutex for server operations
+	serverMux sync.Mutex
+)
 
 func main() {
 	// Define a command-line flag for the port
@@ -58,20 +65,51 @@ func main() {
 		}
 	}
 
+	// Channel to signal when we want to exit
+	exitChan := make(chan struct{})
+
+	// Set up signal handling
 	sigChan := make(chan os.Signal, 1)
 	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
+
+	// Handle shutdown signals
 	go func() {
-		<-sigChan
-		fmt.Println("Shutting down llama-swap")
+		sig := <-sigChan
+		fmt.Printf("Received signal %v, shutting down...\n", sig)
 		proxyManager.Shutdown()
-		os.Exit(0)
+		close(exitChan)
 	}()
 
+	// Start server in main thread
 	fmt.Println("llama-swap listening on " + *listenStr)
-	if err := proxyManager.Run(*listenStr); err != nil {
-		fmt.Printf("Server error: %v\n", err)
-		os.Exit(1)
-	}
+	go func() {
+		for {
+			serverMux.Lock()
+			runChan := make(chan error)
+
+			// Start server in separate goroutine
+			go func() {
+				runChan <- proxyManager.Run(*listenStr)
+			}()
+
+			// Wait for server to complete or config reload
+			err = <-runChan
+
+			if err != nil && err != http.ErrServerClosed {
+				fmt.Printf("Fatal server error: %v\n", err)
+				serverMux.Unlock()
+				close(exitChan)
+				return
+			}
+
+			serverMux.Unlock()
+			// If we get here, it's because of a normal shutdown (config reload)
+			// Just continue the loop to let the new server start
+		}
+	}()
+
+	// Wait for exit signal
+	<-exitChan
 }
 
 // watchConfigFile monitors the configuration file for changes and recreates the ProxyManager with the new config.
@@ -91,15 +129,13 @@ func watchConfigFile(configPath string, pm **proxy.ProxyManager, listenStr *stri
 
 	log.Printf("Watching config file for changes: %s", configPath)
 
-	// Debounce timer
 	var debounceTimer *time.Timer
-	debounceDuration := 2 * time.Second // Adjust as needed
+	debounceDuration := 2 * time.Second
 
 	for {
 		select {
 		case event, ok := <-watcher.Events:
 			if !ok {
-				log.Println("File watcher channel closed.")
 				return
 			}
 			// We only care about writes to the specific config file
@@ -119,18 +155,20 @@ func watchConfigFile(configPath string, pm **proxy.ProxyManager, listenStr *stri
 					}
 
 					// Stop and cleanup old ProxyManager
+					serverMux.Lock()
 					oldPM := *pm
-					log.Println("Shutting down old ProxyManager instance before reload...")
 					oldPM.Shutdown()
 
 					// Create new ProxyManager with new config
 					newPM := proxy.New(newConfig)
 					*pm = newPM
+					serverMux.Unlock()
 
-					// Start serving with new ProxyManager
+					// Start serving with new ProxyManager (outside the lock)
 					go func() {
-						if err := newPM.Run(*listenStr); err != nil {
-							log.Printf("Server error after config reload: %v", err)
+						// Try to start the server
+						if err := newPM.Run(*listenStr); err != nil && err != http.ErrServerClosed {
+							log.Printf("Error starting server after config reload: %v", err)
 						}
 					}()
 

proxy/proxymanager.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"bytes"
+	"context"
 	"encoding/json"
 	"fmt"
 	"io"
@@ -29,6 +30,7 @@ type ProxyManager struct {
 	config           *Config
 	currentProcesses map[string]*Process
 	ginEngine        *gin.Engine
+	server           *http.Server
 
 	// logging
 	proxyLogger    *LogMonitor
@@ -189,7 +191,17 @@ func New(config *Config) *ProxyManager {
 }
 
 func (pm *ProxyManager) Run(addr ...string) error {
-	return pm.ginEngine.Run(addr...)
+	address := ":8080"
+	if len(addr) > 0 {
+		address = addr[0]
+	}
+
+	pm.server = &http.Server{
+		Addr:    address,
+		Handler: pm.ginEngine,
+	}
+
+	return pm.server.ListenAndServe()
 }
 
 func (pm *ProxyManager) HandlerFunc(w http.ResponseWriter, r *http.Request) {
@@ -229,13 +241,22 @@ func (pm *ProxyManager) stopProcesses() {
 	pm.currentProcesses = make(map[string]*Process)
 }
 
-// Shutdown is called to shutdown all upstream processes
-// when llama-swap is shutting down.
+// Shutdown gracefully shuts down the server and all upstream processes
 func (pm *ProxyManager) Shutdown() {
 	pm.Lock()
 	defer pm.Unlock()
 
-	// shutdown process in parallel
+	// First, shutdown the HTTP server gracefully
+	if pm.server != nil {
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		if err := pm.server.Shutdown(ctx); err != nil {
+			pm.proxyLogger.Errorf("HTTP server Shutdown: %v", err)
+		}
+	}
+
+	// Then shutdown all processes in parallel
 	var wg sync.WaitGroup
 	for _, process := range pm.currentProcesses {
 		wg.Add(1)