Storing WebSocket connection state in DynamoDB #

In a serverless WebSocket architecture, no process remembers who is connected. API Gateway hands your functions a connection ID per event, and every broadcast begins with a question the database must answer: which connections are in this room right now? Most teams start with a table keyed by connection ID and scan it for room members. That works in a demo; in production, scans are slow and expensive, disconnected connections pile up because $disconnect is not guaranteed to fire, a single popular room overloads one partition, and a user in three tabs appears three times in presence. The connection registry is the core data structure of a serverless real-time system, and its key design decides its cost and latency.

Root cause #

A connection registry must answer three queries efficiently: members of a room (for broadcasts), rooms of a connection (for cleanup on disconnect), and connections of a user (for presence and direct messages). A table keyed only by connection ID answers the second cheaply and the others only with a scan, which reads the whole table on every broadcast.

Staleness is the second problem. API Gateway invokes $disconnect on a best-effort basis; after abrupt network failures, gateway maintenance or throttling, it may not arrive. Rows for dead connections then accumulate, broadcasts try to post to them, and each attempt returns 410 Gone — wasting a call, adding latency, and inflating presence counts.

Registry queries and how each design answers them A table keyed only by connection id needs scans for room members and user connections; keying by room and connection makes room queries cheap; adding global secondary indexes on connection id and user id makes the other two queries cheap too. Registry queries and how each design answers them Room members Connection's rooms User's connections PK = connectionId only Scan GetItem Scan PK = room, SK = conn Query needs index needs index + GSI on connectionId Query Query GSI needs index + GSI on userId Query Query GSI Query GSI Every query a broadcast or cleanup needs must be a Query, never a Scan
Key by room for fan-out; index by connection and user for cleanup and presence.

Resolution #

Use one item per (room, connection) membership, keyed by room, with global secondary indexes on connection ID and user ID. Give every item a TTL a little beyond the gateway’s two-hour connection cap, so orphaned rows expire even if $disconnect never fires. Prune rows that return 410 Gone during broadcasts. For very large rooms, spread members across several partition keys to avoid a hot partition.

import { DynamoDBClient } from '@aws-sdk/client-dynamodb';
import { DynamoDBDocumentClient, PutCommand, QueryCommand, BatchWriteCommand } from '@aws-sdk/lib-dynamodb';

/*
  Table "rt-registry"
    pk  (S)  room#<roomId>#<shard>      partition: room members, sharded for big rooms
    sk  (S)  conn#<connectionId>
    connectionId (S), userId (S), joinedAt (N), expiresAt (N, TTL attribute)
  GSI "byConnection": pk=connectionId          → rooms of a connection (cleanup)
  GSI "byUser":       pk=userId, sk=connectionId → a user's connections (presence, DMs)
*/
const TABLE = 'rt-registry';
const ROOM_SHARDS = 8;                           // spread big rooms over 8 partitions
const TTL_SECONDS = 2 * 60 * 60 + 600;           // gateway cap (2 h) + margin
const ddb = DynamoDBDocumentClient.from(new DynamoDBClient({}));

const shardOf = (connectionId: string) =>
  [...connectionId].reduce((h, c) => (h * 31 + c.charCodeAt(0)) >>> 0, 0) % ROOM_SHARDS;

export async function joinRoom(roomId: string, connectionId: string, userId: string) {
  await ddb.send(new PutCommand({ TableName: TABLE, Item: {
    pk: `room#${roomId}#${shardOf(connectionId)}`, sk: `conn#${connectionId}`,
    connectionId, userId, joinedAt: Date.now(),
    expiresAt: Math.floor(Date.now() / 1000) + TTL_SECONDS,   // orphan safety net
  } }));
}

// Room members: one Query per shard, in parallel, paginated.
export async function roomMembers(roomId: string): Promise<string[]> {
  const shards = await Promise.all([...Array(ROOM_SHARDS).keys()].map(async (s) => {
    const ids: string[] = [];
    let ExclusiveStartKey: Record<string, unknown> | undefined;
    do {
      const r = await ddb.send(new QueryCommand({
        TableName: TABLE, KeyConditionExpression: 'pk = :pk',
        ExpressionAttributeValues: { ':pk': `room#${roomId}#${s}` },
        ProjectionExpression: 'connectionId', ExclusiveStartKey,
      }));
      for (const it of r.Items ?? []) ids.push(it.connectionId as string);
      ExclusiveStartKey = r.LastEvaluatedKey;
    } while (ExclusiveStartKey);
    return ids;
  }));
  return shards.flat();
}

