优化流式传输&fix bugs

internal/handlers/websocket_handler.go

@@ -0,0 +1,68 @@
+// Filename: internal/handlers/websocket_handler.go
+package handlers
+
+import (
+	"net/http"
+	"sync"
+	"time"
+
+	"github.com/gin-gonic/gin"
+	"github.com/gorilla/websocket"
+	"github.com/sirupsen/logrus"
+)
+
+type connWrapper struct {
+	conn *websocket.Conn
+	mu   sync.Mutex
+}
+
+type WebSocketHandler struct {
+	logger   *logrus.Logger
+	clients  sync.Map
+	upgrader websocket.Upgrader
+}
+
+func NewWebSocketHandler(logger *logrus.Logger) *WebSocketHandler {
+	return &WebSocketHandler{
+		logger: logger,
+		upgrader: websocket.Upgrader{
+			CheckOrigin: func(r *http.Request) bool { return true },
+		},
+	}
+}
+
+func (h *WebSocketHandler) HandleSystemLogs(c *gin.Context) {
+	conn, err := h.upgrader.Upgrade(c.Writer, c.Request, nil)
+	if err != nil {
+		h.logger.WithError(err).Error("Failed to upgrade websocket")
+		return
+	}
+	defer conn.Close()
+
+	clientID := time.Now().UnixNano()
+	h.clients.Store(clientID, &connWrapper{conn: conn})
+	defer h.clients.Delete(clientID)
+
+	for {
+		if _, _, err := conn.ReadMessage(); err != nil {
+			break
+		}
+	}
+}
+
+func (h *WebSocketHandler) BroadcastLog(entry *logrus.Entry) {
+	msg := map[string]interface{}{
+		"timestamp": entry.Time.Format(time.RFC3339),
+		"level":     entry.Level.String(),
+		"message":   entry.Message,
+		"fields":    entry.Data,
+	}
+
+	h.clients.Range(func(key, value interface{}) bool {
+		wrapper := value.(*connWrapper)
+		wrapper.mu.Lock()
+		wrapper.conn.WriteJSON(msg)
+		wrapper.mu.Unlock()
+		return true
+	})
+}