replace SSE streams with single unified one

Author: mostlygeek

proxy/logMonitor_test.go

@@ -3,27 +3,38 @@ package proxy
 import (
 	"bytes"
 	"io"
+	"sync"
 	"testing"
 )
 
 func TestLogMonitor(t *testing.T) {
 	logMonitor := NewLogMonitorWriter(io.Discard)
 
+	// 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)
 
 	defer logMonitor.OnLogData(func(data []byte) {
 		client1Messages = append(client1Messages, data...)
+		wg.Done()
 	})()
 
 	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 all writes to complete
+	wg.Wait()
+
 	// Check the buffer
 	expectedHistory := "123"
 	history := string(logMonitor.GetHistory())

proxy/proxymanager.go

@@ -167,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

proxy/proxymanager_api.go

@@ -1,6 +1,8 @@
 package proxy
 
 import (
+	"context"
+	"encoding/json"
 	"net/http"
 	"sort"
 
@@ -17,9 +19,8 @@ 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)
 	}
 }
 
@@ -73,42 +74,92 @@ func (pm *ProxyManager) getModelStatus() []Model {
 	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")
 
-	// flush the first one
-	c.SSEvent("message", pm.getModelStatus())
-	c.Writer.Flush()
+	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:
+			}
 
-	// send whenever the any process state
-	defer event.On(func(e ProcessStateChangeEvent) {
-		if c != nil && c.Writer != nil {
-			models := pm.getModelStatus()
-			c.SSEvent("message", models)
-			c.Writer.Flush()
 		}
-	})()
+	}
 
-	// resend the models when the config is reloaded
-	defer event.On(func(e ConfigFileChangedEvent) {
-		if c != nil && c.Writer != nil && e.ReloadingState == "end" {
-			models := pm.getModelStatus()
-			c.SSEvent("message", models)
-			c.Writer.Flush()
+	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()
 	})()
 
-	select {
-	case <-c.Request.Context().Done():
-	case <-pm.shutdownCtx.Done():
-	}
+	/**
+	 * 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()
+
+	for {
+		select {
+		case <-c.Request.Context().Done():
+			cancel()
+			return
+		case <-pm.shutdownCtx.Done():
+			cancel()
+			return
+		case msg := <-sendBuffer:
+			c.SSEvent("message", msg)
+			c.Writer.Flush()
+		}
+	}
 }

proxy/proxymanager_loghandlers.go

@@ -1,6 +1,7 @@
 package proxy
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"strings"
@@ -52,53 +53,29 @@ func (pm *ProxyManager) streamLogsHandler(c *gin.Context) {
 		}
 	}
 
+	sendChan := make(chan []byte, 10)
+	ctx, cancel := context.WithCancel(c.Request.Context())
 	defer logger.OnLogData(func(data []byte) {
-		if c != nil && c.Writer != nil {
-			c.Writer.Write(data)
-			flusher.Flush()
+		select {
+		case sendChan <- data:
+		case <-ctx.Done():
+			return
+		default:
 		}
 	})()
 
-	select {
-	case <-c.Request.Context().Done():
-	case <-pm.shutdownCtx.Done():
-	}
-
-}
-
-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
-	}
-
-	// 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()
-		}
-	}
-
-	defer logger.OnLogData(func(data []byte) {
-		if c != nil && c.Writer != nil {
-			c.SSEvent("message", string(data))
-			c.Writer.Flush()
+	for {
+		select {
+		case <-c.Request.Context().Done():
+			cancel()
+			return
+		case <-pm.shutdownCtx.Done():
+			cancel()
+			return
+		case data := <-sendChan:
+			c.Writer.Write(data)
+			flusher.Flush()
 		}
-	})()
-
-	select {
-	case <-c.Request.Context().Done():
-	case <-pm.shutdownCtx.Done():
 	}
 }