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.
This commit is contained in:
@@ -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()
|
||||
}
|
||||
}()
|
||||
}
|
||||
Reference in New Issue
Block a user