MENU
AWS

API GatewayとLambdaでWebSocketを試してみる

shio

こんにちは、しおです。

最近サーバレス構成のことを調べていてAWSのドキュメントにWebSocket APIを利用したチャットアプリの構築というものがあったので、今回はそれを参考に自分でやってみたことをブログに書いてみようと思います。

今回参考にしたチュートリアルはこちらです。
チュートリアル: WebSocket API、Lambda、DynamoDB を使用したサーバーレスチャットアプリケーションの構築
https://docs.aws.amazon.com/ja_jp/apigateway/latest/developerguide/websocket-api-chat-app.html

Contents
  1. 下準備
  2. Serverless Frameworkの設定
    1. providerとfunctions
    2. resources
  3. 実装していく
  4. デプロイする

下準備

まずは下準備として構築に必要なものを用意します。

チュートリアルではCloud Formationのサンプルがありましたが、
今回はServerless Frameworkを利用してリソースを作成しました。
また最近はTypeScriptを良く触っているので、TypeScriptで実装できるように準備します。

というわけでyarnを利用してServerless Frameworkの準備をします。

yarn init 
yarn add -D typescript serverless serverless-plugin-typescript

Serverless Frameworkの設定

Serverless Frameworkはインフラの構成管理やデプロイまで行ってくれるツール・フレームワークです。
Serverless Framework
https://www.serverless.com/

ざっくりとした使い方ですが、serverless.yml というyamlファイルを用意して、そこに定義や設定を記載して管理することができます。
今回はAPI Gateway + Lambda + DynamoDB が必要になるので、下記のようなserverless.ymlを作成しました。
詳細はそれぞれ後述します。

service: websocket-serverless

provider:
  name: aws
  runtime: nodejs18.x
  region: ap-northeast-1
  websocketsApiName: 'websocket-test'
  websocketsApiRouteSelectionExpression: '$request.body.action'
  websocketsDescription: 'websocket api'
  iam:
    role:
      statements:
        - Effect: Allow
          Action: 'dynamodb:*'
          Resource:
            - Fn::GetAtt: [connections, Arn]

plugins:
  - serverless-plugin-typescript

functions:
  connectionHandler:
    handler: src/handler.connectHandler
    events:
      - websocket: $connect
  disconnectionHandler:
    handler: src/handler.disconnectHandler
    events:
      - websocket: $disconnect
  defaultHandler:
    handler: src/handler.defaultHandler
    events:
      - websocket: $default

resources:
  Resources:
    connections:
      Type: AWS::DynamoDB::Table
      Properties:
        TableName: connections
        AttributeDefinitions:
          - AttributeName: connectionId
            AttributeType: S
        KeySchema:
          - AttributeName: connectionId
            KeyType: HASH
        ProvisionedThroughput:
          ReadCapacityUnits: 3
          WriteCapacityUnits: 3

providerとfunctions

まずはproviderとfunctionsのところです。

providerの設定には全体で必要な設定を記載できます。
AWSのリージョンやWebSocket用に必要な項目などを設定しています。
またLambdaに紐付けるIAMロールも設定していて、DynamoDBへのアクセスを許可しています。

functionsはデプロイするLambda関数の定義になります。
handlerには実行するモジュールを設定します。（実装の詳細は後述します）
eventsにはAPI Gatewayのルートを設定します。

functions:
  connectionHandler:
    handler: src/handler.connectHandler
    events:
      - websocket: $connect
  disconnectionHandler:
    handler: src/handler.disconnectHandler
    events:
      - websocket: $disconnect
  defaultHandler:
    handler: src/handler.defaultHandler
    events:
      - websocket: $default

connectionHandler, disconnectionHandler, defaultHandler の3つを定義していますが、これはAPI Gatewayのデフォルトで用意されるルート対応させて定義しています。
それぞれ下記のようなルートです。
$connect はAPIの接続が開始される時のルート
$disconnect は接続が切断される時のルート
$default はどのルートにも一致しない場合に呼び出される

