Create a chat app with a WebSocket API, Lambda and DynamoDB

Introduction

1. WebSocket là gì?

Khác với HTTP truyền thống — nơi client phải liên tục "hỏi" server để biết có dữ liệu mới không (polling) — WebSocket cho phép mở một kết nối hai chiều, liên tục giữa client và server. Một khi kết nối được thiết lập, cả hai phía đều có thể chủ động gửi dữ liệu bất cứ lúc nào mà không cần hỏi trước.

2. AWS API Gateway WebSocket API

API Gateway WebSocket API là dịch vụ managed của AWS cho phép bạn xây dựng ứng dụng realtime mà không cần quản lý server WebSocket. AWS xử lý toàn bộ phần infrastructure — bạn chỉ cần viết Lambda function để xử lý logic.

Điểm đặc biệt là API Gateway WebSocket sử dụng khái niệm Routes để phân luồng xử lý:

• $connect — trigger khi client kết nối. Đây cũng là nơi duy nhất bạn cấu hình được authorization (xác thực diễn ra ngay tại thời điểm kết nối).
• $disconnect — trigger khi client ngắt kết nối (đóng tab, mất mạng, hoặc server chủ động đóng).
• $default — route "dự phòng", trigger khi route selection expression không đánh giá được, khi message là non-JSON, hoặc khi giá trị không khớp bất kỳ custom route nào.
• Custom routes — ví dụ sendmessage, joinroom, v.v. — trigger dựa theo giá trị của một field trong message body.

Lưu ý: Route selection expression chỉ áp dụng cho các message thường mà client gửi lên. Riêng $connect và $disconnect luôn được API Gateway trigger tự động theo vòng đời kết nối, không phụ thuộc vào nội dung body.

3. Route Selection Expression

Khi client gửi message, API Gateway cần biết nên gọi Lambda nào. Đây là lúc Route Selection Expression phát huy tác dụng. Với cấu hình $request.body.action, API Gateway sẽ đọc field action trong body của message để quyết định route.

Ví dụ khi client gửi:

{"action": "sendmessage", "message": "Hello everyone!"}

API Gateway đọc action = "sendmessage" → trigger Lambda của route sendmessage. Nếu giá trị action không khớp custom route nào — hoặc message không phải JSON hợp lệ — request sẽ rơi về route $default.

Chỉ message JSON mới định tuyến được theo nội dung. Message non-JSON luôn được đẩy thẳng tới $default. API Gateway hỗ trợ payload tối đa 128 KB (frame tối đa 32 KB).

4. connectionId — Trái tim của WebSocket API

Mỗi khi một client kết nối, API Gateway cấp một connectionId duy nhất cho kết nối đó, và connectionId này tồn tại xuyên suốt vòng đời của kết nối. Đây là định danh để Lambda có thể push message ngược lại xuống đúng client.

Để gửi dữ liệu về client, backend gọi API Gateway Management API (@connections) bằng lệnh PostToConnection — bản chất là một POST request tới:

POST https://{api-id}.execute-api.{region}.amazonaws.com/{stage}/@connections/{connectionId}

Trong ứng dụng chat, mẫu xử lý điển hình là: route sendmessage query toàn bộ connectionId đang active trong DynamoDB, rồi lặp qua từng cái và gọi PostToConnection để phát message tới mọi client đang kết nối.

Lab Introduction

• AWS experience: Intermediate
• Time to complete: 45 minutes
• AWS Region: US East (N. Virginia) us-east-1
• Services used: API Gateway, Lambda, DynamoDB

Architecture

Hệ thống gồm 4 thành phần chính phối hợp với nhau:

• DynamoDB — lưu trữ connectionId của tất cả client đang online
• Connect Function— Lambda xử lý khi client kết nối, lưu connectionId vào DynamoDB
• Disconnect Function— Lambda xử lý khi client ngắt kết nối, xóa connectionId khỏi DynamoDB
• SendMessage Function— Lambda nhận message từ một client, scan DynamoDB lấy tất cả connectionId đang online rồi broadcast message tới tất cả
• Default Function— Lambda xử lý các message không khớp route nào, trả về thông tin connection

Task Details

