Files
gemini-banlancer/internal/logging/websocket_hook.go
2025-11-25 16:58:15 +08:00

37 lines
847 B
Go

// Filename: internal/logging/websocket_hook.go
package logging
import (
"gemini-balancer/internal/config"
"gemini-balancer/internal/handlers"
"github.com/sirupsen/logrus"
)
type WebSocketHook struct {
broadcaster func(*logrus.Entry)
}
func NewWebSocketHook(broadcaster func(*logrus.Entry)) *WebSocketHook {
return &WebSocketHook{broadcaster: broadcaster}
}
func (h *WebSocketHook) Levels() []logrus.Level {
return logrus.AllLevels
}
func (h *WebSocketHook) Fire(entry *logrus.Entry) error {
if h.broadcaster != nil {
go h.broadcaster(entry)
}
return nil
}
func NewLoggerWithWebSocket(cfg *config.Config) (*logrus.Logger, *handlers.WebSocketHandler) {
logger := NewLogger(cfg)
wsHandler := handlers.NewWebSocketHandler(logger)
hook := NewWebSocketHook(wsHandler.BroadcastLog)
logger.AddHook(hook)
return logger, wsHandler
}