また、独自のルートも定義することが可能です。
その場合はproviderで定義したwebsocketsApiRouteSelectionExpression （ルート選択式）からルートを評価します。
今回だとbody内のプロパティ: actionの値で評価するように設定しています。
（ただ、今回の実装では利用しないです）
websocketsApiRouteSelectionExpression: '$request.body.action'

参考：https://www.serverless.com/framework/docs/providers/aws/events/websocket#routes

https://docs.aws.amazon.com/apigateway/latest/developerguide/apigateway-websocket-api-overview.html

resources

resourcesにはLambda関数とは別で作成したいS3やSNSなどを定義することができます。
今回はDynamoDBを利用するのでそれを定義します。

DynamoDBの目的はWebSocketのconnectionIdを保存しておくためです。
API GatewayではWebSocketのそれぞれのconnectionIdを記憶まではしてくれないので、新規の接続時にDynamoDBに保存しておいて、そこからメッセージ送信の際にIdを取得するようにします。

resources:
  Resources:
    connections:
      Type: AWS::DynamoDB::Table
      Properties:
        TableName: connections
        AttributeDefinitions:
          - AttributeName: connectionId
            AttributeType: S
        KeySchema:
          - AttributeName: connectionId
            KeyType: HASH
        ProvisionedThroughput:
          ReadCapacityUnits: 3
          WriteCapacityUnits: 3

実装していく

下準備ができたので実装していこうと思います。

AWSのSDK for JavaScriptを利用しようとしたところ、v2とv3があるようです。
下記のドキュメントに
https://docs.aws.amazon.com/AWSJavaScriptSDK/latest

AWS SDK for JavaScript v3 is the latest and recommended version, which has been GA since December 2020.

https://docs.aws.amazon.com/AWSJavaScriptSDK/latest/#:~:text=AWS%20SDK%20for%20JavaScript%20v3%20is%20the%20latest%20and%20recommended%20version%2C%20which%20has%20been%20GA%20since%20December%202020.

とあるので、今回はv3を利用していきます。
v3のドキュメントはこちらです。
https://docs.aws.amazon.com/AWSJavaScriptSDK/v3/latest/index.html

まずはSDKの用意ですが、v3の場合はサービスごとに個別のモジュールを入れる必要があります。

yarn add @aws-sdk/client-lambda @aws-sdk/client-dynamodb @aws-sdk/client-apigatewaymanagementapi

// 型定義も追加しておきました
yarn add -D @types/aws-lambda

そうしたら実装していきます。
ソースの構成を先に記載しますが、serverless.ymlのhanlderの設定と合わせるため下記のようにします。

src/
  ├ core/
  │   ├ connectHandler.ts
  │   ├ disconnectHandler.ts
  │   └ defaultHandler.ts
  └ handler.ts
serverless.yml
tsconfig.ts

まずはsrc/handler.tsの実装ですが、これはそれぞれの実装をエクスポートするだけです。

export * from './core/connectHandler';
export * from './core/defaultHandler';
export * from './core/disconnectHandler';

それぞれのハンドラを実装していきます。
まずはconnectHandlerから。
基本的にはチュートリアルを参考にして、接続に合わせてDynamoDBに対してconnectionIdを保存するだけです。

import { DynamoDBClient, PutItemCommand } from '@aws-sdk/client-dynamodb';
import { APIGatewayProxyWebsocketHandlerV2 } from 'aws-lambda';

export const connectHandler: APIGatewayProxyWebsocketHandlerV2 = async (
  event,
  context
) => {
  const client = new DynamoDBClient({});
  const input = {
    TableName: 'connections',
    Item: {
      connectionId: { S: event.requestContext.connectionId },
    },
  };
  const putCommand = new PutItemCommand(input);
  await client.send(putCommand);

  return {
    statusCode: 200,
    body: 'hi',
  };
};