1. Create DynamoDB Table
2. Create Lambda Functions
   • Connect Function
   • Disconnect Function
   • Default Function
   • SendMessage Function
3. Create WebSocket API Gateway
   • Create API
   • Add Routes
   • Attach Integrations
   • Deploy Stage
4. Add IAM Permissions
5. Update Environment Variables
6. Test với WebSocket Tester

1. Create DynamoDB Table

Vào DynamoDB Console → Tables → Create table

Table details

• Table name: ConnectionsTable
• Partition key: connectionId String

Table settings: Customize settings

Table class: DynamoDB Standard

Read/write capacity settings

• Read capacity units:
  • Auto scaling: Off
  • Provisioned capacity units: 5
• Write capacity units:
  • Auto scaling: Off
  • Provisioned capacity units: 5

Các settings còn lại để default

→ Click Create table

2. Create Lambda Functions

Tạo lần lượt 4 Lambda function. Mỗi function đều làm theo quy trình:

Vào Lambda console → Create function → Author from scratch

Basic information

• Function name: xem phần 2.1 → 2.4
• Runtime: Python 3.14
• Architecture: x86_64

Sau khi tạo xong từng function, vào tab Code → xóa code mặc định → paste code tương ứng vào → nhấn Deploy.

2.1 Connect Function

• Function name: Connect

import os
import json
import boto3

dynamodb = boto3.resource('dynamodb')

def lambda_handler(event, context):
    table = dynamodb.Table(os.environ['TABLE_NAME'])
    connection_id = event['requestContext']['connectionId']

    try:
        table.put_item(Item={'connectionId': connection_id})
    except Exception as e:
        print(e)
        return {
            'statusCode': 500,
            'body': json.dumps('Failed to connect.')
        }

    return {
        'statusCode': 200,
        'body': json.dumps('Connected successfully!')
    }

• Environment variables

Vào Configuration → Environment variables → Edit → Add environment variable

2.2 Disconnect Function

• Function name: Disconnect

import os
import json
import boto3

dynamodb = boto3.resource('dynamodb')

def lambda_handler(event, context):
    table = dynamodb.Table(os.environ['TABLE_NAME'])
    connection_id = event['requestContext']['connectionId']

    try:
        table.delete_item(Key={'connectionId': connection_id})
    except Exception as e:
        print(e)
        return {
            'statusCode': 500,
            'body': json.dumps('Failed to disconnect.')
        }

    return {
        'statusCode': 200,
        'body': json.dumps('Disconnected successfully!')
    }

• Environment variables

Vào Configuration → Environment variables → Edit → Add environment variable

2.3 Default Function

• Function name: Default

import os
import json
import boto3

def lambda_handler(event, context):
    connection_id = event['requestContext']['connectionId']

    apigw = boto3.client(
        'apigatewaymanagementapi',
        endpoint_url=os.environ['CONNECTIONS_URL']
    )

    try:
        conn_info = apigw.get_connection(ConnectionId=connection_id)
        conn_info['connectionID'] = connection_id

        apigw.post_to_connection(
            ConnectionId=connection_id,
            Data=(
                'Use the sendmessage route to send a message. Your info: '
                + json.dumps(conn_info, default=str)
            ).encode('utf-8')
        )
    except Exception as e:
        print(e)
        return {
            'statusCode': 500,
            'body': json.dumps('Failed to process default route.')
        }

    return {
        'statusCode': 200,
        'body': json.dumps('Default route handled successfully!')
    }

• Environment variables

Vào Configuration → Environment variables → Edit → Add environment variable

2.4 SendMessage Function

• Function name: SendMessage

import os
import json
import boto3

dynamodb = boto3.resource('dynamodb')