// $disconnect (or a 410 during broadcast): remove every membership of the connection.
export async function removeConnection(connectionId: string) {
  const r = await ddb.send(new QueryCommand({
    TableName: TABLE, IndexName: 'byConnection',
    KeyConditionExpression: 'connectionId = :c', ExpressionAttributeValues: { ':c': connectionId },
    ProjectionExpression: 'pk, sk',
  }));
  const keys = (r.Items ?? []).map((it) => ({ DeleteRequest: { Key: { pk: it.pk, sk: it.sk } } }));
  for (let i = 0; i < keys.length; i += 25) {                  // BatchWrite limit: 25 items
    await ddb.send(new BatchWriteCommand({ RequestItems: { [TABLE]: keys.slice(i, i + 25) } }));
  }
}

// Presence: distinct users with at least one live connection (via the byUser index elsewhere).
export async function isOnline(userId: string) {
  const r = await ddb.send(new QueryCommand({
    TableName: TABLE, IndexName: 'byUser', KeyConditionExpression: 'userId = :u',
    ExpressionAttributeValues: { ':u': userId }, Limit: 1, Select: 'COUNT',
  }));
  return (r.Count ?? 0) > 0;
}

The broadcast function from AWS API Gateway WebSocket APIs calls roomMembers, posts to each, and calls removeConnection for any 410 Gone. Together with TTL and $disconnect, that gives three independent cleanup paths — the same layered approach to stale state used for presence in fixing presence flapping and ghost users.

DynamoDB’s TTL deletion is asynchronous — expired items can linger for a while before removal — so queries that must exclude them should filter on expiresAt > now. Treat TTL as garbage collection, not as a precise expiry.

Items read to find a 12-member room's connections Scanning a table keyed by connection id reads every item — ten thousand to a million — to find a twelve-member room, while querying a room-keyed table reads only the twelve members. Items read to find a 12-member room's connections table size grows, room size stays the same Scan: items read per broadcast Query: items read (room of 12) 0 250k 500k 750k 1.0M 10k connections 100k connections 1M connections
Scan cost grows with the whole fleet; Query cost grows only with the room.

Edge cases #

Hot partitions. A partition key per room concentrates a large room’s writes (joins, heartbeat refreshes) and reads (every broadcast) on one partition. Sharding the room key, as above, spreads the load; choose the shard count from the largest expected room.

Heartbeat refresh writes. Refreshing expiresAt on every heartbeat is a write per connection per interval, which is costly at scale. Rely on the fixed TTL beyond the gateway’s two-hour cap instead — no connection can legitimately outlive it — and refresh only on reconnect.

Consistency. Global secondary indexes are eventually consistent, so a connection added a moment ago may be missing from byConnection queries briefly. For cleanup that is harmless (TTL catches it); for presence, a short delay is acceptable.

Verification #

Check the access patterns with realistic data: load the table with a few hundred thousand memberships across rooms of varied size, then measure broadcast query latency and consumed read capacity for small and large rooms. Every registry operation should be a Query or GetItem; search your code and CloudWatch for any Scan:

# Consumed capacity for a room-members query (should be proportional to room size, not table size).
aws dynamodb query --table-name rt-registry --key-condition-expression 'pk = :pk' \
  --expression-attribute-values '{":pk":{"S":"room#r1#0"}}' --return-consumed-capacity TOTAL \
  --select COUNT

Then test orphan cleanup: open connections, kill the clients abruptly so $disconnect may not fire, and confirm the rows disappear through 410 pruning on the next broadcast or, at worst, through TTL.

Three cleanup paths for one dead connection A dead connection is removed by the disconnect handler when it fires, by the broadcast function when posting returns 410 Gone, or by DynamoDB TTL as a final safety net. Three cleanup paths for one dead connection Gateway Registry Broadcast Lambda $disconnect → removeConnection (sometimes never fires) PostToConnection 410 Gone removeConnection TTL expires leftovers None of the three is reliable alone; together they keep the registry accurate
Belt, braces and a TTL.

Operational checklist #

FAQ #

How should I store API Gateway WebSocket connection IDs? #

In a table keyed by room (with the connection ID as sort key), plus indexes by connection ID and user ID. That makes broadcasts, cleanup and presence each a single Query instead of a table scan.

Why do dead connection IDs accumulate in DynamoDB? #

Because $disconnect is best-effort and does not fire for every abrupt disconnection. Add a TTL to every item and delete rows that return 410 Gone when you post to them.

Should I refresh the TTL on every message? #

No — that turns every heartbeat into a write. Since API Gateway closes connections after two hours, a fixed TTL slightly beyond that is enough, refreshed only when a client reconnects.

Can Redis replace DynamoDB for the registry? #

Yes, and many teams use ElastiCache for lower latency. DynamoDB fits a fully serverless stack with no servers or clusters to manage; Redis requires a managed cluster but makes set operations and expiry simpler.

Back to Serverless & Managed WebSockets.