EC2 SSM と Lambda を使って自動でコマンドを実行して  Redash をいい感じに再起動させる

これは[ミクシィグループ Advent Calendar 2019](https://qiita.com/advent-calendar/2019/mixi) 12日目の記事です。

# はじめに

- Redash v7系を前提に話しています（Redash v8はDockerで動かす必要があってアップデートできてない）。
- 今年の春頃社内向けに書いたものを若干修正しつつ公開するので、ちょっと情報が古かったりします。

# TL;DR

- CloudWatch でEC2インスタンスのCPU使用率を監視
- CPU使用率がしきい値を超えたらAmazon SNS経由で通知
- Amazon SNSをトリガーにLambdaを起動して、SSMを通じて再起動コマンドを実行

# Redash サーバのCPU使用率 100% に悩まされる

私が所属している minimo という部署では、部署内で解析を目的にBIツールとして、Redash を Amazon EC2 にてホスティングしています。

Redash では時間がかかるクエリを実行中に「あ、ちょっとまって、やっぱやめた」というときにクエリの実行をキャンセルすると、UI上は停止したように見えても裏ではジョブキューシステム（Celery）のワーカーが動き続けることがあります。
そうなってしまうと、そのワーカーが暴走を始めCPU使用率100%となり、ほかのクエリを実行できなくなってしまいます。

https://blog.adachin.me/archives/8757 ( さらにこの記事で言及されている https://github.com/getredash/redash/issues/1605 も参照のこと ) supervisord の設定をチューニングするといいらしいですが試せておらず、とりあえず大雑把に Celery のプロセスを pkill して Redash のプロセスも再起動しちゃえという方針で対応しています。
当初は手動で Redash の EC2 インスタンスに ssh してコマンドを叩いていましたが、さすがに面倒になり始め SSM を使ってみたかったのもあり、CPU使用率100%になったら自動的に Lambda 経由で再起動コマンドが叩かれるようにしてみることにしました。


# SSM を使ってみる

## Systems Manager のインスタンスプロファイルの作成

https://docs.aws.amazon.com/ja_jp/systems-manager/latest/userguide/sysman-configuring-access-role.html

デフォルトでは、Systems Manager にはインスタンスでアクションを実行する権限がありません。IAM インスタンスプロファイルを使用してアクセスを許可する必要があります。

### 初回の場合（または新規にRoleを作成する必要がある場合）

ここは端折ります。詳しくはドキュメントを参照してください。

雑に `SSMforEC2` という名前で作成しました。

### すでにRoleが存在する場合

ここも端折ります。詳しくはドキュメントを参照してください。

## ssm-agent のインストール

運用している Redash が動いているサーバは Ubuntu なので、以下のURLを参考に ssm-agent をインストールします。

https://docs.aws.amazon.com/ja_jp/systems-manager/latest/userguide/sysman-manual-agent-install.html#agent-install-ubuntu-deb

以下が実行した手順です。
上記のURLが最新なので、基本的にそちらを参考にするとよさそうです。
現時点で記載されている手順と同じですが、文章などは省略してます。

### 1. 一時ディレクトリを作り、パッケージをダウンロードする
```
$ mkdir /tmp/ssm
$ cd /tmp/ssm
$ wget https://s3.amazonaws.com/ec2-downloads-windows/SSMAgent/latest/debian_amd64/amazon-ssm-agent.deb
```

### 2. ダウンロードしたパッケージをインストールする

```
$ sudo dpkg -i amazon-ssm-agent.deb
```

### 3. ssm-agent が起動しているか確認する

```
$ sudo systemctl status amazon-ssm-agent
```

### 4. 上記の amazon-ssm-agent is stopped、inactive、または disabled が返された場合は、次のコマンドを実行してサービスを開始する

```
$ sudo systemctl enable amazon-ssm-agent
```

### 5. 再度ステータスを確認し、起動を確認する
```
$ sudo systemctl status amazon-ssm-agent
```


## ssm-agent の自動更新設定

SSM エージェント は、Systems Manager を変更したり、新しい機能を追加したりする度に更新されます。インスタンスで最新バージョンの SSM エージェント が常に実行されるように、新しいバージョンが利用可能になると自動的に SSM エージェント を更新する ステートマネージャー の関連付けを作成します。

https://docs.aws.amazon.com/ja_jp/systems-manager/latest/userguide/sysman-state-cli.html

これは手元のマシンで実行します（aws コマンドがインストールされ、設定済み前提）

`Values=InstanceID` の部分は該当するインスタンスのインスタンスIDを指定してください。

```
$ aws ssm create-association --targets Key=instanceids,Values=InstanceID --name AWS-UpdateSSMAgent --schedule-expression "cron(0 0 2 ? * SUN *)"
```

毎週日曜日の午前2:00 (UTC) に実行するようになります。


## AWSに認識されているか確認

[AWS Systems Manager の マネージドインスタンス](https://ap-northeast-1.console.aws.amazon.com/systems-manager/managed-instances?region=ap-northeast-1) に該当のインスタンスが表示されるはず。


## 実行ができるか確認

手元のマシンで実行する（aws コマンドがインストールされ、設定済み前提）

`InstanceID` の部分は該当するインスタンスのインスタンスIDを指定してください。

```
$ aws ssm send-command --instance-ids "InstanceID" --document-name "AWS-RunShellScript" --comment "PWD" --parameters commands=pwd
```

pwd が実行されます。実行は非同期なので、基本的に `Pending` ステータスの状態でレスポンスが返ってきます。

応答例

```
{
    "Command": {
        "CommandId": "xxxxxxxx-xxxx-xxxx-xxxx-25e53d10ef7e",
        "DocumentName": "AWS-RunShellScript",
        "DocumentVersion": "",
        "Comment": "PWD",
        "ExpiresAfter": 1547547818.291,
        "Parameters": {
            "commands": [
                "pwd"
            ]
        },
        "InstanceIds": [
            "i-000000000000"
        ],
        "Targets": [],
        "RequestedDateTime": 1547540618.291,
        "Status": "Pending",
        "StatusDetails": "Pending",
        "OutputS3BucketName": "",
        "OutputS3KeyPrefix": "",
        "MaxConcurrency": "50",
        "MaxErrors": "0",
        "TargetCount": 1,
        "CompletedCount": 0,
        "ErrorCount": 0,
        "DeliveryTimedOutCount": 0,
        "ServiceRole": "",
        "NotificationConfig": {
            "NotificationArn": "",
            "NotificationEvents": [],
            "NotificationType": ""
        },
        "CloudWatchOutputConfig": {
            "CloudWatchLogGroupName": "",
            "CloudWatchOutputEnabled": false
        }
    }
}
```

コマンドの詳細を確認することで、実行状態と実行結果を確認できます。

```
$ aws ssm list-command-invocations --command-id xxxxxxxx-xxxx-xxxx-xxxx-25e53d10ef7e --details
```

コマンドの実行結果がエラーの場合（存在しないコマンドを実行してみました）

```
$ aws ssm send-command --instance-ids "InstanceID" --document-name "AWS-RunShellScript" --comment "exec not found command" --parameters commands=hoge
```

```
{
    "Command": {
        "CommandId": "xxxxxxxx-xxxx-xxxx-xxxx-25e53d10ef7e",
        "DocumentName": "AWS-RunShellScript",
        "DocumentVersion": "",
        "Comment": "exec not found command",
        "ExpiresAfter": 1547548549.885,
        "Parameters": {
            "commands": [
                "hoge"
            ]
        },
        "InstanceIds": [
            "InstanceID"
        ],
        "Targets": [],
        "RequestedDateTime": 1547541349.885,
        "Status": "Pending",
        "StatusDetails": "Pending",
        "OutputS3BucketName": "",
        "OutputS3KeyPrefix": "",
        "MaxConcurrency": "50",
        "MaxErrors": "0",
        "TargetCount": 1,
        "CompletedCount": 0,
        "ErrorCount": 0,
        "DeliveryTimedOutCount": 0,
        "ServiceRole": "",
        "NotificationConfig": {
            "NotificationArn": "",
            "NotificationEvents": [],
            "NotificationType": ""
        },
        "CloudWatchOutputConfig": {
            "CloudWatchLogGroupName": "",
            "CloudWatchOutputEnabled": false
        }
    }
}
```

再度取得

```
{
    "CommandInvocations": [
        {
            "CommandId": "xxxxxxxx-xxxx-xxxx-xxxx-25e53d10ef7e",
            "InstanceId": "InstanceID",
            "InstanceName": "",
            "Comment": "exec not found command",
            "DocumentName": "AWS-RunShellScript",
            "DocumentVersion": "",
            "RequestedDateTime": 1547541350.04,
            "Status": "Failed",
            "StatusDetails": "Failed",
            "StandardOutputUrl": "",
            "StandardErrorUrl": "",
            "CommandPlugins": [
                {
                    "Name": "aws:runShellScript",
                    "Status": "Failed",
                    "StatusDetails": "Failed",
                    "ResponseCode": 127,
                    "ResponseStartDateTime": 1547541350.469,
                    "ResponseFinishDateTime": 1547541350.483,
                    "Output": "\n----------ERROR-------\n/var/lib/amazon/ssm/InstanceID/document/orchestration/xxxxxxxx-xxxx-xxxx-xxxx-25e53d10ef7e/awsrunShellScript/0.awsrunShellScript/_script.sh: 1: /var/lib/amazon/ssm/IncetanceID/document/orchestration/xxxxxxxx-xxxx-xxxx-xxxx-25e53d10ef7e/awsrunShellScript/0.awsrunShellScript/_script.sh: hoge: not found\nfailed to run commands: exit status 127",
                    "StandardOutputUrl": "",
                    "StandardErrorUrl": "",
                    "OutputS3Region": "ap-northeast-1",
                    "OutputS3BucketName": "",
                    "OutputS3KeyPrefix": ""
                }
            ],
            "ServiceRole": "",
            "NotificationConfig": {
                "NotificationArn": "",
                "NotificationEvents": [],
                "NotificationType": ""
            },
            "CloudWatchOutputConfig": {
                "CloudWatchLogGroupName": "",
                "CloudWatchOutputEnabled": false
            }
        }
    ]
}
```

# Lambda で Redash をいい感じに restart させる

Node.js で AWS Lambda function を書きました。

```js:index.js
'use strict';

const AWS = require('aws-sdk');
const {waiter, WaiterError} = require('./waiter.js');
const { IncomingWebhook } = require('@slack/client');

const cmd = '/home/ubuntu/redash_kill_and_restart.sh';
const ec2RunCommand = `/sbin/runuser -l ubuntu -c '${cmd}'`;

const parseSNSMessage = (event) => {
  if (event['Records']
      && event['Records'][0]
      && event['Records'][0]['Sns']
      && event['Records'][0]['Sns']['Message']) {
    return JSON.parse(event['Records'][0]['Sns']['Message']);
  }
  return undefined;
};

const getInstanceIdAndTypeFromEvent = (event) => {
  if (event.instanceId) {
    return {
      instanceId: event.instanceId,
      type: 'normal',
    };
  }

  const message = parseSNSMessage(event);
  if (!message || !message['AlarmName']) {
    return undefined;
  }

  const instanceIdDimension = message['Trigger']['Dimensions'].find(d => d.name === 'InstanceId' && d.value);
  if (!instanceIdDimension) {
    return undefined;
  }

  const instanceId = instanceIdDimension.value;
  return {
    instanceId: instanceIdDimension.value,
    type: 'alarm',
  };
};

exports.handler = (event, context, callback) => {
  const idAndType = getInstanceIdAndTypeFromEvent(event);
  if (!idAndType) {
    return callback(null, 'Success. But `event` is invalid.');
  }

  if (idAndType.type === 'alarm') {
    const message = parseSNSMessage(event);
    if (message['NewStateValue'] != 'ALARM') {
      // 通常の状態に戻るタイミングでもイベントが発行されるので、ALARM 状態のとき以外は処理を行わず 成功 を返す
      return callback(null, 'Success. But "NewStateValue" is not "ALARM", it did not work.');
    }
  }

  const ssm = new AWS.SSM({region: process.env.AWS_SSM_REGION});
  const ssmSendCommandParams = {
    DocumentName: 'AWS-RunShellScript',
    InstanceIds: [
      idAndType.instanceId
    ],
    Parameters: {
        'commands': [
            ec2RunCommand,
        ],
    },
    TimeoutSeconds: 30 // AWS の仕様上、30秒以上の指定が必須
  };
  ssm.sendCommand(ssmSendCommandParams).promise().then(result => {
    console.log(result);
    return waiter(ssm, 'Success', {InstanceId: idAndType.instanceId, CommandId: result['Command']['CommandId']});
  }).then(result => {
    console.log(result);
    const webhook = new IncomingWebhook(process.env.SLACK_WEBHOOK_URL);
    return webhook.send({
      text: 'Redash を再起動しました',
    });
    callback(null, result);
  }).catch(err => {
    console.log(err, err.stack);
    const webhook = new IncomingWebhook(process.env.SLACK_WEBHOOK_URL);
    if (err.data) {
      webhook.send({
        text: 'Redash の再起動に失敗しました',
        attachments: [
          {
            fallback: 'Redash の再起動に失敗しました',
            color: '#d00000',
            text: '```' + err + "\n" + JSON.stringify(err.data) + '```',
          }
        ],
      });
    } else {
      webhook.send({
        text: 'Redash の再起動に失敗しました',
        attachments: [
          {
            fallback: 'Redash の再起動に失敗しました',
            color: '#d00000',
            text: '```' + err + "\n" + err.stack + '```',
          }
        ],
      });
    }
    callback(err);
  });
};
```

以下のようなイベントを受け取るか、CloudWatch の Alarm を Amazon SNS 経由で受け取る（渡ってくるイベントのJSONは省略）かすると、サーバに置いてあるリスタートスクリプトを実行してくれて、その実行結果を取得するまで待ってからその結果をSlackに投げて終了します。


```json:event.json
{
    "instanceId": "i-xxxxxxxx"
}
```


## waiter を独自に実装

Node.js の AWS SDK は、インスタンス等のステータスが指定のステータスになるまで待って Promise を解決してくれる便利なメソッド（ waiter と呼ばれています ）が生えています。

SSM ではコマンドの実行とその結果の取得はそれぞれ非同期に行われるので、コマンドの実行が成功したか、失敗したかを確認するには、そのどちらかのステータスになるまで待つ必要があります。
しかし、前述した waiter は AWS SDK の SSM には実装されていませんでした。
そのため、独自に waiter を実装しています（似せて作ったけど互換性はない）

```js:waiter.js
const waitStatuses = {
  waiters: {
    "Success": {
      delay: 3,
      maxAttempts: 20,
      acceptors: [
        {
          expected: 'Pending',
          state: 'retry',
        },
        {
          expected: 'InProgress',
          state: 'retry',
        },
        {
          expected: 'Delayed',
          state: 'retry',
        },
        {
          expected: 'Success',
          state: 'success',
        },
        {
          expected: 'Cancelled',
          state: 'failed',
        },
        {
          expected: 'TimedOut',
          state: 'failed',
        },
        {
          expected: 'Failed',
          state: 'failed',
        },
        {
          expected: 'Cancelling',
          state: 'failed',
        },
      ]
    }
  }
};

class WaiterError extends Error {
  constructor (data, ...params) {
    if (typeof data === 'string') {
      super(data, ...params);
    } else {
      super(...params);
      this.data = data;
    }
  }
};

const timer = msec => new Promise(resolve => setTimeout(resolve, msec));

const execRequest = (ssm, state, params, count, delay) => {
  const acceptors = waitStatuses.waiters[state].acceptors;
  return ssm.getCommandInvocation(params).promise().then(data => {
    const acceptor = acceptors.find(v => v.expected === data['Status']);
    if (!acceptor) {
      throw new WaiterError(data, 'not found "' + data['Status'] + '" in "' + state + '" acceptors.');
      return;
    }
    if (acceptor.state === 'success') {
      return data;
    }
    if (acceptor.state === 'failed') {
      throw new WaiterError(data, 'status failed.');
      return;
    }
    if (--count <= 0) {
      throw new WaiterError('over maxAttempts');
      return;
    }
    return timer(delay * 1000).then(() => {
      return execRequest(ssm, state, params, count, delay);
    });
  });
};

// SSM は waiter に対応していないので独自に waiter を実装する
// ほぼ独自実装なので互換性がないことに注意
//
// ## 参考
// - https://github.com/aws/aws-sdk-js/blob/c4452adb0fe7cbabfec1379da17d1a316408ff6b/lib/service.js#L251
// - https://github.com/aws/aws-sdk-js/blob/c4452adb0fe7cbabfec1379da17d1a316408ff6b/lib/resource_waiter.js#L149
//
const waiter = (ssm, state, params) => {
  if (!Object.keys(waitStatuses.waiters).some(key => key === state)) {
    throw new WaiterError('"' + state + '" is not found in waiterStatuses.');
    return;
  }

  const wait = waitStatuses.waiters[state];

  const config = {
    delay: wait.delay,
    maxAttempts: wait.maxAttempts,
  };

  if (params && params.$waiter) {
    // deep copy
    params = JSON.parse(JSON.stringify(params));
    if (typeof params.$waiter.delay === 'number') {
      config.delay = params.$waiter.delay;
    }
    if (typeof params.$waiter.maxAttempts === 'number') {
      config.maxAttempts = params.$waiter.maxAttempts;
    }
    delete params.$waiter;
  }
  return timer(1000).then(() => {
    return execRequest(ssm, state, params, config.maxAttempts, config.delay);
  });
}

module.exports = {
  waiter,
  WaiterError,
};
```

# CloudWatch のアラームと Amazon SNS の設定

CloudWatch で Redash サーバのCPU使用率がしきい値を超えたらアラームを発報させ、Amazon SNS で通知が飛ぶように設定します。
ここでは詳しく解説しないのでよしなにやってみてください。私は力尽きました。

その後、Amazon SNS をトリガーに、上記で作った Lambda を実行させれば、Redash のCPU使用率が100%に張り付き始めたら雑に再起動させる仕組みの完成です。