この記事は NTTテクノクロス Advent Calendar 2019の2日目の記事です。

こんにちは。安田と申します。
NTTテクノクロスでAI関連の新製品開発を担当しています。早速本題からズレますけれども、マンガを描くのが最近の息抜きで、次のようなマンガを描いています。

マンガでわかるデータ連携
 マンガでわかるAI時代のエンタープライズ・アーキテクチャ
 マンガでわかるITストラテジー

・・・さて、本記事ではAWSサーバレスで「WebSocket」やる方法を見て行きたいと思います。2018年12月からAWS API GatewayがWebSocket通信を張れるようになっているようで、いつか試してみたいなと思っていまして、試してみました。調べてみると、AWS Lambdaにnode.jsを使うサンプルプログラムはたくさん見つかるのですが、なかなかPythonサンプルが見つからず・・・私はPythonの方が得意なのでPythonコードを書いていきたいと思います。

作るもの

WEBチャットを作ります。そのWebページに接続して、メッセージを書き込むと、そのときそのページに接続している（WebSocketのセッションを張っている）全ての端末にリアルタイムでメッセージが配信されます。

システム構成

AWSのサーバレスの代表的な要素であるAPI Gateway、Lambda、DynamoDBを組み合わせて作ります。公開するWebページはPCローカル上にHTMLファイルを置くだけでも動きますし、S3の静的ホスティングに置いて簡単に公開しても動きます。

コーディング

1. まずは、API Gatewayを用意するぞ！

1-1. 「WebSocket」のAPI Gatewayを新規作成

「WebSocket」を選択して、新しいAPI Gatewayを作成します。

ルート選択式は、ここでは「$request.body.action」としておきました。これは後から変更できます。

1-2. ここから一つひとつルートを作っています

接続したときのルート(connect)、接続が外れたときのルート(disconnnect)、メッセージを送った時のルート(sendMessage)を一つひとつ作っていきます。

・・・ですが、その前にDynamoDBとIAMロールを作らなくては、です。

1-3. デプロイしてステージを作る

まだ何もルートを作っていない段階ですが、デプロイしてステージを作っておきます。

青く消したところは、固有の文字列が入ります。この固有の文字列は、次のIAMロール作成のときに使います。

2. LambdaのためのIAMロールを作る

まず、Lambdaを動かすための"AWSLambdaBasicExecutionRole"と、DynamoDBを動かすための"AmazonDynamoDBFullAccess"を追加します。　※FullAccessである必要はないが、簡単なため。

そのあと、次のようなインラインポリシーを設定しておきます。これで、LambdaからAPI Gattewayにアクセスできます。

インラインポリシー

{
    "Version": "2012-10-17",
    "Statement": [
        {
            "Action": [
                "execute-api:ManageConnections"
            ],
            "Resource": [
                "arn:aws:execute-api:ap-northeast-1:<ここにテナントID>:<ここに、さっき設定したAPI Gatewayの固有文字列を書く>/*"
            ],
            "Effect": "Allow"
        }
    ]
}

上記のロールの画像で青く消しているところは、先にAPI Gatewayをデプロイしたときに出てきた固有の文字列を設定してみてください。

3. 接続情報を管理するテーブル(DynamoDB)を作る

プライマリキーは、文字列"id"としておきましょうか。

4. 接続したときのルート(connect)を作る

4-1. Lambdaを作る

新しいLambda/Python3.8を作ります。先ほど使ったIAMロールを設定します。
プログラムは、異常系を鑑みなければ、実質2行のみなので簡単です。（実際には、異常系を考慮した作りにすべきです。）
接続されたconnection_idを取得して、それをDynamoDBに登録するだけです。

ac19_onConnect/lambda_function.py

import json
import os
import boto3

dynamodb = boto3.resource('dynamodb')
connections = dynamodb.Table(os.environ['CONNECTION_TABLE'])

def lambda_handler(event, context):
    connection_id = event.get('requestContext',{}).get('connectionId')
    result = connections.put_item(Item={ 'id': connection_id })
    return { 'statusCode': 200,'body': 'ok' }

Lambdaの環境変数"CONNECTION_TABLE"には、先ほど作成したDynamoDBの名前を設定しておきます。タイムアウト時間は3秒でも大丈夫でしょうけど、1分くらいに設定しておきました。

4-2. API GatewayとLambdaを統合する