def lambda_handler(event, context):
    table = dynamodb.Table(os.environ['TABLE_NAME'])

    try:
        response = table.scan()
        connections = response.get('Items', [])
    except Exception as e:
        print(e)
        return {
            'statusCode': 500,
            'body': json.dumps('Failed to fetch connections.')
        }

    sender_id = event['requestContext']['connectionId']
    message = json.loads(event['body'])['message']

    apigw = boto3.client(
        'apigatewaymanagementapi',
        endpoint_url=os.environ['CONNECTIONS_URL']
    )

    failed = []
    for item in connections:
        cid = item['connectionId']
        if cid == sender_id:
            continue
        try:
            apigw.post_to_connection(
                ConnectionId=cid,
                Data=message.encode('utf-8')
            )
        except Exception as e:
            print(f'Failed to send to {cid}: {e}')
            failed.append(cid)

    return {
        'statusCode': 200,
        'body': json.dumps({
            'message': 'Message sent!',
            'failed_connections': failed
        })
    }

• Environment variables

Vào Configuration → Environment variables → Edit → Add environment variable

3. Create WebSocket API Gateway

3.1 Create WebSocket API

Vào API Gateway Console → Create API → WebSocket API → Build

API details

• API name: d-websocket-api-chat-app-demo
• IP address type: IPv4
• Route selection expression: request.body.action

→ Click Next

3.2 Add Routes

Predefined routes

• Click Add $connect route
• Click Add $disconnect route
• Click Add $default route

Custom routes

• Click Add custom route
  • Route key: sendmessage

Click Next

3.3 Attach Integrations

Với mỗi route, chọn Create and attach an integration:

Khi AWS hỏi "Add permission to Lambda function?" → click OK. AWS sẽ tự thêm quyền cho API Gateway invoke các Lambda function.

Click Next

3.4 Deploy Stages

Stages:

• Stage name: dev

→ Click Next

3.5 Review and create

→ Click Create and deploy

3.6 Confirm URL sau khi deploy

Vào API Gateway → d-websocket-api-chat-app-demo → Stages → dev, bạn sẽ thấy 2 URL quan trọng:

• WebSocket URL: wss://xxxxxxxxxx.execute-api.us-east-1.amazonaws.com/dev

→ Dùng cho Client kết nối

• @connections URL : https://xxxxxxxxxx.execute-api.us-east-1.amazonaws.com/dev/@connections

→ Dùng cho Lambda push message

Lưu ý: CONNECTIONS_URL trong env var của Lambda là @connections URL không có /@connections ở cuối.

4. Update Environment Variables

Quay lại bước 2, update biến môi trường CONNECTIONS_URL còn thiếu cho 2 Lambda:

• Default Function
• SendMessage Function

5. Add IAM Permissions

Khi tạo Lambda, mình để default tự tạo basic execution role — role này mới chỉ có quyền ghi CloudWatch Logs, chưa có quyền truy cập DynamoDB hay gọi vào API Gateway. Giờ ta cần add thêm permission vào role cho từng Lambda function như sau:

• Connect Function
  • Add: AmazonDynamoDBFullAccess_v2
  • Lý do: chỉ ghi connectionId vào DynamoDB (PutItem). Không gọi vào API Gateway.
• Disconnect Function
  • Add: AmazonDynamoDBFullAccess_v2
  • Lý do: chỉ xóa connectionId khỏi DynamoDB (DeleteItem). Không gọi vào API Gateway.
• Default Function
  • Add: AmazonAPIGatewayInvokeFullAccess
  • (Tùy chọn) AmazonDynamoDBFullAccess_v2 — chỉ thêm nếu code $default của bạn có truy vấn DynamoDB. Nếu $default chỉ gửi lại message báo lỗi "route không hợp lệ" cho chính client gọi (qua PostToConnection) thì không cần DynamoDB.
• SendMessage Function
  • Add: AmazonDynamoDBFullAccess_v2 và AmazonAPIGatewayInvokeFullAccess
  • Lý do: cần Query/Scan DynamoDB để lấy danh sách connectionId đang active, rồi dùng PostToConnection để push message tới từng client.

6. Verify kết quả

Dùng piehost.com/websocket-tester để test hoặc wscat trên AWS CloudShell

6.1 Test $connect Function

Mở tab 1 → điền WebSocket URL → click Connect:

wss://xxxxxxxxxx.execute-api.us-east-1.amazonaws.com/dev

Kết quả: status Connected — Connect Function đã lưu connectionId vào DynamoDB.

Kiểm tra DynamoDB → ConnectionsTable → Explore items → xác nhận có 1 item với connectionId.

