System Design: Chat Application

// WebSocket connection management
import { WebSocketServer, WebSocket } from "ws";

interface ConnectedUser {
  userId: string;
  socket: WebSocket;
  serverId: string; // Which chat server this user is connected to
}

class ChatWebSocketServer {
  private connections = new Map<string, WebSocket>();
  private wss: WebSocketServer;

  constructor(port: number) {
    this.wss = new WebSocketServer({ port });
    this.wss.on("connection", this.handleConnection.bind(this));
  }

  private handleConnection(socket: WebSocket, request: any) {
    const userId = this.authenticateUser(request);
    if (!userId) {
      socket.close(4001, "Unauthorized");
      return;
    }

    // Register connection
    this.connections.set(userId, socket);
    this.updatePresence(userId, "online");

    socket.on("message", (data) => this.handleMessage(userId, data));

    socket.on("close", () => {
      this.connections.delete(userId);
      this.updatePresence(userId, "offline");
    });

    // Heartbeat to detect stale connections
    socket.on("pong", () => { /* connection is alive */ });
  }

  async sendToUser(targetUserId: string, message: any): Promise<boolean> {
    const socket = this.connections.get(targetUserId);
    if (socket && socket.readyState === WebSocket.OPEN) {
      socket.send(JSON.stringify(message));
      return true; // Delivered
    }
    return false; // User not on this server
  }
}

// Message flow: User A sends a message to User B

interface ChatMessage {
  id: string;            // Globally unique ID (Snowflake)
  conversationId: string;
  senderId: string;
  content: string;
  contentType: "text" | "image" | "file";
  timestamp: number;
  status: "sent" | "delivered" | "read";
}

class MessageRouter {
  private redis: RedisClient;
  private kafka: KafkaProducer;
  private pushService: PushNotificationService;

  async routeMessage(message: ChatMessage, recipientId: string): Promise<void> {
    // Step 1: Persist the message
    await this.kafka.publish("messages", {
      key: message.conversationId, // Partition by conversation for ordering
      value: message,
    });

    // Step 2: Find which server the recipient is connected to
    const serverInfo = await this.redis.get(`conn:${recipientId}`);

    if (serverInfo) {
      // User is online - route to their chat server
      const { serverId } = JSON.parse(serverInfo);
      await this.forwardToServer(serverId, recipientId, message);
    } else {
      // User is offline - send push notification
      await this.pushService.send(recipientId, {
        title: `New message from ${message.senderId}`,
        body: message.content.substring(0, 100),
        data: { conversationId: message.conversationId },
      });
    }
  }

  private async forwardToServer(
    serverId: string,
    recipientId: string,
    message: ChatMessage
  ): Promise<void> {
    // Use internal RPC or pub/sub to reach the right server
    await this.redis.publish(`server:${serverId}`, JSON.stringify({
      type: "deliver",
      recipientId,
      message,
    }));
  }
}

// Message storage schema (Cassandra)
// Partition key: conversation_id
// Clustering key: message_id (Snowflake, so it's time-ordered)

// CREATE TABLE messages (
//   conversation_id TEXT,
//   message_id BIGINT,
//   sender_id TEXT,
//   content TEXT,
//   content_type TEXT,
//   created_at TIMESTAMP,
//   PRIMARY KEY (conversation_id, message_id)
// ) WITH CLUSTERING ORDER BY (message_id DESC);

class MessageStore {
  // Get message history for a conversation (paginated)
  async getMessages(
    conversationId: string,
    beforeMessageId?: string,
    limit = 50
  ): Promise<ChatMessage[]> {
    let query = "SELECT * FROM messages WHERE conversation_id = ?";
    const params: any[] = [conversationId];

    if (beforeMessageId) {
      query += " AND message_id < ?";
      params.push(beforeMessageId);
    }

    query += " ORDER BY message_id DESC LIMIT ?";
    params.push(limit);

    return this.cassandra.execute(query, params);
  }

  // Store a new message
  async saveMessage(message: ChatMessage): Promise<void> {
    await this.cassandra.execute(
      "INSERT INTO messages (conversation_id, message_id, sender_id, content, content_type, created_at) VALUES (?, ?, ?, ?, ?, ?)",
      [message.conversationId, message.id, message.senderId, message.content, message.contentType, message.timestamp]
    );
  }
}

interface GroupChat {
  id: string;
  name: string;
  memberIds: string[];
  adminIds: string[];
  createdAt: Date;
}

class GroupMessageHandler {
  async sendGroupMessage(
    groupId: string,
    message: ChatMessage
  ): Promise<void> {
    // Persist the message once (not per member)
    await this.messageStore.saveMessage(message);

    // Get group members
    const members = await this.getGroupMembers(groupId);

    // Deliver to each member (except sender)
    const deliveryPromises = members
      .filter((memberId) => memberId !== message.senderId)
      .map((memberId) => this.router.routeMessage(message, memberId));

    // Fan-out delivery in parallel
    await Promise.allSettled(deliveryPromises);
  }
}

// For large groups (>100 members), consider:
// 1. Batch delivery to reduce per-message overhead
// 2. Use a pub/sub channel per group instead of individual delivery
// 3. Rate limit messages to prevent spam

class PresenceService {
  private redis: RedisClient;

  // Called when user connects
  async setOnline(userId: string): Promise<void> {
    await this.redis.set(`presence:${userId}`, "online");
    // Notify friends/contacts about status change
    await this.broadcastStatusChange(userId, "online");
  }

  // Called when user disconnects
  async setOffline(userId: string): Promise<void> {
    // Don't immediately mark offline (handles brief disconnects)
    // Use a delayed approach
    await this.redis.set(`presence:${userId}`, "offline");
    await this.redis.set(`last_seen:${userId}`, Date.now().toString());
    await this.broadcastStatusChange(userId, "offline");
  }

  // Heartbeat approach: clients send heartbeats every 30 seconds
  // If no heartbeat received for 60 seconds, mark offline
  async heartbeat(userId: string): Promise<void> {
    await this.redis.setex(`heartbeat:${userId}`, 60, "alive");
    await this.redis.set(`presence:${userId}`, "online");
  }

  async isOnline(userId: string): Promise<boolean> {
    return (await this.redis.get(`presence:${userId}`)) === "online";
  }

  // Only broadcast to users who have a conversation with this user
  // Don't broadcast to ALL users (too expensive)
  private async broadcastStatusChange(
    userId: string,
    status: string
  ): Promise<void> {
    const recentContacts = await this.getRecentContacts(userId);
    for (const contactId of recentContacts) {
      await this.router.sendToUser(contactId, {
        type: "presence_update",
        userId,
        status,
        lastSeen: status === "offline" ? Date.now() : undefined,
      });
    }
  }
}