# Async Chat Server — Model Solution ## High-level approach `asyncio.start_server()` accepts a connection-handler coroutine and spawns it as an **independent task per client**. Each task owns a `StreamReader`/`StreamWriter` pair and runs a read loop that consumes one line at a time with `await reader.readline()`. Because the event loop is **single-threaded and cooperative**, all client tasks run on one thread and only yield control at `await` points — so shared in-memory state (the client/room registries) needs **no locks**, as long as we never `await` in the middle of mutating it. The three things the interviewer is grading: 1. **Concurrency model** — one task per client, non-blocking I/O, broadcast that excludes the sender. 2. **Backpressure** — `await writer.drain()` after writes so a fast sender or many recipients can't grow buffers without bound. 3. **Lifecycle / cleanup** — every exit path (clean `/quit`, EOF, `ConnectionResetError`, any exception) removes the client from shared state and closes the socket, in a `finally` block. --- ## Part 1 — Single global room ```python import asyncio clients: dict["Client", None] = {} # connected clients (used as an ordered set) class Client: def __init__(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter): self.reader = reader self.writer = writer self.username: str | None = None async def send(self, text: str) -> None: """Write one line to this client, applying backpressure.""" self.writer.write((text + "\n").encode()) await self.writer.drain() # block this task (not others) if the buffer is full async def broadcast(text: str, *, exclude: "Client | None" = None) -> None: """Send `text` to every client except `exclude`.""" # Snapshot the keys first: send() awaits, and a client may disconnect mid-broadcast, # mutating `clients`. Iterating a live dict while it changes raises RuntimeError. for client in list(clients): if client is exclude: continue try: await client.send(text) except (ConnectionError, asyncio.IncompleteReadError): pass # a dead recipient gets cleaned up by its own handler; don't fail the sender async def read_line(reader: asyncio.StreamReader) -> str | None: """Read one line as str, or None on EOF/disconnect.""" data = await reader.readline() if not data: # b"" => EOF (peer closed) return None return data.decode(errors="replace").rstrip("\r\n") async def handle_client(reader, writer): client = Client(reader, writer) try: # --- login: first line is the username --- await client.send("Enter username:") name = await read_line(reader) if not name: # disconnected before sending a name return client.username = name clients[client] = None await broadcast(f"* {name} joined *", exclude=client) await client.send(f"Welcome, {name}!") # --- chat loop --- while True: line = await read_line(reader) if line is None: # clean EOF / disconnect break if line == "": # ignore blank lines continue await broadcast(f"{client.username}: {line}", exclude=client) except (ConnectionResetError, asyncio.IncompleteReadError): pass # client vanished mid-line; not an error we care about finally: # Cleanup runs on EVERY exit path: removes from registry, notifies peers, closes socket. clients.pop(client, None) if client.username: await broadcast(f"* {client.username} left *", exclude=client) writer.close() try: await writer.wait_closed() except (ConnectionError, asyncio.IncompleteReadError): pass async def main(host="0.0.0.0", port=8888): server = await asyncio.start_server(handle_client, host, port) addr = ", ".join(str(s.getsockname()) for s in server.sockets) print(f"Serving on {addr}") async with server: await server.serve_forever() if __name__ == "__main__": asyncio.run(main()) ``` ### Why this satisfies the Part 1 requirements - **Multiple concurrent clients** — `start_server` runs `handle_client` as a separate task per connection; the loop multiplexes them cooperatively. - **No self-echo** — `broadcast(..., exclude=client)` skips the sender. - **Graceful disconnects** — `read_line` returns `None` on EOF; `ConnectionResetError`/`IncompleteReadError` are caught; the `finally` always removes the client and closes the writer. - **One bad client can't crash the server** — each connection is an isolated task; an exception in one only ends that task. Broadcast also swallows per-recipient send failures so a single dead socket doesn't abort a broadcast to everyone else. - **Test it** by opening several terminals: `nc localhost 8888` (or `telnet`), entering a name in each, and typing — messages appear in the *other* sessions, not your own. --- ## Part 2 — Rooms and commands We replace the single global `clients` set with a **room index** (`rooms: dict[str, set[Client]]`) plus a back-reference (`client.room`) so each client knows its current room. A line starting with `/` is a command; everything else is a chat message scoped to the sender's room. ```python import asyncio rooms: dict[str, set["Client"]] = {} # room name -> members DEFAULT_ROOM = "lobby" class Client: def __init__(self, reader, writer): self.reader = reader self.writer = writer self.username: str | None = None self.room: str | None = None async def send(self, text: str) -> None: self.writer.write((text + "\n").encode()) await self.writer.drain() def join_room(client: "Client", room: str) -> None: """Move client into `room`, leaving its previous room and GC'ing empties.""" leave_room(client) rooms.setdefault(room, set()).add(client) client.room = room def leave_room(client: "Client") -> None: if client.room and client.room in rooms: members = rooms[client.room] members.discard(client) if not members: # garbage-collect empty rooms so `rooms` can't grow forever del rooms[client.room] client.room = None async def broadcast_room(room: str, text: str, *, exclude=None) -> None: for client in list(rooms.get(room, ())): if client is exclude: continue try: await client.send(text) except (ConnectionError, asyncio.IncompleteReadError): pass async def read_line(reader) -> str | None: data = await reader.readline() if not data: return None return data.decode(errors="replace").rstrip("\r\n") async def handle_command(client: "Client", line: str) -> bool: """Handle a /command line. Returns False if the client should disconnect.""" parts = line.split(maxsplit=1) cmd = parts[0].lower() arg = parts[1].strip() if len(parts) > 1 else "" if cmd == "/join": if not arg: await client.send("Usage: /join <room>") return True old = client.room join_room(client, arg) if old and old != arg: await broadcast_room(old, f"* {client.username} left *") await broadcast_room(arg, f"* {client.username} joined *", exclude=client) await client.send(f"You are now in '{arg}'.") elif cmd == "/who": members = sorted(c.username for c in rooms.get(client.room, ()) if c.username) await client.send("In room '%s': %s" % (client.room, ", ".join(members))) elif cmd == "/quit": await client.send("Bye!") return False # signal the read loop to disconnect else: await client.send(f"Unknown command: {cmd}") # treated as an error, not a chat message return True async def handle_client(reader, writer): client = Client(reader, writer) try: await client.send("Enter username:") name = await read_line(reader) if not name: return client.username = name join_room(client, DEFAULT_ROOM) await client.send(f"Welcome, {name}! You are in '{DEFAULT_ROOM}'.") await broadcast_room(DEFAULT_ROOM, f"* {name} joined *", exclude=client) while True: line = await read_line(reader) if line is None: break if line == "": continue if line.startswith("/"): if not await handle_command(client, line): break else: await broadcast_room(client.room, f"{client.username}: {line}", exclude=client) except (ConnectionResetError, asyncio.IncompleteReadError): pass finally: room = client.room leave_room(client) # removes from exactly one room's membership if client.username and room: await broadcast_room(room, f"* {client.username} left *") writer.close() try: await writer.wait_closed() except (ConnectionError, asyncio.IncompleteReadError): pass async def main(host="0.0.0.0", port=8888): server = await asyncio.start_server(handle_client, host, port) print("Serving on", ", ".join(str(s.getsockname()) for s in server.sockets)) async with server: await server.serve_forever() if __name__ == "__main__": asyncio.run(main()) ``` ### Why this satisfies the Part 2 requirements - **Exactly one room at a time** — `join_room` calls `leave_room` first, so the client is in at most one room's set, and `client.room` always points to that set. - **Room-scoped broadcast** — chat and join/leave notices go through `broadcast_room(client.room, ...)`, never to other rooms. - **`/who`** — replies only to the requester with the sorted member list of the requester's room. - **`/quit`** — `handle_command` returns `False`, breaking the read loop and falling into the `finally` cleanup. - **Consistent membership on disconnect** — the `finally` always `leave_room`s, so a dropped connection can't leave a stale entry; empty rooms are deleted so `rooms` can't grow without bound. - **Commands vs. messages** — only a leading `/` is treated as a command; unknown commands return an error to the sender rather than being broadcast as chat. --- ## Addressing the follow-ups **Horizontal scaling across processes/machines.** A single asyncio process is bound by one CPU core and one machine's memory/FD limits. To scale out, put a **pub/sub backbone** behind the servers (Redis Pub/Sub, NATS, or Kafka). Each node owns its locally-connected sockets and subscribes to the channels for the rooms it has members in. When a node receives a chat line, it (a) delivers locally to same-room clients and (b) publishes to the room's channel; every other node delivers the published message to its own local members of that room. Membership/`/who` then needs a shared view — keep room membership in Redis (a `SET` per room, or per-node presence keys with TTL heartbeats) so any node can answer `/who`. A load balancer with sticky/consistent routing (or any routing, since pub/sub bridges nodes) sits in front. Trade-offs: at-most/at-least-once delivery semantics from the bus, ordering only within a single publisher, and presence-cleanup on node crashes (TTLs). **Slow / stuck clients.** `await writer.drain()` already applies backpressure to *one task* — but a client whose buffer is permanently full will make that client's task hang on `drain()`, and during a broadcast that serializes delivery to everyone after it. Mitigations, in order of robustness: - Wrap each per-recipient send in `asyncio.wait_for(client.send(text), timeout=...)`; on `TimeoutError`, disconnect that client rather than blocking the broadcast. - Give each client a **bounded per-client outbound `asyncio.Queue`** with a dedicated writer task; broadcast does a non-blocking `put_nowait`, and on `QueueFull` you drop the message (or the client). This decouples broadcast latency from the slowest consumer entirely. - Cap `StreamWriter` high-water marks and treat repeated drain timeouts as a disconnect signal. **Graceful shutdown on `SIGTERM`.** Install a signal handler via `loop.add_signal_handler(signal.SIGTERM, ...)` that: stops accepting new connections (`server.close()`), broadcasts a "server shutting down" notice, gives in-flight writers a bounded window to `drain()`, then cancels remaining client tasks and `await`s them so each runs its `finally` cleanup. `asyncio.run()` already cancels pending tasks on exit; an explicit handler lets you drain rather than hard-cancel. **Bounding memory and abuse.** `reader.readline()` will buffer until a newline, so enforce a max line length — use `reader.readuntil(b"\n")` and catch `asyncio.LimitOverrunError` (or set the stream's `limit`) to reject/disconnect oversized lines. Cap total connections, rooms, and members per room. Add per-client **rate limiting** (token bucket: N messages/sec) and idle timeouts via `asyncio.wait_for` on the read. Reject duplicate usernames at login if uniqueness is required. ## Complexity & correctness notes - Broadcast is **O(R)** in the number of recipients in the room; join/leave/`/who` are **O(1)** amortized except the `/who` member-list build which is O(members). - No locks are needed because mutation of `rooms`/membership happens between `await` points; the one subtlety — iterating a collection that a concurrent disconnect may mutate — is handled by snapshotting with `list(...)` before any `await`-driven send loop. - Every connection has a single `finally` that is the sole place membership is torn down, which is what keeps state leak-free regardless of how the connection ends.