context.WithTimeout()を使うとタイムアウトが想像以上にきれいに書けた

概要

slack の Real Time Messaging を使って bot を書いてたんですが,しばらくすると websocket が切断されてて bot が落ちてしまってました.今回はこれにタイムアウト処理を追加してリカバリできるようにした防備録です.

問題のコード

func (c *Client) GetMessage() (Message, error) {
    var msg Message
    if err := websocket.JSON.Receive(c.socket, &msg); err != nil { // ← 切断されちゃうとこいつが沈黙する
        return msg, err
    }
    return msg, nil
}

タイムアウト処理を入れる

context.Context にはタイムアウト処理をうまくやってくれる仕組みがあるのでこれを使います.

func (c *Client) GetMessage() (Message, error) {
    ctx, cancel := context.WithTimeout(context.Background(), Timeout)  // ← タイムアウト付きの context を設定
    defer cancel()
    ch := make(chan error, 1)

    var msg Message
    go func() {
        ch <- websocket.JSON.Receive(c.socket, &msg)
    }()

    select {
    case err := <-ch: // websocket からメッセージがとれればここへ来る.関数を抜けるときに cancel() が defer で呼ばれて context が解放される
        return msg, err
    case <-ctx.Done(): // ← websocket が沈黙してても,時間が来たら Done() が close される
        return msg, fmt.Errorf("connection lost timeout")
    }
    return msg, nil
}

これで時間が来たら timeout でエラーが返るようになります.エラーが返ってきたら websocket をつなぎ直せば解決です.

が,これだと websocket が生きててもタイムアウトして再接続しちゃいます.タイムアウトする前に websocket を ping で叩いてみて,反応があったら単に関数を抜けるようにすれば無駄な再接続を防げそうです.

func (c Client) GetMessage() (Message, error) {
    ctx, cancel := context.WithTimeout(context.Background(), c.timeout)
    defer cancel()
    ch := make(chan error, 1)

    go func(ctx context.Context, waiting time.Duration) {
        ctx, cancel := context.WithCancel(ctx)  // ← 親の context から新しい context を作ります
        defer cancel()
        select {
        case <-ctx.Done(): // ← 親の context が close したらここにメッセージが来て何もしないで抜ける
        case <-time.After(waiting): // timeout するちょっと前に ping を打つ
            websocket.JSON.Send(c.socket, &Message{Type: EventTypePing, Time: time.Now().Unix()})
        }
    }(ctx, c.timeout-time.Second)

    var msg Message
    go func() {
        ch <- websocket.JSON.Receive(c.socket, &msg)
    }()

    select {
    case err := <-ch: // ← メッセージが無くても websocket が生きてればタイムアウト前には pong が返ってくるはず
        return msg, err
    case <-ctx.Done():
        return msg, fmt.Errorf("connection lost timeout")
    }
    return msg, nil
}

とりあえずこれで安定して動いているようです.context.Context を使うと複雑な処理になりそうなところが想像以上に見通しのよいコードがかけて幸せです.

( '-`).oO( うまいことできた!と思ってるけど,もっとうまいことやれるような気もしてて,なにか方法があれば教えて欲しいです

f:id:ikawaha:20181018223422p:plain

Happy hacking!