Initial commit

backend/internal/api/stream.go

@@ -0,0 +1,111 @@
+package api
+
+import (
+	"net/http"
+	"strconv"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+)
+
+var upgrader = websocket.Upgrader{
+	CheckOrigin: func(r *http.Request) bool { return true },
+}
+
+type LogHub struct {
+	mu      sync.Mutex
+	streams map[uint][]chan []byte
+}
+
+var Hub = &LogHub{
+	streams: make(map[uint][]chan []byte),
+}
+
+func (h *LogHub) Broadcast(deploymentID uint, p []byte) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	if sinks, ok := h.streams[deploymentID]; ok {
+		for _, sink := range sinks {
+			select {
+			case sink <- p:
+			default:
+			}
+		}
+	}
+}
+
+func (h *LogHub) Subscribe(deploymentID uint) chan []byte {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	ch := make(chan []byte, 256)
+	h.streams[deploymentID] = append(h.streams[deploymentID], ch)
+	return ch
+}
+
+func (h *LogHub) Unsubscribe(deploymentID uint, ch chan []byte) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+	if sinks, ok := h.streams[deploymentID]; ok {
+		for i, sink := range sinks {
+			if sink == ch {
+				h.streams[deploymentID] = append(sinks[:i], sinks[i+1:]...)
+				close(ch)
+				break
+			}
+		}
+		if len(h.streams[deploymentID]) == 0 {
+			delete(h.streams, deploymentID)
+		}
+	}
+}
+
+type StreamWriter struct {
+	DeploymentID uint
+}
+
+func (w *StreamWriter) Write(p []byte) (n int, err error) {
+	c := make([]byte, len(p))
+	copy(c, p)
+	Hub.Broadcast(w.DeploymentID, c)
+	return len(p), nil
+}
+
+func (h *Handler) streamDeploymentLogs(c *gin.Context) {
+	deploymentIDStr := c.Param("id")
+	deploymentID, err := strconv.ParseUint(deploymentIDStr, 10, 64)
+	if err != nil {
+		c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid ID"})
+		return
+	}
+
+	conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		return
+	}
+	defer conn.Close()
+
+	logChan := Hub.Subscribe(uint(deploymentID))
+	defer Hub.Unsubscribe(uint(deploymentID), logChan)
+
+	go func() {
+		for {
+			if _, _, err := conn.NextReader(); err != nil {
+				conn.Close()
+				return
+			}
+		}
+	}()
+
+	for logChunk := range logChan {
+		conn.SetWriteDeadline(time.Now().Add(10 * time.Second))
+		if err := conn.WriteMessage(websocket.TextMessage, logChunk); err != nil {
+			return
+		}
+	}
+}
+
+func (h *Handler) RegisterStreamRoutes(r *gin.Engine) {
+	r.GET("/api/deployments/:id/logs/stream", h.streamDeploymentLogs)
+}