使われていないDigdagワーカーのEC2を自動でTerminateする

この記事はZOZOテクノロジーズ #5 Advent Calendar 2019 9日目の記事になります。

また、今年は全部で5つのAdvent Calendarが公開されています。

ZOZOテクノロジーズ #1 Advent Calendar 2019
ZOZOテクノロジーズ #2 Advent Calendar 2019
ZOZOテクノロジーズ #3 Advent Calendar 2019
ZOZOテクノロジーズ #4 Advent Calendar 2019

概要

Digdagのクラスタを、EC2インスタンスで組む場合ワーカーをオートスケーリンググループに含めることでプログラムから任意の個数のDigdagワーカーを立ち上げることができます。その辺の構成については以下の記事に詳しく載っています。

ここで問題になるのが、Digdagワーカーのスケールインのタイミングです。ワークフローの最後にスケールインする方法が一番簡単ですが、他のワークフローが起動している場合にスケールインしてしまうと実行中のジョブが失敗してしまいます。また、ワークフローが途中で失敗してしまった場合にスケールインし忘れるということも起こりがちです。

そこで、自動でDigdagのジョブを掴んでいないか判断しスケールインをする仕組みを作ったので紹介します。

queued_task_locks

Digdagで利用されるテーブルの中にqueued_task_locksというテーブルが存在します。
このテーブルはジョブキューのロックの役割を担っており、以下のようなフィールドを持っています。

 id:               bigint
 site_id:          integer
 queue_id:         integer
 priority:         integer
 retry_count:      integer
 lock_expire_time: bigint
 lock_agent_id:    text

このフィールド中の lock_agent_id にはジョブを実行中のホスト名を含んだ文字列が入ります。
そのため各EC2ワーカーのホスト名をユニークにしておくことで、どのワーカーが現在Digdagのジョブを実行しているか判断することができます。

自動Terminate

queued_task_locksを見ることでDigdagワーカーがタスクを掴んでいるかどうかを取れることがわかったので、自動でワーカーがTerminateされる仕組みを作って見ましょう。

ソースコード

export PGPASSWORD='ポスグレのパスワード'
POSTGRESQL_HOST_NAME='ポスグレのホスト名'
POSTGRESQL_USER_NAME='ポスグレのユーザー名'

RETRY_COUNT='リトライ回数'
RETRY_INTERVAL='リトライ間隔'
MYHOST=`hostname`

for i in `seq 1 $RETRY_COUNT`
do
  HOST_LIST=`/usr/bin/psql -t -h POSTGRESQL_HOST_NAME -U POSTGRESQL_USER_NAME -c 'select lock_agent_id from queued_task_locks;'`
  for HOST in $HOST_LIST
  do
    if [ $(echo $HOST | grep -e $MYHOST) ]; then
      exit 0
    fi
  done
  sleep $RETRY_INTERVAL
done

MYINSTANCEID=`curl 169.254.169.254/latest/meta-data/instance-id/`
REGION='AWSのリージョン'
AUTOSCALING_GROUP='オートスケーリンググループのARN'

/usr/local/bin/aws autoscaling detach-instances --instance-ids $MYINSTANCEID --auto-scaling-group-name $AUTOSCALING_GROUP --should-decrement-desired-capacity --region $REGION
/usr/local/bin/aws ec2 terminate-instances --instance-ids $MYINSTANCEID --region $REGION

まず最初にリトライ回数とリトライ間隔を決めて定期的にqueued_task_locksを見に行きます。そこで自分のホストが見つかった場合はスクリプトを終了します。

もしリトライ回数だけqueued_task_locksを見に行ってそれでも自分のホストが見つからない場合は自分をオートスケーリンググループから切り離し、Terminateを行います。
オートスケーリンググループから切り離すときに--should-decrement-desired-capacity をつけることでdesired-capacity が勝手にデクリメントされます。

また、スケールインとTerminateの順番も重要です。先にTerminateをしてしまうとスケールインが実行される前にTerminateがスタートしてしまいます。

最後にこのスクリプトをcronで定期的に実行することで使われていないワーカーが勝手にTerminateされスケールインされます。

cronの間隔・リトライ回数・リトライ間隔を自分のワークフローにあった値にするのがきもとなります。

まとめ

以上のように、queued_task_locksを利用することでワーカー自身がジョブを掴んでいるか判断し使われていない場合に自分でスケールイン・Terminateをする仕組みを紹介しました。

今日で3日続いたDigdagネタは終わりになります。明日は@gold-kouさんによる「NLB配下でgRPC通信するときに考えるALPN対応状況」です。