Add Event Bus (#184)

Major internal refactor to use an event bus to pass event/messages along. These changes are largely invisible user facing but sets up internal design for real time stats and information.

- `--watch-config` logic refactored for events
- remove multiple SSE api endpoints, replaced with /api/events
- keep all functionality essentially the same
- UI/backend sync is in near real time now

proxy/events.go

@@ -0,0 +1,49 @@
+package proxy
+
+// package level registry of the different event types
+
+const ProcessStateChangeEventID = 0x01
+const ChatCompletionStatsEventID = 0x02
+const ConfigFileChangedEventID = 0x03
+const LogDataEventID = 0x04
+
+type ProcessStateChangeEvent struct {
+	ProcessName string
+	NewState    ProcessState
+	OldState    ProcessState
+}
+
+func (e ProcessStateChangeEvent) Type() uint32 {
+	return ProcessStateChangeEventID
+}
+
+type ChatCompletionStats struct {
+	TokensGenerated int
+}
+
+func (e ChatCompletionStats) Type() uint32 {
+	return ChatCompletionStatsEventID
+}
+
+type ReloadingState int
+
+const (
+	ReloadingStateStart ReloadingState = iota
+	ReloadingStateEnd
+)
+
+type ConfigFileChangedEvent struct {
+	ReloadingState ReloadingState
+}
+
+func (e ConfigFileChangedEvent) Type() uint32 {
+	return ConfigFileChangedEventID
+}
+
+type LogDataEvent struct {
+	Data []byte
+}
+
+func (e LogDataEvent) Type() uint32 {
+	return LogDataEventID
+}

proxy/logMonitor.go

@@ -2,10 +2,13 @@ package proxy
 
 import (
 	"container/ring"
+	"context"
 	"fmt"
 	"io"
 	"os"
 	"sync"
+
+	"github.com/kelindar/event"
 )
 
 type LogLevel int
@@ -18,7 +21,7 @@ const (
 )
 
 type LogMonitor struct {
-	clients  map[chan []byte]bool
+	eventbus *event.Dispatcher
 	mu       sync.RWMutex
 	buffer   *ring.Ring
 	bufferMu sync.RWMutex
@@ -37,11 +40,11 @@ func NewLogMonitor() *LogMonitor {
 
 func NewLogMonitorWriter(stdout io.Writer) *LogMonitor {
 	return &LogMonitor{
-		clients: make(map[chan []byte]bool),
-		buffer:  ring.New(10 * 1024), // keep 10KB of buffered logs
-		stdout:  stdout,
-		level:   LevelInfo,
-		prefix:  "",
+		eventbus: event.NewDispatcher(),
+		buffer:   ring.New(10 * 1024), // keep 10KB of buffered logs
+		stdout:   stdout,
+		level:    LevelInfo,
+		prefix:   "",
 	}
 }
 
@@ -81,34 +84,14 @@ func (w *LogMonitor) GetHistory() []byte {
 	return history
 }
 
-func (w *LogMonitor) Subscribe() chan []byte {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-
-	ch := make(chan []byte, 100)
-	w.clients[ch] = true
-	return ch
-}
-
-func (w *LogMonitor) Unsubscribe(ch chan []byte) {
-	w.mu.Lock()
-	defer w.mu.Unlock()
-
-	delete(w.clients, ch)
-	close(ch)
+func (w *LogMonitor) OnLogData(callback func(data []byte)) context.CancelFunc {
+	return event.Subscribe(w.eventbus, func(e LogDataEvent) {
+		callback(e.Data)
+	})
 }
 
 func (w *LogMonitor) broadcast(msg []byte) {
-	w.mu.RLock()
-	defer w.mu.RUnlock()
-
-	for client := range w.clients {
-		select {
-		case client <- msg:
-		default:
-			// If client buffer is full, skip
-		}
-	}
+	event.Publish(w.eventbus, LogDataEvent{Data: msg})
 }
 
 func (w *LogMonitor) SetPrefix(prefix string) {

proxy/logMonitor_test.go

@@ -10,38 +10,29 @@ import (
 func TestLogMonitor(t *testing.T) {
 	logMonitor := NewLogMonitorWriter(io.Discard)
 
-	// Test subscription
-	client1 := logMonitor.Subscribe()
-	client2 := logMonitor.Subscribe()
-
-	defer logMonitor.Unsubscribe(client1)
-	defer logMonitor.Unsubscribe(client2)
+	// A WaitGroup is used to wait for all the expected writes to complete
+	var wg sync.WaitGroup
 
 	client1Messages := make([]byte, 0)
 	client2Messages := make([]byte, 0)
 
-	var wg sync.WaitGroup
-	wg.Add(1)
+	defer logMonitor.OnLogData(func(data []byte) {
+		client1Messages = append(client1Messages, data...)
+		wg.Done()
+	})()
 
-	go func() {
-		defer wg.Done()
-		for {
-			select {
-			case data := <-client1:
-				client1Messages = append(client1Messages, data...)
-			case data := <-client2:
-				client2Messages = append(client2Messages, data...)
-			default:
-				return
-			}
-		}
-	}()
+	defer logMonitor.OnLogData(func(data []byte) {
+		client2Messages = append(client2Messages, data...)
+		wg.Done()
+	})()
+
+	wg.Add(6) // 2 x 3 writes
 
 	logMonitor.Write([]byte("1"))
 	logMonitor.Write([]byte("2"))
 	logMonitor.Write([]byte("3"))
 
-	// Wait for the goroutine to finish
+	// wait for all writes to complete
 	wg.Wait()
 
 	// Check the buffer

proxy/process.go

@@ -13,6 +13,8 @@ import (
 	"sync"
 	"syscall"
 	"time"
+
+	"github.com/kelindar/event"
 )
 
 type ProcessState string
@@ -127,6 +129,7 @@ func (p *Process) swapState(expectedState, newState ProcessState) (ProcessState,
 
 	p.state = newState
 	p.proxyLogger.Debugf("<%s> swapState() State transitioned from %s to %s", p.ID, expectedState, newState)
+	event.Emit(ProcessStateChangeEvent{ProcessName: p.ID, NewState: newState, OldState: expectedState})
 	return p.state, nil
 }
 

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
@@ -157,9 +167,7 @@ func (pm *ProxyManager) setupGinEngine() {
 	// in proxymanager_loghandlers.go
 	pm.ginEngine.GET("/logs", pm.sendLogsHandlers)
 	pm.ginEngine.GET("/logs/stream", pm.streamLogsHandler)
-	pm.ginEngine.GET("/logs/streamSSE", pm.streamLogsHandlerSSE)
 	pm.ginEngine.GET("/logs/stream/:logMonitorID", pm.streamLogsHandler)
-	pm.ginEngine.GET("/logs/streamSSE/:logMonitorID", pm.streamLogsHandlerSSE)
 
 	/**
 	 * User Interface Endpoints
@@ -261,6 +269,7 @@ func (pm *ProxyManager) Shutdown() {
 		}(processGroup)
 	}
 	wg.Wait()
+	pm.shutdownCancel()
 }
 
 func (pm *ProxyManager) swapProcessGroup(requestedModel string) (*ProcessGroup, string, error) {

proxy/proxymanager_api.go

@@ -1,25 +1,28 @@
 package proxy
 
 import (
+	"context"
+	"encoding/json"
 	"net/http"
 	"sort"
-	"time"
 
 	"github.com/gin-gonic/gin"
+	"github.com/kelindar/event"
 )
 
 type Model struct {
-	Id    string `json:"id"`
-	State string `json:"state"`
+	Id          string `json:"id"`
+	Name        string `json:"name"`
+	Description string `json:"description"`
+	State       string `json:"state"`
 }
 
 func addApiHandlers(pm *ProxyManager) {
 	// Add API endpoints for React to consume
 	apiGroup := pm.ginEngine.Group("/api")
 	{
-		apiGroup.GET("/models", pm.apiListModels)
-		apiGroup.GET("/modelsSSE", pm.apiListModelsSSE)
 		apiGroup.POST("/models/unload", pm.apiUnloadAllModels)
+		apiGroup.GET("/events", pm.apiSendEvents)
 	}
 }
 
@@ -65,37 +68,102 @@ func (pm *ProxyManager) getModelStatus() []Model {
 			}
 		}
 		models = append(models, Model{
-			Id:    modelID,
-			State: state,
+			Id:          modelID,
+			Name:        pm.config.Models[modelID].Name,
+			Description: pm.config.Models[modelID].Description,
+			State:       state,
 		})
 	}
 
 	return models
 }
 
-func (pm *ProxyManager) apiListModels(c *gin.Context) {
-	c.JSON(http.StatusOK, pm.getModelStatus())
+type messageType string
+
+const (
+	msgTypeModelStatus messageType = "modelStatus"
+	msgTypeLogData     messageType = "logData"
+)
+
+type messageEnvelope struct {
+	Type messageType `json:"type"`
+	Data string      `json:"data"`
 }
 
-// stream the models as a SSE
-func (pm *ProxyManager) apiListModelsSSE(c *gin.Context) {
+// sends a stream of different message types that happen on the server
+func (pm *ProxyManager) apiSendEvents(c *gin.Context) {
 	c.Header("Content-Type", "text/event-stream")
 	c.Header("Cache-Control", "no-cache")
 	c.Header("Connection", "keep-alive")
 	c.Header("X-Content-Type-Options", "nosniff")
 
-	notify := c.Request.Context().Done()
+	sendBuffer := make(chan messageEnvelope, 25)
+	ctx, cancel := context.WithCancel(c.Request.Context())
+	sendModels := func() {
+		data, err := json.Marshal(pm.getModelStatus())
+		if err == nil {
+			msg := messageEnvelope{Type: msgTypeModelStatus, Data: string(data)}
+			select {
+			case sendBuffer <- msg:
+			case <-ctx.Done():
+				return
+			default:
+			}
+
+		}
+	}
+
+	sendLogData := func(source string, data []byte) {
+		data, err := json.Marshal(gin.H{
+			"source": source,
+			"data":   string(data),
+		})
+		if err == nil {
+			select {
+			case sendBuffer <- messageEnvelope{Type: msgTypeLogData, Data: string(data)}:
+			case <-ctx.Done():
+				return
+			default:
+			}
+		}
+	}
+
+	/**
+	 * Send updated models list
+	 */
+	defer event.On(func(e ProcessStateChangeEvent) {
+		sendModels()
+	})()
+	defer event.On(func(e ConfigFileChangedEvent) {
+		sendModels()
+	})()
+
+	/**
+	 * Send Log data
+	 */
+	defer pm.proxyLogger.OnLogData(func(data []byte) {
+		sendLogData("proxy", data)
+	})()
+	defer pm.upstreamLogger.OnLogData(func(data []byte) {
+		sendLogData("upstream", data)
+	})()
+
+	// send initial batch of data
+	sendLogData("proxy", pm.proxyLogger.GetHistory())
+	sendLogData("upstream", pm.upstreamLogger.GetHistory())
+	sendModels()
 
-	// Stream new events
 	for {
 		select {
-		case <-notify:
+		case <-c.Request.Context().Done():
+			cancel()
 			return
-		default:
-			models := pm.getModelStatus()
-			c.SSEvent("message", models)
+		case <-pm.shutdownCtx.Done():
+			cancel()
+			return
+		case msg := <-sendBuffer:
+			c.SSEvent("message", msg)
 			c.Writer.Flush()
-			<-time.After(1000 * time.Millisecond)
 		}
 	}
 }

proxy/proxymanager_loghandlers.go

@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"strings"
@@ -34,10 +35,7 @@ func (pm *ProxyManager) streamLogsHandler(c *gin.Context) {
 		c.String(http.StatusBadRequest, err.Error())
 		return
 	}
-	ch := logger.Subscribe()
-	defer logger.Unsubscribe(ch)
 
-	notify := c.Request.Context().Done()
 	flusher, ok := c.Writer.(http.Flusher)
 	if !ok {
 		c.AbortWithError(http.StatusInternalServerError, fmt.Errorf("streaming unsupported"))
@@ -55,57 +53,28 @@ func (pm *ProxyManager) streamLogsHandler(c *gin.Context) {
 		}
 	}
 
-	// Stream new logs
+	sendChan := make(chan []byte, 10)
+	ctx, cancel := context.WithCancel(c.Request.Context())
+	defer logger.OnLogData(func(data []byte) {
+		select {
+		case sendChan <- data:
+		case <-ctx.Done():
+			return
+		default:
+		}
+	})()
+
 	for {
 		select {
-		case msg := <-ch:
-			_, err := c.Writer.Write(msg)
-			if err != nil {
-				// just break the loop if we can't write for some reason
-				return
-			}
+		case <-c.Request.Context().Done():
+			cancel()
+			return
+		case <-pm.shutdownCtx.Done():
+			cancel()
+			return
+		case data := <-sendChan:
+			c.Writer.Write(data)
 			flusher.Flush()
-		case <-notify:
-			return
-		}
-	}
-}
-
-func (pm *ProxyManager) streamLogsHandlerSSE(c *gin.Context) {
-	c.Header("Content-Type", "text/event-stream")
-	c.Header("Cache-Control", "no-cache")
-	c.Header("Connection", "keep-alive")
-	c.Header("X-Content-Type-Options", "nosniff")
-
-	logMonitorId := c.Param("logMonitorID")
-	logger, err := pm.getLogger(logMonitorId)
-	if err != nil {
-		c.String(http.StatusBadRequest, err.Error())
-		return
-	}
-	ch := logger.Subscribe()
-	defer logger.Unsubscribe(ch)
-
-	notify := c.Request.Context().Done()
-
-	// Send history first if not skipped
-	_, skipHistory := c.GetQuery("no-history")
-	if !skipHistory {
-		history := logger.GetHistory()
-		if len(history) != 0 {
-			c.SSEvent("message", string(history))
-			c.Writer.Flush()
-		}
-	}
-
-	// Stream new logs
-	for {
-		select {
-		case msg := <-ch:
-			c.SSEvent("message", string(msg))
-			c.Writer.Flush()
-		case <-notify:
-			return
 		}
 	}
 }