feat: add WebSocket connection manager infrastructure

Add ConnectionManager for real-time WebSocket communication: - Message types for progress updates, annotations, conflicts - Broadcast message structure with source device tracking - Device connection tracking with user and device metadata - Automatic broadcast loop with concurrent message delivery - Connection management (add, remove, get by ID/user) - Stale connection cleanup (2-minute timeout) - Connection statistics by device type - Background cleanup task runs every minute This implements the core WebSocket infrastructure needed for Week 9 of the Universal Sync Implementation Guide.
2026-01-30 21:46:53 -05:00
parent 25057cf33a
commit 9bfe14bb38
3 changed files with 210 additions and 0 deletions
@@ -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
@@ -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=
@@ -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()
+		}
+	}()
+}