feat: add remote worker system for distributed translation
- Add remote translator worker (Docker) that connects via WebSocket - Add backend Go handlers for remote worker CRUD (REST + WebSocket) - Add frontend Admin panel for managing remote workers - Add ProtectedRoute component to secure admin routes - Add PostgreSQL schema for remote_workers table - Fix duplicate Load function in config.go - Fix frontend duplicate fetchStatus function - Fix Dockerfiles missing go.sum (use go mod download + tidy) Features: - Workers connect to server via WebSocket (no ports needed) - API key authentication per worker - Job assigner automatically distributes translation jobs - Worker status tracking (online/offline/disabled) - Timeout handling for stuck jobs
This commit is contained in:
parent
b8e48f3199
commit
8093cc1e40
21 changed files with 2065 additions and 58 deletions
187
backend/internal/handlers/remote_worker.go
Normal file
187
backend/internal/handlers/remote_worker.go
Normal file
|
|
@ -0,0 +1,187 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"crypto/rand"
|
||||
"encoding/hex"
|
||||
"log"
|
||||
"net/http"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/rss2/backend/internal/db"
|
||||
"github.com/rss2/backend/internal/models"
|
||||
)
|
||||
|
||||
func CreateRemoteWorker(c *gin.Context) {
|
||||
var req struct {
|
||||
Name string `json:"name" binding:"required"`
|
||||
Capabilities string `json:"capabilities"`
|
||||
}
|
||||
|
||||
if err := c.ShouldBindJSON(&req); err != nil {
|
||||
c.JSON(http.StatusBadRequest, gin.H{"error": "Invalid request", "message": err.Error()})
|
||||
return
|
||||
}
|
||||
|
||||
if req.Capabilities == "" {
|
||||
req.Capabilities = "cpu"
|
||||
}
|
||||
|
||||
apiKey := generateAPIKey()
|
||||
|
||||
ctx := c.Request.Context()
|
||||
var workerID int
|
||||
err := db.GetPool().QueryRow(ctx, `
|
||||
INSERT INTO remote_workers (name, api_key, capabilities, status, created_at)
|
||||
VALUES ($1, $2, $3, 'offline', NOW())
|
||||
RETURNING id
|
||||
`, req.Name, apiKey, req.Capabilities).Scan(&workerID)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("Error creating remote worker: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to create worker"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusCreated, gin.H{
|
||||
"id": workerID,
|
||||
"name": req.Name,
|
||||
"api_key": apiKey,
|
||||
"capabilities": req.Capabilities,
|
||||
"status": "offline",
|
||||
})
|
||||
}
|
||||
|
||||
func ListRemoteWorkers(c *gin.Context) {
|
||||
ctx := c.Request.Context()
|
||||
|
||||
rows, err := db.GetPool().Query(ctx, `
|
||||
SELECT id, name, api_key, capabilities, status, last_seen, created_at
|
||||
FROM remote_workers
|
||||
ORDER BY created_at DESC
|
||||
`)
|
||||
if err != nil {
|
||||
log.Printf("Error listing remote workers: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to list workers"})
|
||||
return
|
||||
}
|
||||
defer rows.Close()
|
||||
|
||||
var workers []models.RemoteWorker
|
||||
for rows.Next() {
|
||||
var w models.RemoteWorker
|
||||
if err := rows.Scan(&w.ID, &w.Name, &w.APIKey, &w.Capabilities, &w.Status, &w.LastSeen, &w.CreatedAt); err != nil {
|
||||
log.Printf("Error scanning worker: %v", err)
|
||||
continue
|
||||
}
|
||||
workers = append(workers, w)
|
||||
}
|
||||
|
||||
if workers == nil {
|
||||
workers = []models.RemoteWorker{}
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, workers)
|
||||
}
|
||||
|
||||
func GetRemoteWorker(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
var w models.RemoteWorker
|
||||
err := db.GetPool().QueryRow(ctx, `
|
||||
SELECT id, name, api_key, capabilities, status, last_seen, created_at
|
||||
FROM remote_workers
|
||||
WHERE id = $1
|
||||
`, id).Scan(&w.ID, &w.Name, &w.APIKey, &w.Capabilities, &w.Status, &w.LastSeen, &w.CreatedAt)
|
||||
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Worker not found"})
|
||||
return
|
||||
}
|
||||
|
||||
var stats struct {
|
||||
JobsCompleted int `json:"jobs_completed"`
|
||||
JobsPending int `json:"jobs_pending"`
|
||||
}
|
||||
|
||||
db.GetPool().QueryRow(ctx, `
|
||||
SELECT COUNT(*) FROM traducciones WHERE worker_id = $1 AND status = 'done'
|
||||
`, id).Scan(&stats.JobsCompleted)
|
||||
|
||||
db.GetPool().QueryRow(ctx, `
|
||||
SELECT COUNT(*) FROM traducciones WHERE worker_id = $1 AND status = 'pending'
|
||||
`, id).Scan(&stats.JobsPending)
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{
|
||||
"worker": w,
|
||||
"stats": stats,
|
||||
})
|
||||
}
|
||||
|
||||
func DeleteRemoteWorker(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
result, err := db.GetPool().Exec(ctx, `
|
||||
DELETE FROM remote_workers WHERE id = $1
|
||||
`, id)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("Error deleting remote worker: %v", err)
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to delete worker"})
|
||||
return
|
||||
}
|
||||
|
||||
if result.RowsAffected() == 0 {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Worker not found"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"message": "Worker deleted"})
|
||||
}
|
||||
|
||||
func ToggleRemoteWorker(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
var currentStatus string
|
||||
err := db.GetPool().QueryRow(ctx, `SELECT status FROM remote_workers WHERE id = $1`, id).Scan(¤tStatus)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusNotFound, gin.H{"error": "Worker not found"})
|
||||
return
|
||||
}
|
||||
|
||||
newStatus := "disabled"
|
||||
if currentStatus == "disabled" {
|
||||
newStatus = "offline"
|
||||
}
|
||||
|
||||
_, err = db.GetPool().Exec(ctx, `UPDATE remote_workers SET status = $1 WHERE id = $2`, newStatus, id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to update worker status"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"status": newStatus})
|
||||
}
|
||||
|
||||
func RegenerateAPIKey(c *gin.Context) {
|
||||
id := c.Param("id")
|
||||
ctx := c.Request.Context()
|
||||
|
||||
newAPIKey := generateAPIKey()
|
||||
|
||||
_, err := db.GetPool().Exec(ctx, `UPDATE remote_workers SET api_key = $1 WHERE id = $2`, newAPIKey, id)
|
||||
if err != nil {
|
||||
c.JSON(http.StatusInternalServerError, gin.H{"error": "Failed to regenerate API key"})
|
||||
return
|
||||
}
|
||||
|
||||
c.JSON(http.StatusOK, gin.H{"api_key": newAPIKey})
|
||||
}
|
||||
|
||||
func generateAPIKey() string {
|
||||
bytes := make([]byte, 32)
|
||||
rand.Read(bytes)
|
||||
return hex.EncodeToString(bytes)
|
||||
}
|
||||
297
backend/internal/handlers/worker_ws.go
Normal file
297
backend/internal/handlers/worker_ws.go
Normal file
|
|
@ -0,0 +1,297 @@
|
|||
package handlers
|
||||
|
||||
import (
|
||||
"context"
|
||||
"encoding/json"
|
||||
"log"
|
||||
"net/http"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gin-gonic/gin"
|
||||
"github.com/gorilla/websocket"
|
||||
"github.com/rss2/backend/internal/db"
|
||||
"github.com/rss2/backend/internal/models"
|
||||
)
|
||||
|
||||
var upgrader = websocket.Upgrader{
|
||||
CheckOrigin: func(r *http.Request) bool { return true },
|
||||
}
|
||||
|
||||
type WSWorker struct {
|
||||
conn *websocket.Conn
|
||||
workerID int
|
||||
capabilities string
|
||||
workerName string
|
||||
lastHeartbeat time.Time
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
var (
|
||||
workers = make(map[int]*WSWorker)
|
||||
workersByAPI = make(map[string]*WSWorker)
|
||||
workersMu sync.RWMutex
|
||||
)
|
||||
|
||||
func HandleWorkerWS(c *gin.Context) {
|
||||
apiKey := c.Query("api_key")
|
||||
if apiKey == "" {
|
||||
apiKey = c.GetHeader("X-API-Key")
|
||||
}
|
||||
|
||||
if apiKey == "" {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "API key required"})
|
||||
return
|
||||
}
|
||||
|
||||
ctx := c.Request.Context()
|
||||
|
||||
var workerID int
|
||||
var capabilities string
|
||||
err := db.GetPool().QueryRow(ctx, `
|
||||
SELECT id, capabilities FROM remote_workers
|
||||
WHERE api_key = $1 AND status != 'disabled'
|
||||
`, apiKey).Scan(&workerID, &capabilities)
|
||||
|
||||
if err != nil {
|
||||
c.JSON(http.StatusUnauthorized, gin.H{"error": "Invalid API key"})
|
||||
return
|
||||
}
|
||||
|
||||
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
|
||||
if err != nil {
|
||||
log.Printf("WebSocket upgrade error: %v", err)
|
||||
return
|
||||
}
|
||||
|
||||
wsWorker := &WSWorker{
|
||||
conn: conn,
|
||||
workerID: workerID,
|
||||
capabilities: capabilities,
|
||||
lastHeartbeat: time.Now(),
|
||||
}
|
||||
|
||||
workersMu.Lock()
|
||||
workers[workerID] = wsWorker
|
||||
workersByAPI[apiKey] = wsWorker
|
||||
workersMu.Unlock()
|
||||
|
||||
db.GetPool().Exec(ctx, `
|
||||
UPDATE remote_workers SET status = 'online', last_seen = NOW() WHERE id = $1
|
||||
`, workerID)
|
||||
|
||||
go heartbeatLoop(wsWorker)
|
||||
go writeLoop(wsWorker)
|
||||
readLoop(wsWorker, apiKey)
|
||||
}
|
||||
|
||||
func readLoop(wsWorker *WSWorker, apiKey string) {
|
||||
defer func() {
|
||||
cleanupWorker(wsWorker, apiKey)
|
||||
}()
|
||||
|
||||
for {
|
||||
var msg models.WSClientMessage
|
||||
err := wsWorker.conn.ReadJSON(&msg)
|
||||
if err != nil {
|
||||
log.Printf("Worker %d read error: %v", wsWorker.workerID, err)
|
||||
break
|
||||
}
|
||||
|
||||
switch msg.Type {
|
||||
case "register":
|
||||
wsWorker.mu.Lock()
|
||||
if msg.WorkerName != "" {
|
||||
wsWorker.workerName = msg.WorkerName
|
||||
}
|
||||
if msg.Capabilities != "" {
|
||||
wsWorker.capabilities = msg.Capabilities
|
||||
}
|
||||
wsWorker.mu.Unlock()
|
||||
|
||||
sendWS(wsWorker, models.WSServerMessage{Type: "ack", Job: nil})
|
||||
|
||||
case "heartbeat":
|
||||
wsWorker.mu.Lock()
|
||||
wsWorker.lastHeartbeat = time.Now()
|
||||
wsWorker.mu.Unlock()
|
||||
sendWS(wsWorker, models.WSServerMessage{Type: "ack", Job: nil})
|
||||
|
||||
case "result":
|
||||
if msg.Result != nil {
|
||||
handleTranslationResult(wsWorker.workerID, msg.Result)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func writeLoop(wsWorker *WSWorker) {
|
||||
ticker := time.NewTicker(30 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
if err := wsWorker.conn.WriteMessage(websocket.PingMessage, nil); err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func heartbeatLoop(wsWorker *WSWorker) {
|
||||
ticker := time.NewTicker(10 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
wsWorker.mu.Lock()
|
||||
elapsed := time.Since(wsWorker.lastHeartbeat)
|
||||
wsWorker.mu.Unlock()
|
||||
|
||||
if elapsed > 60*time.Second {
|
||||
log.Printf("Worker %d heartbeat timeout", wsWorker.workerID)
|
||||
wsWorker.conn.Close()
|
||||
return
|
||||
}
|
||||
|
||||
sendWS(wsWorker, models.WSServerMessage{Type: "ping", Job: nil})
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func sendWS(wsWorker *WSWorker, msg models.WSServerMessage) {
|
||||
data, _ := json.Marshal(msg)
|
||||
wsWorker.conn.WriteMessage(websocket.TextMessage, data)
|
||||
}
|
||||
|
||||
func cleanupWorker(wsWorker *WSWorker, apiKey string) {
|
||||
workersMu.Lock()
|
||||
delete(workers, wsWorker.workerID)
|
||||
delete(workersByAPI, apiKey)
|
||||
workersMu.Unlock()
|
||||
|
||||
ctx := context.Background()
|
||||
db.GetPool().Exec(ctx, `
|
||||
UPDATE remote_workers SET status = 'offline', last_seen = NOW() WHERE id = $1
|
||||
`, wsWorker.workerID)
|
||||
|
||||
wsWorker.conn.Close()
|
||||
|
||||
db.GetPool().Exec(ctx, `
|
||||
UPDATE traducciones
|
||||
SET status = 'pending', worker_id = NULL, assigned_at = NULL
|
||||
WHERE worker_id = $1 AND status = 'assigned'
|
||||
`, wsWorker.workerID)
|
||||
|
||||
log.Printf("Worker %d disconnected", wsWorker.workerID)
|
||||
}
|
||||
|
||||
func handleTranslationResult(workerID int, result *models.TranslationResult) {
|
||||
ctx := context.Background()
|
||||
|
||||
if result.Error != "" {
|
||||
db.GetPool().Exec(ctx, `
|
||||
UPDATE traducciones
|
||||
SET status = 'error', worker_id = NULL, assigned_at = NULL
|
||||
WHERE id = $1
|
||||
`, result.JobID)
|
||||
log.Printf("Job %d failed: %s", result.JobID, result.Error)
|
||||
return
|
||||
}
|
||||
|
||||
_, err := db.GetPool().Exec(ctx, `
|
||||
UPDATE traducciones
|
||||
SET titulo_trad = $1, resumen_trad = $2, status = 'done',
|
||||
worker_id = $3, assigned_at = NULL
|
||||
WHERE id = $4
|
||||
`, result.TitleTr, result.SummaryTr, workerID, result.JobID)
|
||||
|
||||
if err != nil {
|
||||
log.Printf("Error updating translation result: %v", err)
|
||||
}
|
||||
|
||||
db.GetPool().Exec(ctx, `
|
||||
UPDATE remote_workers SET last_seen = NOW() WHERE id = $1
|
||||
`, workerID)
|
||||
|
||||
log.Printf("Job %d completed by worker %d", result.JobID, workerID)
|
||||
}
|
||||
|
||||
func AssignJobToWorker(workerID int) *models.TranslationJob {
|
||||
ctx := context.Background()
|
||||
|
||||
workersMu.RLock()
|
||||
wsWorker, exists := workers[workerID]
|
||||
workersMu.RUnlock()
|
||||
|
||||
if !exists {
|
||||
return nil
|
||||
}
|
||||
|
||||
wsWorker.mu.Lock()
|
||||
workerCapabilities := wsWorker.capabilities
|
||||
wsWorker.mu.Unlock()
|
||||
|
||||
_ = workerCapabilities
|
||||
|
||||
var job models.TranslationJob
|
||||
err := db.GetPool().QueryRow(ctx, `
|
||||
SELECT t.id, t.noticia_id, t.lang_from, t.lang_to, n.titulo, n.resumen
|
||||
FROM traducciones t
|
||||
JOIN noticias n ON n.id = t.noticia_id
|
||||
WHERE t.status = 'pending'
|
||||
AND (t.worker_id IS NULL OR t.status = 'assigned' AND t.assigned_at < NOW() - INTERVAL '5 minutes')
|
||||
AND t.lang_to = 'es'
|
||||
AND (t.titulo_trad IS NULL OR t.resumen_trad IS NULL)
|
||||
ORDER BY n.fecha DESC
|
||||
LIMIT 1
|
||||
FOR UPDATE SKIP LOCKED
|
||||
`).Scan(&job.ID, &job.NewsID, &job.LangFrom, &job.LangTo, &job.Title, &job.Summary)
|
||||
|
||||
if err != nil {
|
||||
return nil
|
||||
}
|
||||
|
||||
db.GetPool().Exec(ctx, `
|
||||
UPDATE traducciones
|
||||
SET status = 'assigned', worker_id = $1, assigned_at = NOW()
|
||||
WHERE id = $2
|
||||
`, workerID, job.ID)
|
||||
|
||||
sendWS(wsWorker, models.WSServerMessage{
|
||||
Type: "job",
|
||||
Job: &job,
|
||||
})
|
||||
|
||||
return &job
|
||||
}
|
||||
|
||||
func StartJobAssigner() {
|
||||
go func() {
|
||||
ticker := time.NewTicker(5 * time.Second)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-ticker.C:
|
||||
workersMu.RLock()
|
||||
for workerID := range workers {
|
||||
workersMu.RUnlock()
|
||||
AssignJobToWorker(workerID)
|
||||
workersMu.RLock()
|
||||
}
|
||||
workersMu.RUnlock()
|
||||
|
||||
ctx := context.Background()
|
||||
db.GetPool().Exec(ctx, `
|
||||
UPDATE traducciones
|
||||
SET status = 'pending', worker_id = NULL, assigned_at = NULL
|
||||
WHERE status = 'assigned'
|
||||
AND assigned_at < NOW() - INTERVAL '10 minutes'
|
||||
`)
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue