Design a Chat System

WebSocket Endpoints:
wss://chat.example.com/ws?token=JWT_TOKEN

HTTP REST API:
GET /api/v1/chats                    # List user's conversations  
GET /api/v1/chats/{chatId}/messages  # Message history with pagination
POST /api/v1/chats                   # Create new chat/group
POST /api/v1/chats/{chatId}/members  # Add members to group
PUT /api/v1/users/{userId}/presence  # Update online status
POST /api/v1/media/upload            # Upload files/images

{
  "type": "message",
  "chatId": "chat_123",
  "messageId": "msg_abc",
  "senderId": "user_456",
  "content": "Hello world!",
  "timestamp": 1703875200000,
  "messageType": "text"
}

{
  "type": "typing_indicator", 
  "chatId": "chat_123",
  "userId": "user_456",
  "isTyping": true
}

{
  "type": "presence_update",
  "userId": "user_456", 
  "status": "online",
  "lastSeen": 1703875200000
}

{
  "type": "delivery_receipt",
  "messageId": "msg_abc",
  "status": "delivered", // delivered, read
  "userId": "user_789",
  "timestamp": 1703875205000
}

[Client Apps] ←→ [Load Balancer] ←→ [WebSocket Servers] ←→ [Message Queue]
                                           |                      |
                                           ↓                      ↓
[Connection Manager] ←→ [Presence Service] ←→ [Chat Service] ←→ [Message DB]
           |                                      |                |
           ↓                                      ↓                ↓
[Redis Sessions] ←→ [Notification Service] ←→ [File Storage] ←→ [User DB]

@Service
public class WebSocketConnectionManager {
    private final Map<String, Set<WebSocketSession>> userConnections = new ConcurrentHashMap<>();
    private final RedisTemplate<String, String> redis;
    
    public void addConnection(String userId, WebSocketSession session) {
        // Track local connections
        userConnections.computeIfAbsent(userId, k -> ConcurrentHashMap.newKeySet())
                      .add(session);
        
        // Register server instance for this user in Redis
        String serverInstance = InetAddress.getLocalHost().getHostName();
        redis.opsForSet().add("user_connections:" + userId, serverInstance);
        redis.expire("user_connections:" + userId, Duration.ofHours(1));
    }
    
    public void removeConnection(String userId, WebSocketSession session) {
        Set<WebSocketSession> sessions = userConnections.get(userId);
        if (sessions != null) {
            sessions.remove(session);
            if (sessions.isEmpty()) {
                userConnections.remove(userId);
                // Remove server registration if user has no more connections
                String serverInstance = InetAddress.getLocalHost().getHostName();
                redis.opsForSet().remove("user_connections:" + userId, serverInstance);
            }
        }
    }
    
    public Set<String> getServerInstancesForUser(String userId) {
        return redis.opsForSet().members("user_connections:" + userId);
    }
}

class ChatWebSocket {
  constructor(url) {
    this.ws = new WebSocket(url);
    this.pingInterval = null;
    this.missedPings = 0;
    this.maxMissedPings = 3;
    this.setupHeartbeat();
  }
  
  setupHeartbeat() {
    this.pingInterval = setInterval(() => {
      if (this.ws.readyState === WebSocket.OPEN) {
        this.ws.send(JSON.stringify({ type: 'ping' }));
        this.missedPings++;
        
        if (this.missedPings > this.maxMissedPings) {
          console.log('Connection seems dead, reconnecting...');
          this.reconnect();
        }
      }
    }, 30000); // Send ping every 30 seconds
  }
  
  onMessage(event) {
    const message = JSON.parse(event.data);
    
    if (message.type === 'pong') {
      this.missedPings = 0; // Reset missed ping counter
      return;
    }
    
    // Handle other message types...
  }
  
  reconnect() {
    this.ws.close();
    clearInterval(this.pingInterval);
    
    // Exponential backoff for reconnection
    setTimeout(() => {
      this.ws = new WebSocket(this.url);
      this.setupHeartbeat();
      this.setupMessageHandlers();
    }, Math.min(1000 * Math.pow(2, this.reconnectAttempts), 30000));
  }
}

@Component
public class MessageRoutingService {
    private final MessageQueue messageQueue;
    private final WebSocketConnectionManager connectionManager;
    
    public void routeMessage(ChatMessage message) {
        List<String> recipients = getChatParticipants(message.getChatId());
        
        for (String recipient : recipients) {
            if (recipient.equals(message.getSenderId())) {
                continue; // Don't send back to sender
            }
            
            Set<String> serverInstances = connectionManager.getServerInstancesForUser(recipient);
            
            if (serverInstances.isEmpty()) {
                // User is offline, queue for push notification
                enqueueNotification(recipient, message);
            } else {
                // Route to all server instances where user has connections
                for (String serverInstance : serverInstances) {
                    RouteMessage routeMsg = RouteMessage.builder()
                        .serverInstance(serverInstance)
                        .userId(recipient)
                        .message(message)
                        .build();
                    
                    messageQueue.send("websocket_routing", routeMsg);
                }
            }
        }
    }
}

@Entity
public class ChatMessage {
    @Id
    private String messageId;
    private String chatId;
    private String senderId;
    private String content;
    private Long globalSequence;     // Global ordering across all chats
    private Long chatSequence;       // Per-chat ordering
    private Instant timestamp;
    private MessageStatus status;
    
    // getters/setters...
}

@Service
public class MessageSequenceService {
    private final RedisTemplate<String, String> redis;
    
    public Long generateChatSequence(String chatId) {
        // Atomic increment per chat
        return redis.opsForValue().increment("chat_sequence:" + chatId);
    }
    
    public Long generateGlobalSequence() {
        // Global sequence for message ordering across entire system
        return redis.opsForValue().increment("global_message_sequence");
    }
    
    public void saveMessage(ChatMessage message) {
        // Assign sequence numbers before saving
        message.setChatSequence(generateChatSequence(message.getChatId()));
        message.setGlobalSequence(generateGlobalSequence());
        message.setTimestamp(Instant.now());
        
        messageRepository.save(message);
    }
}

class MessageBuffer {
  constructor() {
    this.messages = new Map(); // messageId -> message
    this.nextExpectedSequence = 1;
    this.deliveredMessages = [];
  }
  
  addMessage(message) {
    this.messages.set(message.messageId, message);
    this.processBuffer();
  }
  
  processBuffer() {
    while (true) {
      // Find message with next expected sequence number
      const nextMessage = Array.from(this.messages.values())
        .find(msg => msg.chatSequence === this.nextExpectedSequence);
      
      if (!nextMessage) {
        break; // Wait for missing message
      }
      
      // Deliver message in order
      this.deliveredMessages.push(nextMessage);
      this.messages.delete(nextMessage.messageId);
      this.nextExpectedSequence++;
      
      this.renderMessage(nextMessage);
    }
  }
  
  handleMissingMessages() {
    // Request missing messages from server if gap detected
    if (this.messages.size > 0) {
      const minSequence = Math.min(...Array.from(this.messages.values())
        .map(msg => msg.chatSequence));
      
      if (minSequence > this.nextExpectedSequence) {
        this.requestMissingMessages(this.nextExpectedSequence, minSequence - 1);
      }
    }
  }
}

@Entity
public class MessageDeliveryReceipt {
    @Id
    private String id;
    private String messageId;
    private String userId;
    private DeliveryStatus status; // SENT, DELIVERED, READ
    private Instant timestamp;
    
    // getters/setters...
}

@Service
public class DeliveryReceiptService {
    
    public void markAsDelivered(String messageId, String userId) {
        MessageDeliveryReceipt receipt = MessageDeliveryReceipt.builder()
            .messageId(messageId)
            .userId(userId)
            .status(DeliveryStatus.DELIVERED)
            .timestamp(Instant.now())
            .build();
        
        deliveryReceiptRepository.save(receipt);
        
        // Send receipt back to message sender
        broadcastDeliveryReceipt(messageId, userId, DeliveryStatus.DELIVERED);
    }
    
    public void markAsRead(String messageId, String userId) {
        // Update existing receipt or create new one
        MessageDeliveryReceipt receipt = deliveryReceiptRepository
            .findByMessageIdAndUserId(messageId, userId)
            .orElse(new MessageDeliveryReceipt());
        
        receipt.setMessageId(messageId);
        receipt.setUserId(userId);
        receipt.setStatus(DeliveryStatus.READ);
        receipt.setTimestamp(Instant.now());
        
        deliveryReceiptRepository.save(receipt);
        broadcastDeliveryReceipt(messageId, userId, DeliveryStatus.READ);
    }
    
    private void broadcastDeliveryReceipt(String messageId, String userId, 
                                        DeliveryStatus status) {
        // Find original message sender
        ChatMessage originalMessage = messageRepository.findById(messageId);
        if (originalMessage != null) {
            DeliveryReceiptEvent event = DeliveryReceiptEvent.builder()
                .messageId(messageId)
                .userId(userId)
                .status(status)
                .timestamp(Instant.now())
                .build();
            
            // Send receipt to original sender
            webSocketService.sendToUser(originalMessage.getSenderId(), event);
        }
    }
}

@Service
public class OfflineMessageService {
    private final NotificationService notificationService;
    private final MessageRepository messageRepository;
    
    public void handleOfflineDelivery(String userId, ChatMessage message) {
        // Store message for offline user
        OfflineMessage offlineMsg = OfflineMessage.builder()
            .userId(userId)
            .messageId(message.getMessageId())
            .chatId(message.getChatId())
            .timestamp(Instant.now())
            .delivered(false)
            .build();
        
        offlineMessageRepository.save(offlineMsg);
        
        // Send push notification
        PushNotification notification = PushNotification.builder()
            .userId(userId)
            .title(getSenderName(message.getSenderId()))
            .body(truncateMessage(message.getContent()))
            .badge(getUnreadCount(userId))
            .data(Map.of(
                "chatId", message.getChatId(),
                "messageId", message.getMessageId()
            ))
            .build();
        
        notificationService.sendPushNotification(notification);
    }
    
    public List<ChatMessage> deliverOfflineMessages(String userId) {
        List<OfflineMessage> offlineMessages = offlineMessageRepository
            .findByUserIdAndDeliveredFalse(userId);
        
        List<String> messageIds = offlineMessages.stream()
            .map(OfflineMessage::getMessageId)
            .collect(toList());
        
        List<ChatMessage> messages = messageRepository.findByIdIn(messageIds);
        
        // Mark as delivered
        offlineMessages.forEach(msg -> {
            msg.setDelivered(true);
            msg.setDeliveredAt(Instant.now());
        });
        offlineMessageRepository.saveAll(offlineMessages);
        
        return messages;
    }
}

CREATE TABLE chat_messages (
    message_id VARCHAR(36) PRIMARY KEY,
    chat_id VARCHAR(36) NOT NULL,
    sender_id VARCHAR(36) NOT NULL,
    content TEXT NOT NULL,
    message_type ENUM('text', 'image', 'file', 'system') DEFAULT 'text',
    chat_sequence BIGINT NOT NULL,
    global_sequence BIGINT NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP,
    
    INDEX idx_chat_sequence (chat_id, chat_sequence),
    INDEX idx_global_sequence (global_sequence),
    INDEX idx_sender_time (sender_id, created_at)
);

CREATE TABLE chats (
    chat_id VARCHAR(36) PRIMARY KEY,
    chat_type ENUM('direct', 'group') NOT NULL,
    name VARCHAR(255),
    description TEXT,
    created_by VARCHAR(36) NOT NULL,
    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    last_message_at TIMESTAMP,
    last_message_id VARCHAR(36),
    
    INDEX idx_created_by (created_by),
    INDEX idx_last_message (last_message_at)
);

CREATE TABLE chat_participants (
    chat_id VARCHAR(36),
    user_id VARCHAR(36),
    joined_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
    role ENUM('admin', 'member') DEFAULT 'member',
    last_read_sequence BIGINT DEFAULT 0,
    
    PRIMARY KEY (chat_id, user_id),
    INDEX idx_user_chats (user_id, joined_at)
);

@Service
public class ChatShardingService {
    private static final int SHARD_COUNT = 16;
    
    public int getShardForChat(String chatId) {
        // Use consistent hashing based on chat ID
        return Math.abs(chatId.hashCode()) % SHARD_COUNT;
    }
    
    public int getShardForUser(String userId) {
        // User data can be sharded separately from messages
        return Math.abs(userId.hashCode()) % SHARD_COUNT;
    }
    
    public String getMessageTableName(String chatId) {
        int shard = getShardForChat(chatId);
        return "chat_messages_shard_" + shard;
    }
    
    // For cross-shard queries (user's chat list), use a mapping table
    public List<ChatSummary> getUserChats(String userId) {
        // Query user's chat participation across all shards
        List<ChatParticipant> participations = chatParticipantRepository
            .findByUserId(userId);
        
        // Group by shard to minimize database connections
        Map<Integer, List<String>> chatsByShard = participations.stream()
            .collect(groupingBy(
                p -> getShardForChat(p.getChatId()),
                mapping(ChatParticipant::getChatId, toList())
            ));
        
        List<ChatSummary> allChats = new ArrayList<>();
        for (Map.Entry<Integer, List<String>> entry : chatsByShard.entrySet()) {
            DataSource shard = getShardDataSource(entry.getKey());
            List<ChatSummary> shardChats = queryChatSummaries(shard, entry.getValue());
            allChats.addAll(shardChats);
        }
        
        return allChats.stream()
            .sorted((a, b) -> b.getLastMessageAt().compareTo(a.getLastMessageAt()))
            .collect(toList());
    }
}

@RestController
public class ChatHistoryController {
    
