feat: finish project basic structure and core functions
This commit is contained in:
@@ -0,0 +1,295 @@
|
||||
package komari
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"nukumizu-backend/config"
|
||||
"nukumizu-backend/internal/node"
|
||||
"nukumizu-backend/postLog"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
// WSMessage is the structure of messages received from Komari WebSocket.
|
||||
type WSMessage struct {
|
||||
Status string `json:"status"`
|
||||
Data struct {
|
||||
Online []string `json:"online"`
|
||||
Data map[string]node.Report `json:"data"`
|
||||
} `json:"data"`
|
||||
}
|
||||
|
||||
// WSClient manages a WebSocket connection to Komari for real-time status updates.
|
||||
type WSClient struct {
|
||||
komariURL string
|
||||
conn *websocket.Conn
|
||||
connMu sync.Mutex
|
||||
reconnectCount atomic.Int32
|
||||
maxRetries int
|
||||
stopCh chan struct{}
|
||||
running atomic.Bool
|
||||
onReconnectFail func() // Callback when all reconnection attempts fail.
|
||||
}
|
||||
|
||||
// NewWSClient creates a new WebSocket client for Komari.
|
||||
func NewWSClient(komariURL string, maxRetries int) *WSClient {
|
||||
return &WSClient{
|
||||
komariURL: komariURL,
|
||||
maxRetries: maxRetries,
|
||||
stopCh: make(chan struct{}),
|
||||
}
|
||||
}
|
||||
|
||||
// SetOnReconnectFail sets the callback invoked when all reconnection attempts fail.
|
||||
func (w *WSClient) SetOnReconnectFail(cb func()) {
|
||||
w.onReconnectFail = cb
|
||||
}
|
||||
|
||||
// Connect establishes the WebSocket connection and starts the read loop.
|
||||
func (w *WSClient) Connect() error {
|
||||
postLog.Info("Connecting to Komari WebSocket at " + w.komariURL)
|
||||
|
||||
conn, err := w.dial()
|
||||
if err != nil {
|
||||
postLog.Error("Failed to connect to Komari WebSocket: " + err.Error())
|
||||
return err
|
||||
}
|
||||
|
||||
w.connMu.Lock()
|
||||
w.conn = conn
|
||||
w.connMu.Unlock()
|
||||
|
||||
w.reconnectCount.Store(0)
|
||||
w.running.Store(true)
|
||||
|
||||
// Request initial state.
|
||||
if err := conn.WriteMessage(websocket.TextMessage, []byte("get")); err != nil {
|
||||
postLog.Error("Failed to send 'get' message: " + err.Error())
|
||||
return err
|
||||
}
|
||||
|
||||
postLog.Info("Connected to Komari WebSocket")
|
||||
go w.readLoop()
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
func (w *WSClient) dial() (*websocket.Conn, error) {
|
||||
// Build the WebSocket URL from the Komari HTTP URL.
|
||||
wsURL := w.komariURL
|
||||
if len(wsURL) > 7 && wsURL[:7] == "http://" {
|
||||
wsURL = "ws://" + wsURL[7:]
|
||||
} else if len(wsURL) > 8 && wsURL[:8] == "https://" {
|
||||
wsURL = "wss://" + wsURL[8:]
|
||||
}
|
||||
wsURL = wsURL + "/api/clients"
|
||||
|
||||
header := http.Header{}
|
||||
dialer := websocket.Dialer{
|
||||
HandshakeTimeout: 10 * time.Second,
|
||||
}
|
||||
|
||||
conn, _, err := dialer.Dial(wsURL, header)
|
||||
return conn, err
|
||||
}
|
||||
|
||||
func (w *WSClient) readLoop() {
|
||||
defer func() {
|
||||
if r := recover(); r != nil {
|
||||
postLog.Error(fmt.Sprintf("WebSocket read loop panic recovered: %v", r))
|
||||
}
|
||||
w.running.Store(false)
|
||||
}()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-w.stopCh:
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
w.connMu.Lock()
|
||||
conn := w.conn
|
||||
w.connMu.Unlock()
|
||||
|
||||
if conn == nil {
|
||||
return
|
||||
}
|
||||
|
||||
_, data, err := conn.ReadMessage()
|
||||
if err != nil {
|
||||
postLog.Warning("Komari WebSocket read error: " + err.Error())
|
||||
w.tryReconnect()
|
||||
return
|
||||
}
|
||||
|
||||
var msg WSMessage
|
||||
if err := json.Unmarshal(data, &msg); err != nil {
|
||||
postLog.Warning("Failed to parse Komari WebSocket message: " + err.Error())
|
||||
continue
|
||||
}
|
||||
|
||||
if msg.Status != "success" {
|
||||
postLog.Warning("Komari WebSocket message with non-success status")
|
||||
continue
|
||||
}
|
||||
|
||||
// Update the node tracker with the received data.
|
||||
tracker := node.GetTracker()
|
||||
if tracker != nil {
|
||||
tracker.UpdateStatus(msg.Data.Online, msg.Data.Data)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (w *WSClient) tryReconnect() {
|
||||
count := w.reconnectCount.Add(1)
|
||||
if int(count) > w.maxRetries {
|
||||
postLog.Error(fmt.Sprintf("Komari WebSocket reconnection failed after %d attempts", w.maxRetries))
|
||||
if w.onReconnectFail != nil {
|
||||
w.onReconnectFail()
|
||||
}
|
||||
// Reset count and start a long-delay retry.
|
||||
w.reconnectCount.Store(0)
|
||||
time.Sleep(30 * time.Second)
|
||||
go w.tryReconnect()
|
||||
return
|
||||
}
|
||||
|
||||
// Exponential backoff.
|
||||
delay := time.Duration(1<<(count-1)) * time.Second
|
||||
if delay > 30*time.Second {
|
||||
delay = 30 * time.Second
|
||||
}
|
||||
|
||||
postLog.Info(fmt.Sprintf("Attempting Komari WebSocket reconnect %d/%d in %v...",
|
||||
count, w.maxRetries, delay))
|
||||
|
||||
time.Sleep(delay)
|
||||
|
||||
conn, err := w.dial()
|
||||
if err != nil {
|
||||
postLog.Warning(fmt.Sprintf("Komari WebSocket reconnect attempt %d failed: %v", count, err))
|
||||
go w.tryReconnect()
|
||||
return
|
||||
}
|
||||
|
||||
w.connMu.Lock()
|
||||
if w.conn != nil {
|
||||
w.conn.Close()
|
||||
}
|
||||
w.conn = conn
|
||||
w.connMu.Unlock()
|
||||
|
||||
w.reconnectCount.Store(0)
|
||||
|
||||
// Request full state after reconnect.
|
||||
if err := conn.WriteMessage(websocket.TextMessage, []byte("get")); err != nil {
|
||||
postLog.Error("Failed to send 'get' after reconnect: " + err.Error())
|
||||
return
|
||||
}
|
||||
|
||||
// Re-fetch node list after reconnect.
|
||||
client := GetClient()
|
||||
if client != nil {
|
||||
nodes, err := client.FetchNodes()
|
||||
if err != nil {
|
||||
postLog.Error("Failed to refresh node list after reconnect: " + err.Error())
|
||||
} else {
|
||||
tracker := node.GetTracker()
|
||||
if tracker != nil {
|
||||
nodeNames := make(map[string]string, len(nodes))
|
||||
for _, n := range nodes {
|
||||
nodeNames[n.UUID] = n.Name
|
||||
}
|
||||
tracker.UpdateNodeList(nodeNames)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
postLog.Info("Komari WebSocket reconnected successfully")
|
||||
go w.readLoop()
|
||||
}
|
||||
|
||||
// Stop closes the WebSocket connection and stops the read loop.
|
||||
func (w *WSClient) Stop() {
|
||||
w.running.Store(false)
|
||||
close(w.stopCh)
|
||||
|
||||
w.connMu.Lock()
|
||||
defer w.connMu.Unlock()
|
||||
if w.conn != nil {
|
||||
w.conn.Close()
|
||||
w.conn = nil
|
||||
}
|
||||
}
|
||||
|
||||
// IsRunning returns whether the WebSocket client is actively running.
|
||||
func (w *WSClient) IsRunning() bool {
|
||||
return w.running.Load()
|
||||
}
|
||||
|
||||
// --- Global Komari client management ---
|
||||
|
||||
var globalClient *Client
|
||||
var globalWSClient *WSClient
|
||||
var globalClientMu sync.Mutex
|
||||
|
||||
// InitClient initializes the global Komari HTTP client.
|
||||
func InitClient(komariURL string) {
|
||||
globalClientMu.Lock()
|
||||
defer globalClientMu.Unlock()
|
||||
globalClient = NewClient(komariURL)
|
||||
}
|
||||
|
||||
// GetClient returns the global Komari HTTP client.
|
||||
func GetClient() *Client {
|
||||
return globalClient
|
||||
}
|
||||
|
||||
// InitWSClient initializes the global Komari WebSocket client.
|
||||
func InitWSClient(komariURL string, maxRetries int) {
|
||||
globalClientMu.Lock()
|
||||
defer globalClientMu.Unlock()
|
||||
globalWSClient = NewWSClient(komariURL, maxRetries)
|
||||
}
|
||||
|
||||
// GetWSClient returns the global Komari WebSocket client.
|
||||
func GetWSClient() *WSClient {
|
||||
return globalWSClient
|
||||
}
|
||||
|
||||
// LoginAndStart performs the Komari login and returns an error if it fails.
|
||||
func LoginAndStart() error {
|
||||
cfg := config.GetConfig()
|
||||
client := GetClient()
|
||||
if client == nil {
|
||||
return fmt.Errorf("komari client not initialized")
|
||||
}
|
||||
|
||||
if err := client.Login(cfg.Komari.Account.Username, cfg.Komari.Account.Password); err != nil {
|
||||
return fmt.Errorf("komari login failed: %w", err)
|
||||
}
|
||||
|
||||
// Fetch initial node list.
|
||||
nodes, err := client.FetchNodes()
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to fetch nodes from komari: %w", err)
|
||||
}
|
||||
|
||||
tracker := node.GetTracker()
|
||||
if tracker != nil {
|
||||
nodeNames := make(map[string]string, len(nodes))
|
||||
for _, n := range nodes {
|
||||
nodeNames[n.UUID] = n.Name
|
||||
}
|
||||
tracker.UpdateNodeList(nodeNames)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
Reference in New Issue
Block a user