6.2 Test $sendmessage Function

• Mở tab 2 mới, connect cùng URL trên (bây giờ DynamoDB có 2 connectionId).

• Từ tab 2, gửi message:

{"action": "sendmessage", "message": "Hello everyone!"}

• Kết quả: Tab 1 nhận được Hello everyone!

Tại sao tab 1 không nhận lại? Vì SendMessage Function bỏ qua connectionId của người gửi (if cid == sender_id: continue). Đây là behavior của broadcast chat — chỉ gửi tới người khác.

6.3 Test $disconnect Function

Đóng tab 1 → kiểm tra DynamoDB → connectionId của tab 1 đã bị xóa, chỉ còn connectionId của tab 2.

6.4 Test $default Function

Từ tab 2, gửi một message có action không hợp lệ:

{"message": "This will trigger $default"}

Kết quả: nhận về thông tin connection của chính mình.

7. Clean up

• Delete Lambda Functions:
  • Connect Function
  • Disconnect Function
  • Default Function
  • SendMessage Function
• Delete API Gateway: d-websocket-api-chat-app-demo
• Delete DynamoDB Table: ConnectionsTable
• Delete IAM Roles được tạo tự động cho các Lambda

Challenge

Bạn đã xây dựng được chat app realtime serverless hoàn chỉnh! Hãy thử nâng cấp với những thử thách sau:

• Bảo mật WebSocket với Lambda Authorizer + Cognito — WebSocket API không hỗ trợ Cognito Authorizer built-in như REST API. Bạn cần tạo một Lambda Authorizer tùy chỉnh để verify Cognito JWT token tại route $connect. Token thường được truyền qua query string vì WebSocket handshake không hỗ trợ custom header.
• Xử lý stale connection — khi client mất mạng đột ngột, $disconnect không được gọi, connectionId cũ vẫn nằm trong DynamoDB. Hãy bắt GoneException (HTTP 410) trong SendMessage Function để tự động xóa stale connection. Ngoài ra, thêm DynamoDB TTL với thời gian tối đa 2 giờ (giới hạn của API Gateway WebSocket) làm lưới an toàn phòng trường hợp Lambda bị lỗi không xóa kịp.
• Thêm tính năng username — lưu thêm username vào DynamoDB khi connect (truyền qua query string) và hiển thị [username]: message khi broadcast.
• Nâng cấp thành multi-room chat và tối ưu DynamoDB — hiện tại tất cả mọi người đều nhận message của nhau vì chỉ có 1 phòng duy nhất. Hãy thử mở rộng thành nhiều phòng (general, sport, tech, ...) bằng cách lưu thêm roomId vào DynamoDB khi connect. Khi có nhiều phòng, hãy thay scan() bằng query() kết hợp GSI theo roomId để chỉ đọc đúng connection trong phòng cần broadcast — tránh đọc toàn bộ bảng mỗi lần gửi message.

Tổng Kết

Qua bài lab này, chúng ta đã từng bước xây dựng một chat app realtime hoàn chỉnh theo kiến trúc serverless:

• Tạo DynamoDB Table để làm "sổ danh bạ" lưu tất cả connectionId đang online.
• Viết 4 Lambda function xử lý từng giai đoạn của WebSocket lifecycle: connect, disconnect, sendmessage, và default.
• Deploy WebSocket API Gateway với 4 routes tương ứng, mỗi route gắn với một Lambda.
• Cấp IAM permission để Lambda có thể push message ngược lại xuống client qua PostToConnection.
• Test end-to-end với 2 tab WebSocket tester, xác minh broadcast hoạt động đúng.

Kết luận: API Gateway WebSocket API kết hợp với Lambda và DynamoDB là bộ ba serverless lý tưởng cho các ứng dụng realtime — không cần server thường trực, không cần quản lý connection pool, scale tự động và chỉ trả tiền khi có request thực sự.

Tài liệu Tham khảo

• API Gateway WebSocket APIs
• Tutorial: Create a WebSocket chat app with a WebSocket API, Lambda and DynamoDB
• ApiGatewayManagementApi
• Manage connected users and client apps: $connect and $disconnect routes