Goで重い処理をtimeoutさせる ~その２~

前回はGoでtimeoutさせる処理に関して書きました。

yuutookun.hatenablog.com

しかし、前述の記事で書いているのはレスポンスが返らない場合でした。多くの場合、レスポンスが必要だと思うので、レスポンスを付与するパターンを紹介します。

functionの返り値が値とエラーの二つになっていますが、ここは type result を変更すれば増やすことができます。ただ、レスポンスをstructにしておけば、このメソッドで事足りると思います。

流れとしては、

引数の関数の結果格納用のチャネルを作成
goroutineで引数の関数を実行。結果はチャネルに。
for ループのselectでcontextを監視して
チャネルに結果が入れば、結果を返す
context.Doneになれば、空の結果を返す

という感じです。 responseのresultは汎用性のために interface{} で定義しています。

関数終了前にcloseしたりすると、内部でpanicが発生してしまうので、チャネルのcloseに関しては goroutineの終了と一緒にcloseしてます。

以下実際のコードになります

var FailedGetChannel = errors.New("failed to get Channel")

type result struct {
    value interface{}
    err   error
}

func main() {
    f := func() (interface{}, error) {
        time.Sleep(10 * time.Second)
        return []int{1, 2, 3}, nil
    }
    ctx, cancel := context.WithTimeout(context.Background(), 2 * time.Second)
    defer cancel()
    inter, err := ExecWithContext(ctx, f)
    if inter == nil { // interがnullだとcastでpanicになるため
        fmt.Println(err)
        return
    }
    fmt.Println(inter.([]int), err)
}

// ExecWithContext f()の結果をcontextキャンセルされるまで待つ
func ExecWithContext(ctx context.Context, f func() (interface{}, error)) (interface{}, error) {
    resultCh := make(chan result) // f()の実行結果を入れるresult チャネル

    go func() {
        defer close(resultCh)
        resultCh <- func() result {
            i, e := f()
            return result{value: i, err: e}
        }()
    }()
    return waitResult(ctx, resultCh)
}

// contextがtimeoutになるまでchannelの結果が返るのを待つ
func waitResult(ctx context.Context, ch chan result) (interface{}, error) {
    var i result
    for {
        select {
        case <-ctx.Done():
            return i.value, ctx.Err()
        case i, ok := <-ch:
            if !ok {
                return nil, FailedGetChannel
            }
            return i.value, i.err
        default:
        }
        time.Sleep(1 * time.Millisecond)
    }
}