先ほど作ったAPI Gatewayのconnectルートと、上記のLambda"ac19_onConnect"を統合します。

以降で説明するdisconnectとsendMessageのルートをLambdaと統合する際も、同様の手順となります。統合したら、ステージのデプロイも忘れずに。

5. 接続解除したときのルート(disconnect)を作る

connectルートと同様の手順で、まずLambdaのプログラムを作って、API Gatewayのdisconnectルートと統合します。異常系を考えなければプログラムは簡単で、接続解除されたconnection_idを取得して、それをDynamoDBから削除するだけです。Lambdaの他の設定（ロール、タイムアウト、環境変数）は、先のonConnectのプログラムと同じ。

ac19_onDisconnect/lambda_function.py

import json
import os
import logging
import boto3

dynamodb = boto3.resource('dynamodb')
connections = dynamodb.Table(os.environ['CONNECTION_TABLE'])


def lambda_handler(event, context):
    connection_id = event.get('requestContext',{}).get('connectionId')
    result = connections.delete_item(Key={ 'id': connection_id })
    return { 'statusCode': 200, 'body': 'ok' }

6. sendMessageを受信したときのルート(sendMessage)を作る

sendMessageだけ少し行数がありますが、やっていることは簡単です。
sendMessageを受信したら、DynamoDBに登録されている各接続に対して、それぞれメッセージを配信しているだけです。これまでと同様の手順で、まずLambdaのプログラムを作って、API GatewayにsendMessageルートを作成して、そのルートと統合します。その後、忘れずにAPI Getewayをステージにデプロイしておいてください。Lambdaの他の設定（ロール、タイムアウト、環境変数）は、先の2つのLambdaと同じ。

ac19_sendMessage/lambda_function.py

import json
import os
import sys
import logging
import boto3
import botocore

dynamodb = boto3.resource('dynamodb')
connections = dynamodb.Table(os.environ['CONNECTION_TABLE'])


def lambda_handler(event, context):

    post_data = json.loads(event.get('body', '{}')).get('data')
    print(post_data)
    domain_name = event.get('requestContext',{}).get('domainName')
    stage       = event.get('requestContext',{}).get('stage')

    items = connections.scan(ProjectionExpression='id').get('Items')
    if items is None:
        return { 'statusCode': 500,'body': 'something went wrong' }

    apigw_management = boto3.client('apigatewaymanagementapi',
                                    endpoint_url=F"https://{domain_name}/{stage}")
    for item in items:
        try:
            print(item)
            _ = apigw_management.post_to_connection(ConnectionId=item['id'],
                                                         Data=post_data)
        except:
            pass
    return { 'statusCode': 200,'body': 'ok' }

7. HTMLのWebページを作る

プログラム内の"wss://"のところには、さっき設定したAPI Gatewayの固有文字列を書いておいてください。このHTMLはPCローカルに置いても動きますし、AWS S3の静的ホスティングなどで公開しても動かせます。

index.html

<!DOCTYPE html>
<html>
<head>
    <meta charset="UTF-8">
    <title>「WebSocket」お試しチャット！</title>
</head>
<body>
	<H3>「WebSocket」お試しチャット！</H3>
    <input id="input" type="text" />
    <button onclick="send()">送信</button>
    <pre id="output"></pre>
    <script>
        var input = document.getElementById('input');
        var output = document.getElementById('output');
        var socket = new WebSocket("wss://<ここに、さっき設定したAPI Gatewayの固有文字列を書く>.execute-api.ap-northeast-1.amazonaws.com/v01");

        socket.onopen = function() {
           output.innerHTML += "接続できました！\n";
        };

        socket.onmessage = function(e) {
            output.innerHTML += "受信：" + e.data + "\n";
        };

        function send() {
            socket.send(JSON.stringify(
                {
                    "action":"sendMessage",
                    "data": input.value
                }
            ));
            input.value = "";
        };
    </script>
</body>
</html>

よし！これで完成です！試してみましょう！

参考

主に、AWSのページを見ながら作りました。
https://aws.amazon.com/jp/blogs/news/announcing-websocket-apis-in-amazon-api-gateway/

おわりに

AWSサーバレスでWebSocketを張れるようになったので、色々なものが作れそうでワクワクしています。それでは、NTTテクノクロス Advent Calendar 2019 の3日目も、引き続きお楽しみください！