API Gateway の WebsocketAPI を Serverless で実装してみる

概要

API Gateway の WebsocketAPI を ServerlessFramework で実装してみました。

こちらのレポジトリにて公開されている Node.js で書かれたコードを Ruby で書き直しています。

掲載したコードはこちらでも公開しています。

WebsocketAPI in API Gateway

今回は $connectと$disconnect、カスタムルートであるsendmessageの３つを Lambda 関数として作成します。

フォルダ構成

.
├── onconnect
│   └── handler.rb
├── ondisconnect
│   └── handler.rb
├── sendmessage
│   └── handler.rb
└── serverless.yml

$connect ルート

クライアントが API と接続を開始した時に呼び出されます。引数で受け取った event オブジェクトからコネクション ID を取り出して DynamoDB に登録しています。

require 'aws-sdk-dynamodb'
require 'aws-sdk-apigatewaymanagementapi'

class WebsocketApi
  @dynamoDB = Aws::DynamoDB::Resource.new(region: 'ap-northeast-1')
  @table = @dynamoDB.table(ENV['TABLE_NAME'])

  def self.connect(event:, context:)
    params = {
      item: {
        connectionId: event['requestContext']['connectionId']
      }
    }
    begin
      @table.put_item(params)
      return {
        statusCode: 200,
        body: 'Connected'
      }
    rescue StandardError => e
      puts "Failed to connect: #{e}"
      return {
        statusCode: 500
      }
    end
  end
end

$disconnect ルート

クライアントまたはサーバーが API から切断した時に呼び出されます。接続時に登録したコネクション ID を削除しています。

require 'aws-sdk-dynamodb'
require 'aws-sdk-apigatewaymanagementapi'

class WebsocketApi
  @dynamoDB = Aws::DynamoDB::Resource.new(region: 'ap-northeast-1')
  @table = @dynamoDB.table(ENV['TABLE_NAME'])

  def self.disconnect(event:, context:)
    params = {
      key: {
        connectionId: event['requestContext']['connectionId']
      }
    }
    begin
      @table.delete_item(params)
      return {
        statusCode: 200,
        body: 'Disconnected'
      }
    rescue StandardError => e
      puts "Failed to disconnect: #{e}"
      return {
        statusCode: 500
      }
    end
 end
end

sendmessage ルート

メッセージを送信するためのカスタムルートです。 DynamoDB からコネクション ID を全て取り出し、それぞれに対してメッセージを送信しています。

require 'aws-sdk-dynamodb'
require 'aws-sdk-apigatewaymanagementapi'

class WebsocketApi
  @dynamoDB = Aws::DynamoDB::Resource.new(region: 'ap-northeast-1')
  @table = @dynamoDB.table(ENV['TABLE_NAME'])

  def self.send_message(event:, context:)
    begin
      resp = @table.scan(
        projection_expression: 'connectionId'
      )
    rescue Aws::DynamoDB::Errors::ServiceError => e
      return {
        statusCode: 500,
        body: e
      }
    end

    api_gw = Aws::ApiGatewayManagementApi::Client.new(
      endpoint: 'https://' + event['requestContext']['domainName'] + '/' + event['requestContext']['stage']
    )

    resp.items.map do |item|
      begin
        api_gw.post_to_connection(
          connection_id: item['connectionId'], data: JSON.parse(event['body'])['data']
        )
      rescue Aws::ApiGatewayManagementApi::Errors::Http410Error => e
        puts "Found stale connection, deleting #{item['connectionId']}"
        @table.delete_item(key: { connectionId: item['connectionId'] })
      rescue StandardError => e
        throw e
      end
    end
    { statusCode: 200, body: 'Data sent' }
  end
end

デプロイする

AWS へのデプロイには Serverless を使用します。利用したことがない場合はnpm install -g serverlessでインストールしてください。以前に使用したことがある場合でもバージョンが 1.38 未満だと WebsocketAPI に対応していないので注意してください。

provider項目のwebsocketsApiRouteSelectionExpressionはメッセージを送信する時に必要になる値です。デプロイ後に WebsocketAPI をテストする段階で説明します。

iamRoleStatementsは Lambda の実行に必要な権限を定義しています。 functionsで先ほど書いたコードを Lambda 関数としてデプロイするように定義しています。 resources以下で定義しているのは DynamoDB の設定です。

service: simple-websocket-chat-ruby

provider:
  name: aws
  region: ap-northeast-1
  runtime: ruby2.7
  websocketsApiName: SimpleChatWebSocket
  websocketsApiRouteSelectionExpression: $request.body.action
  environment:
    TABLE_NAME: simple_chatconnections

  iamRoleStatements:
    - Effect: Allow
      Action:
        - "dynamodb:PutItem"
        - "dynamodb:GetItem"
        - "dynamodb:DeleteItem"
        - "dynamodb:Scan"
      Resource:
        - Fn::GetAtt: [SimpleChatTable, Arn]
    - Effect: Allow
      Action:
        - "execute-api:ManageConnections"
      Resource:
        - "arn:aws:execute-api:*:*:**/@connections/*"

functions:
  connect:
    handler: onconnect/handler.WebsocketApi.connect
    events:
      - websocket:
          route: $connect
  disconnect:
    handler: ondisconnect/handler.WebsocketApi.disconnect
    events:
      - websocket:
          route: $disconnect
  sendMessage:
    handler: sendmessage/handler.WebsocketApi.send_message
    events:
      - websocket:
          route: sendmessage

resources:
  Resources:
    SimpleChatTable:
      Type: "AWS::DynamoDB::Table"
      Properties:
        TableName: simple_chatconnections
        AttributeDefinitions:
          - AttributeName: connectionId
            AttributeType: S
        KeySchema:
          - AttributeName: connectionId
            KeyType: HASH
        ProvisionedThroughput:
          ReadCapacityUnits: 1
          WriteCapacityUnits: 1

WebsocketAPI のテスト

wscat を使用してコマンドラインから WebsocketAPI のテストをしてみます。

$ npm install -g wscat
$ wscat -c wss://{YOUR-API-ID}.execute-api.{YOUR-REGION}.amazonaws.com/{STAGE}
connected (press CTRL+C to quit)
> {"action":"sendmessage", "data":"hello world"}
< hello world

送信したメッセージに含まれている"action":"sendmessage"という値はデプロイ時に定義した項目と対応しています。 websocketsApiRouteSelectionExpressionの$request.body.actionと、functions項目にあるsendMessageの route: sendmessageがそれにあたります。

例えばそれぞれが$request.body.message、route: pushmessageであるとします。するとメッセージ送信時に入力すべき値は"message":"pushmessage"ということです。

おわり

AWS 上に作成されたリソースは以下のコマンドを実行することで削除できます。

$ sls remove

概要#

WebsocketAPI in API Gateway#

フォルダ構成#

$connect ルート#

$disconnect ルート#

sendmessage ルート#

デプロイする#

WebsocketAPI のテスト#

おわり#

概要