Ktor で WebSocket サーバーを実装する

Ubie Advent Calendar 4 日目の記事です。まったく業務とは関係ないんですが、 https://byou.chat/ を作った時に Ktor で WebSocket をやったので、簡単な使い方と仕組みを残しておきます。

TL;DR

Main.kt

fun Application.main() {
    install(WebSockets) {
        timeout = Duration.ofSeconds(5)
        pingPeriod = Duration.ofMinutes(1)
    }

    routing {
        webSocket("/") {
            val uniqueId = generateNonce()
            println("Connection to $uniqueId established")

            incoming.consumeEach { frame ->
                if (frame is Frame.Text) {
                    outgoing.send(Frame.Text("Response for message: ${frame.readText()}"))
                    println("Message \"${frame.readText()}\" received from $uniqueId")
                    if (frame.readText() == "close") {
                        close(CloseReason(CloseReason.Codes.NORMAL, "Closed by server"))
                    }
                }
            }

            println("Connection to $uniqueId closed")
        }
    }
}

導入

Ktor は公式で WebSocket をサポートしており、下記のようにモジュールを読み込むだけで有効化することができます。

build.gradle.kts

dependencies {
    implementation("io.ktor:ktor-server-core:$ktorVersion")
    implementation("io.ktor:ktor-server-netty:$ktorVersion")
    implementation("io.ktor:ktor-websockets:$ktorVersion")
}

Main.kt

fun Application.main() {
    install(WebSockets) {
        timeout = Duration.ofSeconds(5)
        pingPeriod = Duration.ofMinutes(1)
    }
}

実装

ルーティング

ルーティングは通常の HTTP リクエストと同様に行うことができます。

Main.kt

fun Application.main() {
    // 省略
    routing {
        webSocket("/") {
            // ...
        }
    }
}

コネクションの確立

webSocket() {} ブロックはコネクションの確立と同時に実行されます。
generateNonce() で固有の ID(nonce)を生成することができます。これはコネクションに紐づいたデータの保持とかに使います。

webSocket("/") {
    val uniqueId = generateNonce()
    println("Connection to $uniqueId established")
}

フレームを受け取る

WebSocket ではフレームという単位でデータを送受信します。
Ktor では Channel という Coroutine の機能を用いて Ktor とアプリケーション間でフレームのやりとりをします。

webSocket("/") {
    val uniqueId = generateNonce()
    println("Connection to $uniqueId established")

    incoming.consumeEach { frame ->
        if (frame is Frame.Text) {
            println("Message \"${frame.readText()}\" received from $uniqueId")
        }
    }
}

incoming の型は ReceiveChannel<Frame> となっており、 ReceiveChannel#consumeEach はコネクションのクローズ等で明示的に Channel がキャンセルされるまで、受信したフレームに対して処理を実行し続ける suspend function です。

WebSocket の仕様では、Ping/Pong フレームでコネクションを確立したり Close フレームでコネクションを閉じたりする必要があるのですが、Ktor ではそのような制御フレームは隠蔽されていて、Text や Binary のようなデータフレームのみを扱うだけで済みます。

制御フレームも自分でハンドルしたい場合は、 webSocket(path) { ... } の代わりに webSocketRaw(path) { ... } を使います。

フレームを送る

フレームの送信は outgoing チャンネルを用います。 型は SendChannel<Frame> となっており、 outgoing.send(Frame) でフレームを Channel に積むと、Ktor が順に送信します。

webSocket("/") {
    val uniqueId = generateNonce()
    println("Connection to $uniqueId established")

    incoming.consumeEach { frame ->
        if (frame is Frame.Text) {
            println("Message \"${frame.readText()}\" received from $uniqueId")
            outgoing.send(Frame.Text("Response for message: ${frame.readText()}"))
        }
    }
}

コネクションのクローズ

コネクションがクローズすると、 incoming チャンネルがキャンセルされ、中断を抜けて consumeEach の後に処理が進みます。
ですので、そこでクローズ時の処理を行うことができます。

webSocket("/") {
    val uniqueId = generateNonce()
    println("Connection to $uniqueId established")

    incoming.consumeEach { frame ->
        // 略
    }

    println("Connection to $uniqueId closed")
}

また、 close(CloseReason) を用いてサーバー側からクローズすることもできます。

webSocket("/") {
    val uniqueId = generateNonce()
    println("Connection to $uniqueId established")

    incoming.consumeEach { frame ->
        if (frame is Frame.Text && frame.readText() == "close") {
            close(CloseReason(CloseReason.Codes.NORMAL, "Closed by server"))
        }
    }

    println("Connection to $uniqueId closed")
}

おわりに

見ての通り Ktor はやっぱりめちゃめちゃ薄いので、ここでフレーム受け取ったあとは好きに設計できて良いですね。手っ取り早く Text フレームで json 送ったりちゃんとプロトコル設計して Binary フレームで送ったり、いろいろ選択肢があって楽しい、けどムズい！