クエリ実行中断の実装

はじめに

こんにちは、DevOps グループの中島です。

当社ではデータ分析のために Redash を利用しており、
DevOps グループではその管理・運用を行っています。
日々、Redash が重いとかクエリが返ってこない等の問い合わせに対応する中で、
その原因の一部について調査した内容を紹介します。

発生した問題

事象としては発行したクエリが返ってこないという問い合わせで、
調べてみると接続先のデータベース (Aurora MySQL 5.8) の
CPU 使用率が 100% 付近になっていました。
show processlist で実行中のクエリを確認すると、
長時間実行中のクエリが複数存在しており、これが原因でした。

問題の対応

運用者の権限であれば、Redash の画面から実行中のクエリを見ることができ、
ここでクエリ実行のキャンセルをすることが出来ます。
しかし、キャンセルボタンを押してもなかなか消えません。
結局データベースに直接接続して、対象のクエリを KILL することで
CPU 使用率は正常に戻りました。

クエリが中断されない理由

ではなぜキャンセルボタンを押しても実行中のクエリが消えないのか
という点が気になってくるところです。

Redash のコードを見てみたところ、キャンセル処理は SQL の実行を担当する
ワーカープロセスに対して SIGINT を送信することで行っていました。
MySQL 用のクエリ実行処理を抜粋するとおおむね以下のようになっています。
(当社で利用している過去バージョンのものです)

try:
    connection = MySQLdb.connect(...),
    cursor = connection.cursor()
    cursor.execute(query)

    data = cursor.fetchall()

    while cursor.nextset():
        data = cursor.fetchall()
    cursor.close()

except KeyboardInterrupt:
    cursor.close()
    error = "Query cancelled by user."
    json_data = None

finally:
    if connection:
        connection.close()

キャンセル実行時には SIGINT が KeyboardInterrupt を発生させ、
cursor と connection を閉じて終了するという意図が読み取れます。
しかし、クエリ実行が中断されないということは、
このコードがうまく動作していないと考えられます。

そこで、以下のような実験用のコードを書いて試してみます。
データベースに select sleep(100) を発行し、
SIGINT を受け取ったら終了するコードです。

import MySQLdb

conn = MySQLdb.connect(...)
cur = conn.cursor()

try:
    sql = "select sleep(100)"
    cur.execute(sql)
    rows = cur.fetchall()
    for row in rows:
         print(row)
except KeyboardInterrupt:
    print('KeyboardInterrupt')

Ctrl + C で SIGINT を送信したところ、プログラムが反応せず、
except KeyboardInterrupt に入ることはありませんでした。
このため、Redash 上からキャンセルボタンを押しても、
クエリの実行は止まらなかったと考えることができます。

MySQLdb は内部で MySQL C API を呼び出すことで MySQL にクエリを発行していますが、
C の処理に制御がある状態では SIGINT のシグナルが受け付けられない、
あるいは C の方にシグナルの処理を明示的に書く必要があるのかもしれません。
上記コードの try の内部において、C で書いた無限ループを呼び出すコードを
試したところ、同様の挙動となったためそのように推測しました。

クエリ実行中断の方法？

話は変わりますが、最近似たような問題がサービス中に発生しました。
内容としては、リクエストを受け付けたアプリケーションが、
データベースへのクエリ実行に時間がかかりすぎていて、
アプリケーションの入り口で設定したタイムアウトになったあとも
クエリがそのまま実行され続け、同じリクエストが何度も行われ
データベースの CPU 使用率が高騰するというものです。

この場合、タイムアウトとなったあとにアプリケーションからクエリを
中断させれば良いのでは、と単純にはそう考えられます*1 。

ここで疑問に思ったのは、そもそもクエリの実行中断というのは
一般的にどのようにするものなのか、ということです。
私が想像していたのは MySQL のプロトコルであれば、張ったコネクション上で
それ用のコマンドを送信すればサーバがクエリ実行を
中断してくれるようなものだと思っていました。

クエリ実行中断の一般的な実装

そこで、代表的な例として mysql コマンドではどのようになっているかを
見てみたところ、クエリ実行中のコネクションとは別の新たなコネクションを確立して、
KILL QUERY コマンドを発行することでクエリを中断していました。(参考)
もう少しエレガントな方法でやっているのかと思っていたのに、
案外泥臭いことをやっているなという感想を持ちました。

一方、PostgreSQL ではどうなっているかを見てみたところ、
C の API である libpq では中断用の API が PQcancel として存在しています。
ただ、この中の処理でも同様に新しいコネクションを確立して
クエリを KILL する実装がされていました。(参考)
こちらは、ドキュメントが存在していて、それによると、
クエリ実行中も入力を随時確認しないといけなくなるので、
処理効率上そのようにしているとのことです。