feat: implement live-sync for tasks, workspace note sharing, and document locking system

internal/ws/hub.go

@@ -12,8 +12,8 @@ import (
 )
 
 var upgrader = websocket.Upgrader{
-	ReadBufferSize:  1024,
-	WriteBufferSize: 1024,
+	ReadBufferSize:  2048,
+	WriteBufferSize: 2048,
 	CheckOrigin: func(r *http.Request) bool {
 		return true // Self-hosted reverse proxy compatible
 	},
@@ -27,53 +27,273 @@ type Client struct {
 	send    chan []byte
 }
 
+type docLockItem struct {
+	docID     string
+	userID    string
+	userName  string
+	lockedAt  time.Time
+	expiresAt time.Time
+}
+
 type Hub struct {
-	mu    sync.RWMutex
-	rooms map[string]map[*Client]bool
+	mu         sync.RWMutex
+	spaceRooms map[string]map[*Client]bool
+	userRooms  map[string]map[*Client]bool
+	allClients map[*Client]bool
+	locks      map[string]*docLockItem
 }
 
 func NewHub() *Hub {
-	return &Hub{
-		rooms: make(map[string]map[*Client]bool),
+	h := &Hub{
+		spaceRooms: make(map[string]map[*Client]bool),
+		userRooms:  make(map[string]map[*Client]bool),
+		allClients: make(map[*Client]bool),
+		locks:      make(map[string]*docLockItem),
 	}
+	go h.cleanupLocksLoop()
+	return h
 }
 
 func (h *Hub) Register(c *Client) {
 	h.mu.Lock()
-	if h.rooms[c.spaceID] == nil {
-		h.rooms[c.spaceID] = make(map[*Client]bool)
+	h.allClients[c] = true
+	if c.spaceID != "" {
+		if h.spaceRooms[c.spaceID] == nil {
+			h.spaceRooms[c.spaceID] = make(map[*Client]bool)
+		}
+		h.spaceRooms[c.spaceID][c] = true
 	}
-	h.rooms[c.spaceID][c] = true
-	onlineCount := len(h.rooms[c.spaceID])
+	if c.userID != "" && c.userID != "anonymous" {
+		if h.userRooms[c.userID] == nil {
+			h.userRooms[c.userID] = make(map[*Client]bool)
+		}
+		h.userRooms[c.userID][c] = true
+	}
+	onlineCount := len(h.allClients)
 	h.mu.Unlock()
 
-	h.Broadcast(c.spaceID, "", "USER_PRESENCE", map[string]interface{}{
-		"onlineCount": onlineCount,
-	})
-}
-
-func (h *Hub) Unregister(c *Client) {
-	h.mu.Lock()
-	var onlineCount int
-	if room, ok := h.rooms[c.spaceID]; ok {
-		if _, exists := room[c]; exists {
-			delete(room, c)
-			close(c.send)
-		}
-		onlineCount = len(room)
-		if onlineCount == 0 {
-			delete(h.rooms, c.spaceID)
-		}
-	}
-	h.mu.Unlock()
-
-	if onlineCount > 0 {
+	if c.spaceID != "" {
 		h.Broadcast(c.spaceID, "", "USER_PRESENCE", map[string]interface{}{
 			"onlineCount": onlineCount,
 		})
 	}
 }
 
+func (h *Hub) Unregister(c *Client) {
+	h.mu.Lock()
+	delete(h.allClients, c)
+	if c.spaceID != "" {
+		if room, ok := h.spaceRooms[c.spaceID]; ok {
+			delete(room, c)
+			if len(room) == 0 {
+				delete(h.spaceRooms, c.spaceID)
+			}
+		}
+	}
+	userRemaining := 0
+	if c.userID != "" && c.userID != "anonymous" {
+		if uRoom, ok := h.userRooms[c.userID]; ok {
+			delete(uRoom, c)
+			userRemaining = len(uRoom)
+			if userRemaining == 0 {
+				delete(h.userRooms, c.userID)
+			}
+		}
+	}
+	close(c.send)
+
+	// If this user has no other active client connections, release their document locks
+	var releasedDocIDs []string
+	if userRemaining == 0 && c.userID != "" {
+		now := time.Now()
+		for docID, l := range h.locks {
+			if l.userID == c.userID || l.expiresAt.Before(now) {
+				releasedDocIDs = append(releasedDocIDs, docID)
+				delete(h.locks, docID)
+			}
+		}
+	}
+	h.mu.Unlock()
+
+	// Broadcast unlock events for released locks
+	for _, docID := range releasedDocIDs {
+		h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID})
+	}
+}
+
+// Document Locking Methods
+
+func (h *Hub) AcquireDocLock(docID, userID, userName string, ttl time.Duration) (*domain.DocLockInfo, bool, *domain.DocLockInfo) {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	now := time.Now()
+	existing, exists := h.locks[docID]
+	if exists && existing.expiresAt.After(now) {
+		if existing.userID != userID {
+			// Locked by another user
+			info := &domain.DocLockInfo{
+				DocID:     existing.docID,
+				UserID:    existing.userID,
+				UserName:  existing.userName,
+				ExpiresAt: existing.expiresAt.Format(time.RFC3339),
+			}
+			return nil, false, info
+		}
+		// Same user extending lock
+		existing.expiresAt = now.Add(ttl)
+		if userName != "" {
+			existing.userName = userName
+		}
+		info := &domain.DocLockInfo{
+			DocID:     existing.docID,
+			UserID:    existing.userID,
+			UserName:  existing.userName,
+			ExpiresAt: existing.expiresAt.Format(time.RFC3339),
+		}
+		return info, true, nil
+	}
+
+	// Create fresh lock
+	item := &docLockItem{
+		docID:     docID,
+		userID:    userID,
+		userName:  userName,
+		lockedAt:  now,
+		expiresAt: now.Add(ttl),
+	}
+	h.locks[docID] = item
+	info := &domain.DocLockInfo{
+		DocID:     item.docID,
+		UserID:    item.userID,
+		UserName:  item.userName,
+		ExpiresAt: item.expiresAt.Format(time.RFC3339),
+	}
+	return info, true, nil
+}
+
+func (h *Hub) HeartbeatDocLock(docID, userID string, ttl time.Duration) bool {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	now := time.Now()
+	item, ok := h.locks[docID]
+	if !ok || item.userID != userID || item.expiresAt.Before(now) {
+		return false
+	}
+	item.expiresAt = now.Add(ttl)
+	return true
+}
+
+func (h *Hub) ReleaseDocLock(docID, userID string) bool {
+	h.mu.Lock()
+	defer h.mu.Unlock()
+
+	item, ok := h.locks[docID]
+	if !ok {
+		return false
+	}
+	if item.userID == userID || userID == "admin" {
+		delete(h.locks, docID)
+		return true
+	}
+	return false
+}
+
+func (h *Hub) GetDocLock(docID string) *domain.DocLockInfo {
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	item, ok := h.locks[docID]
+	if !ok || item.expiresAt.Before(time.Now()) {
+		return nil
+	}
+	return &domain.DocLockInfo{
+		DocID:     item.docID,
+		UserID:    item.userID,
+		UserName:  item.userName,
+		ExpiresAt: item.expiresAt.Format(time.RFC3339),
+	}
+}
+
+func (h *Hub) cleanupLocksLoop() {
+	ticker := time.NewTicker(4 * time.Second)
+	defer ticker.Stop()
+
+	for range ticker.C {
+		h.mu.Lock()
+		now := time.Now()
+		var expiredDocIDs []string
+		for docID, l := range h.locks {
+			if l.expiresAt.Before(now) {
+				expiredDocIDs = append(expiredDocIDs, docID)
+				delete(h.locks, docID)
+			}
+		}
+		h.mu.Unlock()
+
+		for _, docID := range expiredDocIDs {
+			h.BroadcastToAll("DOC_UNLOCKED", map[string]string{"docId": docID})
+		}
+	}
+}
+
+// Broadcasting Methods
+
+func (h *Hub) BroadcastToAll(event string, payload interface{}) {
+	msg := domain.WSMessage{
+		Event:     event,
+		Payload:   payload,
+		Timestamp: domain.NowISO(),
+	}
+	data, err := json.Marshal(msg)
+	if err != nil {
+		return
+	}
+
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	for client := range h.allClients {
+		select {
+		case client.send <- data:
+		default:
+		}
+	}
+}
+
+func (h *Hub) BroadcastToUser(userID, event string, payload interface{}) {
+	msg := domain.WSMessage{
+		Event:     event,
+		Payload:   payload,
+		Timestamp: domain.NowISO(),
+	}
+	data, err := json.Marshal(msg)
+	if err != nil {
+		return
+	}
+
+	h.mu.RLock()
+	defer h.mu.RUnlock()
+
+	if room, ok := h.userRooms[userID]; ok {
+		for client := range room {
+			select {
+			case client.send <- data:
+			default:
+			}
+		}
+	}
+}
+
+func (h *Hub) BroadcastTaskEvent(userID, event string, payload interface{}) {
+	h.BroadcastToUser(userID, event, payload)
+}
+
+func (h *Hub) BroadcastDocEvent(docID, event string, payload interface{}) {
+	h.BroadcastToAll(event, payload)
+}
+
 func (h *Hub) Broadcast(spaceID, listID, event string, payload interface{}) {
 	msg := domain.WSMessage{
 		Event:     event,
@@ -90,7 +310,7 @@ func (h *Hub) Broadcast(spaceID, listID, event string, payload interface{}) {
 	h.mu.RLock()
 	defer h.mu.RUnlock()
 
-	room, ok := h.rooms[spaceID]
+	room, ok := h.spaceRooms[spaceID]
 	if !ok {
 		return
 	}
@@ -115,7 +335,7 @@ func (h *Hub) ServeWS(w http.ResponseWriter, r *http.Request, spaceID, userID st
 		conn:    conn,
 		spaceID: spaceID,
 		userID:  userID,
-		send:    make(chan []byte, 64),
+		send:    make(chan []byte, 128),
 	}
 
 	h.Register(client)
@@ -129,7 +349,7 @@ func (c *Client) readPump() {
 		_ = c.conn.Close()
 	}()
 
-	c.conn.SetReadLimit(8192)
+	c.conn.SetReadLimit(16384)
 	_ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
 	c.conn.SetPongHandler(func(string) error {
 		_ = c.conn.SetReadDeadline(time.Now().Add(60 * time.Second))
@@ -137,10 +357,42 @@ func (c *Client) readPump() {
 	})
 
 	for {
-		_, _, err := c.conn.ReadMessage()
+		_, message, err := c.conn.ReadMessage()
 		if err != nil {
 			break
 		}
+
+		// Handle client-initiated lock and unlock messages
+		var incoming struct {
+			Event   string          `json:"event"`
+			DocID   string          `json:"docId"`
+			Payload json.RawMessage `json:"payload"`
+		}
+		if json.Unmarshal(message, &incoming) == nil && c.userID != "" && c.userID != "anonymous" {
+			docID := incoming.DocID
+			if docID == "" && len(incoming.Payload) > 0 {
+				var p struct {
+					DocID string `json:"docId"`
+				}
+				_ = json.Unmarshal(incoming.Payload, &p)
+				docID = p.DocID
+			}
+
+			if docID != "" {
+				switch incoming.Event {
+				case "LOCK_DOC":
+					if info, ok, _ := c.hub.AcquireDocLock(docID, c.userID, "", 30*time.Second); ok && info != nil {
+						c.hub.BroadcastDocEvent(docID, "DOC_LOCKED", info)
+					}
+				case "HEARTBEAT_DOC":
+					c.hub.HeartbeatDocLock(docID, c.userID, 30*time.Second)
+				case "UNLOCK_DOC":
+					if c.hub.ReleaseDocLock(docID, c.userID) {
+						c.hub.BroadcastDocEvent(docID, "DOC_UNLOCKED", map[string]string{"docId": docID})
+					}
+				}
+			}
+		}
 	}
 }