Updates from my site

api/sse.go

@@ -1,15 +1,19 @@
 package api
 
 import (
+	"errors"
 	"fmt"
-	"log"
 	"time"
 
 	"github.com/labstack/echo/v4"
 	"goth.stack/lib"
 )
 
-func SSE(c echo.Context) error {
+func SSE(c echo.Context, pubSub lib.PubSub) error {
+	if pubSub == nil {
+		return errors.New("pubSub is nil")
+	}
+
 	channel := c.QueryParam("channel")
 	if channel == "" {
 		channel = "default"
@@ -18,41 +22,29 @@ func SSE(c echo.Context) error {
 	// Use the request context, which is cancelled when the client disconnects
 	ctx := c.Request().Context()
 
-	pubsub, _ := lib.Subscribe(lib.RedisClient, channel)
+	pubsub, err := pubSub.SubscribeToChannel(channel)
+	if err != nil {
+		return fmt.Errorf("failed to subscribe to channel: %w", err)
+	}
 
-	c.Response().Header().Set(echo.HeaderContentType, "text/event-stream")
-	c.Response().Header().Set(echo.HeaderConnection, "keep-alive")
-	c.Response().Header().Set(echo.HeaderCacheControl, "no-cache")
+	lib.SetSSEHeaders(c)
 
 	// Create a ticker that fires every 15 seconds
-	ticker := time.NewTicker(30 * time.Second)
+	ticker := lib.CreateTickerAndKeepAlive(c, 30*time.Second)
 	defer ticker.Stop()
 
+	// Create a client channel and add it to the SSE server
+	client := make(chan string)
+	lib.SSEServer.AddClient(channel, client)
+	defer lib.SSEServer.RemoveClient(channel, client)
+
+	go lib.HandleIncomingMessages(c, pubsub, client)
+
 	for {
 		select {
 		case <-ctx.Done():
 			// If the client has disconnected, stop the loop
 			return nil
-		case <-ticker.C:
-			// Every 30 seconds, send a comment to keep the connection alive
-			if _, err := c.Response().Write([]byte(": keep-alive\n\n")); err != nil {
-				return err
-			}
-			c.Response().Flush()
-		default:
-			// Handle incoming messages as before
-			msg, err := pubsub.ReceiveMessage(ctx)
-			if err != nil {
-				log.Printf("Failed to receive message: %v", err)
-				continue
-			}
-
-			data := fmt.Sprintf("data: %s\n\n", msg.Payload)
-			if _, err := c.Response().Write([]byte(data)); err != nil {
-				return err
-			}
-
-			c.Response().Flush()
 		}
 	}
 }