diff --git a/go.mod b/go.mod index 9b11431..b120d18 100644 --- a/go.mod +++ b/go.mod @@ -22,6 +22,7 @@ require ( github.com/gabriel-vasile/mimetype v1.4.12 // indirect github.com/go-playground/locales v0.14.1 // indirect github.com/go-playground/universal-translator v0.18.1 // indirect + github.com/gorilla/websocket v1.5.3 // indirect github.com/jackc/pgpassfile v1.0.0 // indirect github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a // indirect github.com/jackc/puddle/v2 v2.2.1 // indirect diff --git a/go.sum b/go.sum index 73654d5..242f5c7 100644 --- a/go.sum +++ b/go.sum @@ -24,6 +24,8 @@ github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI= github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY= github.com/google/uuid v1.4.0 h1:MtMxsa51/r9yyhkyLsVeVt0B+BGQZzpQiTQ4eHZ8bc4= github.com/google/uuid v1.4.0/go.mod h1:TIyPZe4MgqvfeYDBFedMoGGpEw/LqOeaOT+nhxU+yHo= +github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg= +github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE= github.com/jackc/pgpassfile v1.0.0 h1:/6Hmqy13Ss2zCq62VdNG8tM1wchn8zjSGOBJ6icpsIM= github.com/jackc/pgpassfile v1.0.0/go.mod h1:CEx0iS5ambNFdcRtxPj5JhEz+xB6uRky5eyVu/W2HEg= github.com/jackc/pgservicefile v0.0.0-20221227161230-091c0ba34f0a h1:bbPeKD0xmW/Y25WS6cokEszi5g+S0QxI/d45PkRi7Nk= diff --git a/internal/sync/websocket.go b/internal/sync/websocket.go new file mode 100644 index 0000000..2451b84 --- /dev/null +++ b/internal/sync/websocket.go @@ -0,0 +1,207 @@ +package sync + +import ( + "log" + "sync" + "time" + + "github.com/google/uuid" +) + +// Message types for WebSocket communication +const ( + MessageTypeProgressUpdate = "progress_update" + MessageTypeAnnotationUpdate = "annotation_update" + MessageTypeConflict = "conflict" + MessageTypeSyncComplete = "sync_complete" + MessageTypeHeartbeat = "heartbeat" + MessageTypeInitial = "initial_state" +) + +// BroadcastMessage represents a message to broadcast to connected clients +type BroadcastMessage struct { + Type string `json:"type"` + Timestamp string `json:"timestamp"` + Data interface{} `json:"data"` + Source *SourceDevice `json:"source_device,omitempty"` +} + +// SourceDevice represents the device that sent the update +type SourceDevice struct { + ID string `json:"id"` + Name string `json:"name"` + Type string `json:"type"` +} + +// DeviceConnection represents an active WebSocket connection +type DeviceConnection struct { + ID string + UserID string + DeviceType string + DeviceName string + Connected time.Time + LastPing time.Time + Send chan BroadcastMessage + Disconnected chan struct{} +} + +// ConnectionManager manages WebSocket connections +type ConnectionManager struct { + connections map[string]*DeviceConnection + broadcast chan BroadcastMessage + mu sync.RWMutex +} + +// NewConnectionManager creates a new connection manager +func NewConnectionManager() *ConnectionManager { + cm := &ConnectionManager{ + connections: make(map[string]*DeviceConnection), + broadcast: make(chan BroadcastMessage, 100), + } + + // Start broadcast goroutine + go cm.broadcastLoop() + + return cm +} + +// broadcastLoop handles broadcasting messages to all connected clients +func (m *ConnectionManager) broadcastLoop() { + for { + select { + case msg := <-m.broadcast: + m.mu.RLock() + for _, conn := range m.connections { + select { + case conn.Send <- msg: + default: + // Channel full, skip this connection + } + } + m.mu.RUnlock() + } + } +} + +// Broadcast sends a message to all connected clients +func (m *ConnectionManager) Broadcast(msg BroadcastMessage) { + select { + case m.broadcast <- msg: + default: + log.Println("Broadcast channel full, message dropped") + } +} + +// BroadcastProgressUpdate broadcasts a progress update to all connected clients +func (m *ConnectionManager) BroadcastProgressUpdate(bookID uuid.UUID, percentage float64, source SourceDevice) { + msg := BroadcastMessage{ + Type: MessageTypeProgressUpdate, + Timestamp: time.Now().Format(time.RFC3339), + Data: map[string]interface{}{ + "book_id": bookID.String(), + "percentage": percentage, + }, + Source: &source, + } + m.Broadcast(msg) +} + +// BroadcastAnnotationUpdate broadcasts an annotation update to all connected clients +func (m *ConnectionManager) BroadcastAnnotationUpdate(bookID uuid.UUID, annotationType string, data interface{}, source SourceDevice) { + msg := BroadcastMessage{ + Type: MessageTypeAnnotationUpdate, + Timestamp: time.Now().Format(time.RFC3339), + Data: map[string]interface{}{ + "book_id": bookID.String(), + "annotation_type": annotationType, + "data": data, + }, + Source: &source, + } + m.Broadcast(msg) +} + +// AddConnection adds a new WebSocket connection +func (m *ConnectionManager) AddConnection(conn *DeviceConnection) { + m.mu.Lock() + defer m.mu.Unlock() + m.connections[conn.ID] = conn + log.Printf("WebSocket: Connection added for device %s (user: %s, type: %s)", conn.DeviceName, conn.UserID, conn.DeviceType) +} + +// RemoveConnection removes a WebSocket connection +func (m *ConnectionManager) RemoveConnection(connID string) { + m.mu.Lock() + defer m.mu.Unlock() + if conn, exists := m.connections[connID]; exists { + close(conn.Disconnected) + delete(m.connections, connID) + log.Printf("WebSocket: Connection removed for device %s", conn.DeviceName) + } +} + +// GetConnection retrieves a connection by ID +func (m *ConnectionManager) GetConnection(connID string) (*DeviceConnection, bool) { + m.mu.RLock() + defer m.mu.RUnlock() + conn, exists := m.connections[connID] + return conn, exists +} + +// GetUserConnections returns all connections for a specific user +func (m *ConnectionManager) GetUserConnections(userID string) []*DeviceConnection { + m.mu.RLock() + defer m.mu.RUnlock() + + var conns []*DeviceConnection + for _, conn := range m.connections { + if conn.UserID == userID { + conns = append(conns, conn) + } + } + return conns +} + +// GetConnectionCount returns the number of active connections +func (m *ConnectionManager) GetConnectionCount() int { + m.mu.RLock() + defer m.mu.RUnlock() + return len(m.connections) +} + +// CleanupStaleConnections removes connections that haven't sent a ping in 2 minutes +func (m *ConnectionManager) CleanupStaleConnections() { + m.mu.Lock() + defer m.mu.Unlock() + + staleThreshold := time.Now().Add(-2 * time.Minute) + for id, conn := range m.connections { + if conn.LastPing.Before(staleThreshold) { + log.Printf("WebSocket: Cleaning up stale connection for %s", conn.DeviceName) + close(conn.Disconnected) + delete(m.connections, id) + } + } +} + +// GetConnectionStats returns statistics about connections +func (m *ConnectionManager) GetConnectionStats() map[string]int { + m.mu.RLock() + defer m.mu.RUnlock() + + stats := make(map[string]int) + for _, conn := range m.connections { + stats[conn.DeviceType]++ + } + return stats +} + +// StartCleanupTask starts a background task to cleanup stale connections +func (m *ConnectionManager) StartCleanupTask() { + ticker := time.NewTicker(1 * time.Minute) + go func() { + for range ticker.C { + m.CleanupStaleConnections() + } + }() +}