    @GetMapping("/api/v1/chats/{chatId}/messages")
    public MessageHistoryResponse getMessages(
            @PathVariable String chatId,
            @RequestParam(required = false) String cursor,
            @RequestParam(defaultValue = "50") int limit) {
        
        // Validate user access to chat
        validateUserAccess(chatId, getCurrentUserId());
        
        Long sequenceStart = null;
        if (cursor != null) {
            sequenceStart = decodeCursor(cursor);
        }
        
        List<ChatMessage> messages = messageService.getChatHistory(
            chatId, sequenceStart, limit + 1); // Fetch one extra for hasMore
        
        boolean hasMore = messages.size() > limit;
        if (hasMore) {
            messages = messages.subList(0, limit);
        }
        
        String nextCursor = null;
        if (hasMore && !messages.isEmpty()) {
            ChatMessage lastMessage = messages.get(messages.size() - 1);
            nextCursor = encodeCursor(lastMessage.getChatSequence());
        }
        
        return MessageHistoryResponse.builder()
            .messages(messages)
            .hasMore(hasMore)
            .nextCursor(nextCursor)
            .build();
    }
    
    private String encodeCursor(Long sequence) {
        return Base64.getEncoder().encodeToString(sequence.toString().getBytes());
    }
    
    private Long decodeCursor(String cursor) {
        byte[] decoded = Base64.getDecoder().decode(cursor);
        return Long.parseLong(new String(decoded));
    }
}

@Service
public class TypingIndicatorService {
    private final RedisTemplate<String, String> redis;
    private static final Duration TYPING_TIMEOUT = Duration.ofSeconds(10);
    
    public void startTyping(String chatId, String userId) {
        String key = "typing:" + chatId;
        
        // Add user to typing set with expiration
        redis.opsForSet().add(key, userId);
        redis.expire(key, TYPING_TIMEOUT);
        
        // Broadcast typing indicator to chat participants
        TypingIndicatorEvent event = TypingIndicatorEvent.builder()
            .chatId(chatId)
            .userId(userId)
            .isTyping(true)
            .build();
        
        chatEventService.broadcastToChatParticipants(chatId, event, userId);
    }
    
    public void stopTyping(String chatId, String userId) {
        String key = "typing:" + chatId;
        redis.opsForSet().remove(key, userId);
        
        TypingIndicatorEvent event = TypingIndicatorEvent.builder()
            .chatId(chatId)
            .userId(userId)
            .isTyping(false)
            .build();
        
        chatEventService.broadcastToChatParticipants(chatId, event, userId);
    }
    
    public Set<String> getTypingUsers(String chatId) {
        String key = "typing:" + chatId;
        return redis.opsForSet().members(key);
    }
}

class TypingManager {
  constructor(chatId, webSocket) {
    this.chatId = chatId;
    this.ws = webSocket;
    this.typingTimer = null;
    this.isCurrentlyTyping = false;
  }
  
  handleKeypress() {
    if (!this.isCurrentlyTyping) {
      this.startTyping();
    }
    
    // Reset the timer on each keypress
    clearTimeout(this.typingTimer);
    this.typingTimer = setTimeout(() => {
      this.stopTyping();
    }, 3000); // Stop typing after 3 seconds of inactivity
  }
  
  startTyping() {
    this.isCurrentlyTyping = true;
    this.ws.send(JSON.stringify({
      type: 'typing_start',
      chatId: this.chatId
    }));
  }
  
  stopTyping() {
    if (this.isCurrentlyTyping) {
      this.isCurrentlyTyping = false;
      this.ws.send(JSON.stringify({
        type: 'typing_stop',
        chatId: this.chatId
      }));
    }
  }
}

@RestController
public class MediaController {
    private final MediaStorageService storageService;
    private final VirusScanService virusScanService;
    
    @PostMapping("/api/v1/media/upload")
    public MediaUploadResponse uploadMedia(@RequestParam("file") MultipartFile file,
                                         @RequestParam("chatId") String chatId) {
        
        // Validate file size and type
        validateFile(file);
        
        // Scan for viruses (async for large files)
        virusScanService.scanFile(file);
        
        // Generate unique file ID and upload to storage
        String fileId = UUID.randomUUID().toString();
        String storageUrl = storageService.uploadFile(file, fileId);
        
        // Create media message
        ChatMessage mediaMessage = ChatMessage.builder()
            .messageId(UUID.randomUUID().toString())
            .chatId(chatId)
            .senderId(getCurrentUserId())
            .messageType(MessageType.MEDIA)
            .content(createMediaContent(file, storageUrl))
            .build();
        
        messageService.sendMessage(mediaMessage);
        
        return MediaUploadResponse.builder()
            .fileId(fileId)
            .url(storageUrl)
            .messageId(mediaMessage.getMessageId())
            .build();
    }
    
    private String createMediaContent(MultipartFile file, String url) {
        MediaContent content = MediaContent.builder()
            .filename(file.getOriginalFilename())
            .mimeType(file.getContentType())
            .size(file.getSize())
            .url(url)
            .build();
        
        return objectMapper.writeValueAsString(content);
    }
}

Key Exchange Flow:
1. Client A generates ephemeral key pair
2. Client A requests Client B's public key from server
3. Clients perform key agreement (ECDH)
4. Messages encrypted with derived symmetric key
5. Server only stores encrypted message blobs