fix blocked shutdown due to ui

Author: mostlygeek

proxy/proxymanager.go

@@ -2,6 +2,7 @@ package proxy
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"mime/multipart"
@@ -33,6 +34,10 @@ type ProxyManager struct {
 	muxLogger      *LogMonitor
 
 	processGroups map[string]*ProcessGroup
+
+	// shutdown signaling
+	shutdownCtx    context.Context
+	shutdownCancel context.CancelFunc
 }
 
 func New(config Config) *ProxyManager {
@@ -63,6 +68,8 @@ func New(config Config) *ProxyManager {
 		upstreamLogger.SetLogLevel(LevelInfo)
 	}
 
+	shutdownCtx, shutdownCancel := context.WithCancel(context.Background())
+
 	pm := &ProxyManager{
 		config:    config,
 		ginEngine: gin.New(),
@@ -72,6 +79,9 @@ func New(config Config) *ProxyManager {
 		upstreamLogger: upstreamLogger,
 
 		processGroups: make(map[string]*ProcessGroup),
+
+		shutdownCtx:    shutdownCtx,
+		shutdownCancel: shutdownCancel,
 	}
 
 	// create the process groups
@@ -261,6 +271,7 @@ func (pm *ProxyManager) Shutdown() {
 		}(processGroup)
 	}
 	wg.Wait()
+	pm.shutdownCancel()
 }
 
 func (pm *ProxyManager) swapProcessGroup(requestedModel string) (*ProcessGroup, string, error) {

proxy/proxymanager_api.go

@@ -96,5 +96,9 @@ func (pm *ProxyManager) apiListModelsSSE(c *gin.Context) {
 		}
 	})()
 
-	<-c.Request.Context().Done()
+	select {
+	case <-c.Request.Context().Done():
+	case <-pm.shutdownCtx.Done():
+	}
+
 }

proxy/proxymanager_loghandlers.go

@@ -52,14 +52,18 @@ func (pm *ProxyManager) streamLogsHandler(c *gin.Context) {
 		}
 	}
 
-	cancelFn := logger.OnLogData(func(data []byte) {
+	defer logger.OnLogData(func(data []byte) {
 		if c != nil && c.Writer != nil {
 			c.Writer.Write(data)
 			flusher.Flush()
 		}
-	})
-	defer cancelFn()
-	<-c.Request.Context().Done()
+	})()
+
+	select {
+	case <-c.Request.Context().Done():
+	case <-pm.shutdownCtx.Done():
+	}
+
 }
 
 func (pm *ProxyManager) streamLogsHandlerSSE(c *gin.Context) {
@@ -85,14 +89,17 @@ func (pm *ProxyManager) streamLogsHandlerSSE(c *gin.Context) {
 		}
 	}
 
-	cancelFn := logger.OnLogData(func(data []byte) {
+	defer logger.OnLogData(func(data []byte) {
 		if c != nil && c.Writer != nil {
 			c.SSEvent("message", string(data))
 			c.Writer.Flush()
 		}
-	})
-	defer cancelFn()
-	<-c.Request.Context().Done()
+	})()
+
+	select {
+	case <-c.Request.Context().Done():
+	case <-pm.shutdownCtx.Done():
+	}
 }
 
 // getLogger searches for the appropriate logger based on the logMonitorId