そしてdisconnectHandlerはconnectの逆で、
切断に合わせてDynamoDBから該当のconnectionIdを削除します。

import { DeleteItemCommand, DynamoDBClient } from '@aws-sdk/client-dynamodb';
import { APIGatewayProxyWebsocketHandlerV2 } from 'aws-lambda';

export const disconnectHandler: APIGatewayProxyWebsocketHandlerV2 = async (
  event,
  context
) => {
  const client = new DynamoDBClient({});
  const input = {
    TableName: 'connections',
    Key: {
      connectionId: { S: event.requestContext.connectionId },
    },
  };
  const deleteCommand = new DeleteItemCommand(input);
  await client.send(deleteCommand);

  return {
    statusCode: 200,
    body: 'bye',
  };
};

最後にdefaultHandlerです。
チュートリアルではメッセージ送信用の独自のルートを別途用意していますが、今回はdefaultのルートで実装してしまいます。
やっていることは、DynamoDBに保存されているconnectinIdを全て取得して、それぞれsocketに対してリクエストを送信するようにしています。
（自身のconnectionには送らないように判定を入れてます）

import { ApiGatewayManagementApi } from '@aws-sdk/client-apigatewaymanagementapi';
import { DynamoDBClient, ScanCommand } from '@aws-sdk/client-dynamodb';
import { APIGatewayProxyWebsocketHandlerV2 } from 'aws-lambda';

export const defaultHandler: APIGatewayProxyWebsocketHandlerV2 = async (
  event,
  context
) => {
  try {
    const client = new DynamoDBClient({});
    const input = {
      TableName: 'connections',
    };
    const scanCommand = new ScanCommand(input);
    const connectionIds = await client.send(scanCommand);

    const { domainName, stage } = event.requestContext;
    const apigateway = new ApiGatewayManagementApi({
      endpoint: `https://${domainName}/${stage}`,
    });

    await Promise.all(
      connectionIds.Items?.map(async (item) => {
        const connectionId = item.connectionId.S;
        if (connectionId !== event.requestContext.connectionId) {
          await apigateway.postToConnection({
            ConnectionId: connectionId,
            Data: new TextEncoder().encode(event.body),
          });
        }
      }) ?? []
    );

    return {
      statusCode: 200,
      body: 'success',
    };
  } catch (error) {
    return {
      statusCode: 500,
      body: 'error',
    };
  }
};

ここまでで一通りの実装は完了です。

デプロイする

あとは実際にAWS上にデプロイしてみましょう。

Serverless Frameworkを利用したデプロイは、コマンドを叩くだけで完了です。
フレームワークの裏側でCloudFormationのスタックを作成してくれています。

デプロイのコマンドは下記です
profileの設定などはdefaultを利用しにいきますが、コマンド引数として --aws-profileを渡すと切り替えることができます。
またAPI Gatewayのステージはdevがデフォルトですが、こちらも--stageを渡して設定することも可能です。
serverless.ymlに上記の設定を記載することも可能なので、適宜設定を追加してください。

yarn sls deploy
yarn sls deploy --aws-profile foobar

問題がなければデプロイが完了して、API Gatewayや紐づくLambda関数が作成されているはずです。

実際に機能しているかどうか確認してみましょう。
チュートリアルではwscatというのが紹介されているので、それを利用します。

wscatの参考
https://docs.aws.amazon.com/ja_jp/apigateway/latest/developerguide/apigateway-how-to-call-websocket-api-wscat.html

wscatを利用して接続してみます。
API Gatewayのurlは実際のもので読み替えてください。

wscat -c wss://****.ap-northeast-1.amazonaws.com/dev

複数のターミナルから接続するとSocketで通信することができるはずです。
今回はルートの設定をしていないので、内容に関わらず全てdefaultHandlerの処理が実行されて他の接続に対してメッセージが送信されます。
下記のような形で通信が確認できると思います。

完成！！

AUTHOR
shio
shio
記事URLをコピーしました