add sse handle with std lib

broker.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"html/template"
 	"log/slog"
+	"net/http"
 	"strings"
 	"sync"
 	"time"
@@ -120,17 +121,42 @@ func (b *SSEBroker) Broadcast(event string, data ...any) {
 	b.broadcast <- Message{Event: event, Data: payload}
 }
 
-func (b *SSEBroker) HandleFiberCtxSSE(c *fiber.Ctx) error {
+// Interface para abstraer las diferencias entre Fiber y HTTP estándar
+type SSEWriter interface {
+	Write(data []byte) (int, error)
+	Flush() error
+}
 
-	c.Set("Content-Type", "text/event-stream")
+// Wrapper para http.ResponseWriter
-	c.Set("Cache-Control", "no-cache")
+type httpSSEWriter struct {
-	c.Set("Connection", "keep-alive")
+	w http.ResponseWriter
-	c.Set("Transfer-Encoding", "chunked")
+}
 
-	clientID := uuid.New().String()
+func (h *httpSSEWriter) Write(data []byte) (int, error) {
-	client := b.RegisterClient(clientID)
+	return h.w.Write(data)
+}
 
-	c.Status(fiber.StatusOK).Context().SetBodyStreamWriter(fasthttp.StreamWriter(func(w *bufio.Writer) {
+func (h *httpSSEWriter) Flush() error {
+	if flusher, ok := h.w.(http.Flusher); ok {
+		flusher.Flush()
+	}
+	return nil
+}
+
+// Wrapper para bufio.Writer
+type fiberSSEWriter struct {
+	w *bufio.Writer
+}
+
+func (f *fiberSSEWriter) Write(data []byte) (int, error) {
+	return f.w.Write(data)
+}
+
+func (f *fiberSSEWriter) Flush() error {
+	return f.w.Flush()
+}
+
+func (b *SSEBroker) handleSSEConnection(clientID string, client *Client, writer SSEWriter) {
 	slog.Info("SSE connection established", "client_id", clientID)
 
 	defer func() {
@@ -138,9 +164,9 @@ func (b *SSEBroker) HandleFiberCtxSSE(c *fiber.Ctx) error {
 		slog.Info("SSE connection closed", "client_id", clientID)
 	}()
 
-		fmt.Fprintf(w, "event: client_id\n")
+	fmt.Fprintf(writer, "event: client_id\n")
-		fmt.Fprintf(w, "data: %s\n\n", clientID)
+	fmt.Fprintf(writer, "data: %s\n\n", clientID)
-		w.Flush()
+	writer.Flush()
 
 	ticker := time.NewTicker(time.Minute)
 	defer ticker.Stop()
@@ -163,37 +189,65 @@ func (b *SSEBroker) HandleFiberCtxSSE(c *fiber.Ctx) error {
 				data = string(jsonData)
 			}
 
-				fmt.Fprintf(w, "event: %s\n", message.Event)
+			fmt.Fprintf(writer, "event: %s\n", message.Event)
 			scanner := bufio.NewScanner(strings.NewReader(data))
 			for scanner.Scan() {
-					fmt.Fprintf(w, "data: %s\n", scanner.Text())
+				fmt.Fprintf(writer, "data: %s\n", scanner.Text())
 			}
-				fmt.Fprint(w, "\n")
+			fmt.Fprint(writer, "\n")
 
-				if err := w.Flush(); err != nil {
+			if err := writer.Flush(); err != nil {
 				slog.Warn("flush error, closing client", "client_id", clientID)
 				return
 			}
+
 		case <-ticker.C:
-				fmt.Fprintf(w, "event: ping\n")
+			fmt.Fprintf(writer, "event: ping\n")
-				fmt.Fprintf(w, "data: %s\n\n", "ping")
+			fmt.Fprintf(writer, "data: %s\n\n", "ping")
+			writer.Flush()
 
 		case <-client.Close:
 			return
+		}
+	}
+}
 
-			default:
+func (b *SSEBroker) HandleFiberCtxSSE(c *fiber.Ctx) error {
-				err := w.Flush()
+	c.Set("Content-Type", "text/event-stream")
-				if err != nil {
+	c.Set("Cache-Control", "no-cache")
-					slog.Warn("error while flushing", "error", err)
+	c.Set("Connection", "keep-alive")
-					return
+	c.Set("Transfer-Encoding", "chunked")
-				}
+
-			}
+	clientID := uuid.New().String()
-		}
+	client := b.RegisterClient(clientID)
+
+	c.Status(fiber.StatusOK).Context().SetBodyStreamWriter(fasthttp.StreamWriter(func(w *bufio.Writer) {
+		writer := &fiberSSEWriter{w: w}
+		b.handleSSEConnection(clientID, client, writer)
 	}))
 
 	return nil
 }
 
+func (b *SSEBroker) HandleHTTPSSE(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "text/event-stream")
+	w.Header().Set("Cache-Control", "no-cache")
+	w.Header().Set("Connection", "keep-alive")
+	w.Header().Set("Transfer-Encoding", "chunked")
+
+	clientID := uuid.New().String()
+	client := b.RegisterClient(clientID)
+
+	writer := &httpSSEWriter{w: w}
+
+	go func() {
+		b.handleSSEConnection(clientID, client, writer)
+	}()
+
+	<-r.Context().Done()
+	slog.Info("client disconnected", "client_id", clientID)
+}
+
 func (b *SSEBroker) Shutdown() {
 	close(b.done)
 }