go routine のキャンセル

Go routine のキャンセルは、Package contextを使うと大変便利。デッドラインとか、タイムアウトとか、非同期のキャンセル系をうまく扱うためのものらしい。

早速単純なサンプルを書いてみた。

package main

import (
	"context"
	"fmt"
	"sync"
	"time"
)

var (
	wg sync.WaitGroup
)

func work(ctx context.Context) error {
	defer wg.Done()
	for {
		select {
		case <-time.After(1 * time.Second):
			fmt.Println("I'm working days and nights!")
		case <-ctx.Done():
			fmt.Println("I've got a cancel from client! Ops")
			return ctx.Err()
		}
	}
}

func main() {
	ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second)
	defer cancel()

	fmt.Println("Hey, Let's do some work, together")
	wg.Add(1)
	go work(ctx)
	wg.Wait()

	fmt.Println("Finished! Enjoy your weekend!")
}

凄く単純で、time.After(1 * time.Second) で、一秒ごとに、case 文の中身が実行される。ただし、呼び元で、 ctx, cancel := context.WithTimeout(context.Background(), 4*time.Second) によって、４秒後には、context に対してキャンセルが発行される。ちなみに、context.Background() はブランクのコンテキスト。他にも、TODO とかあるけど、それは、どんなコンテキストをつかったらいいかまだ分からないときようなので、基本的に Background() でよさげ。

さて、実行結果

$ go run sample/cancelsample/main.go
Hey, Let's do some work, together
I'm working days and nights!
I'm working days and nights!
I'm working days and nights!
I've got a cancel from client! Ops
Finished! Enjoy your weekend!

ばっちり予想通り。

サーバーのレスポンスが遅いとキャンセルする

さて、次は、ちょい難しいものを。サーバー側で、ランダムに、レスポンスが遅いもの、早いものがあるという仕組みを作る。

package main

import (
	"fmt"
	"math/rand"
	"net/http"
	"time"

	"github.com/gorilla/mux"
)

func randomServer(w http.ResponseWriter, r *http.Request) {
	rand.Seed(time.Now().UnixNano())
	value := rand.Intn(3)
	w.WriteHeader(http.StatusOK)
	if value == 0 {
		time.Sleep(time.Second * 6)
		w.Write([]byte(fmt.Sprintf("You've got a slow response")))
	} else {
		w.Write([]byte(fmt.Sprintf("You've got a quick response")))
	}
}

func main() {
	router := mux.NewRouter()
	router.NotFoundHandler = http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		w.WriteHeader(http.StatusNotFound)
		w.Write([]byte(fmt.Sprintf("%s not found\n", r.URL)))
	})
	router.HandleFunc("/api/{name}", randomServer).Methods("GET")
	fmt.Println("...Server stareted")
	fmt.Println("try http://localhost:39090/api/hello")
	http.ListenAndServe(":39090", router)

}

これは、単純なサーバーで、今回は、github.com/gorilla/mux というサードバーティのものを使って書いてみた。サーバーは単純に疑似乱数をつかって、1/3 ぐらいの確率で、サーバーのレスポンスが遅くなるという単純なもの。実行しておく。

$ go run server/cancel/main.go
...Server stareted
try http://localhost:39090/api/hello

一方で、先ほどのキャンセルの仕組みを使った、クライアントを書いておく。ちょっとだけ複雑。

package main

import (
	"context"
	"fmt"
	"io/ioutil"
	"net/http"
	"sync"
	"time"
)

var (
	wg sync.WaitGroup
)

func work(ctx context.Context) error {
	defer wg.Done()
	c := make(chan struct {
		r   *http.Response
		err error
	}, 1)

	tr := &http.Transport{}
	client := &http.Client{Transport: tr}

	req, _ := http.NewRequest("GET", "http://localhost:39090/api/hello", nil)
	go func() {
		resp, err := client.Do(req)
		pack := struct {
			r   *http.Response
			err error
		}{resp, err}
		c <- pack
	}()

	select {
	case <-ctx.Done():
		tr.CancelRequest(req)
		<-c
		fmt.Println("Client Canceled!")
		return ctx.Err()
	case ret := <-c:
		err := ret.err
		resp := ret.r
		if err != nil {
			fmt.Println("Error", err)
			return err
		}
		defer resp.Body.Close()
		out, _ := ioutil.ReadAll(resp.Body)
		fmt.Printf("Server said: %s\n", out)
	}
	return nil
}

func retrive() error {
	fmt.Println("calling http rest service!")
	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()
	wg.Add(1)
	go work(ctx)
	wg.Wait()
	fmt.Println("Finished calling!")
	return ctx.Err()
}

func main() {
	fmt.Println("*** Start Calling the server!")
	for {
		err := retrive()
		if err != nil {
			fmt.Println("The end!")
			return
		}
	}
}

本質的には先ほどと変わらない。ポイントは、

	ctx, cancel := context.WithTimeout(context.Background(), time.Second*5)
	defer cancel()

でコンテキストを定義する。タイムアウト５秒に設定。　ちなみに、defer cancel() は、コンテキストのリソースを開放するために、キャンセル系のところには必要になる。今回、retrive() を書いているので、この、defer を使いたかったため。

func work(ctx context.Context) error

タイムアウトしたい対象の go routine には、コンテキストを渡しておく。

go routine を実行したら

	select {
	case <-ctx.Done():
		tr.CancelRequest(req)
		<-c
		fmt.Println("Client Canceled!")
		return ctx.Err()
	case ret := <-c:
		err := ret.err
		resp := ret.r
		if err != nil {
			fmt.Println("Error", err)
			return err
		}
		defer resp.Body.Close()
		out, _ := ioutil.ReadAll(resp.Body)
		fmt.Printf("Server said: %s\n", out)
	}

select 文を使って、チャネルが返ってくるのを待合せる。それと同時に、context がキャンセルされるケースも待ち合わせておく。とても簡単ですね。
じゃあ、実行。乱数なので、場合によって、異なりますが。

$ go run client/cancelclient/main.go
*** Start Calling the server!
calling http rest service!
Server said: You've got a quick response
Finished calling!
calling http rest service!
Client Canceled!
Finished calling!
The end!

一発目で当たるとこんな感じ

$ go run client/cancelclient/main.go
*** Start Calling the server!
calling http rest service!
Client Canceled!
Finished calling!
The end!

ちなみに、これはサンプルなのでこう書いていますが、本来は、http.Client にタイムアウトあります。How to set timeout for http.Get() requests in Golang?

まとめ

Context を使うと、C# でいう Cancellation Token みたいな感じのことができる。ちなみに、今回のサンプルは GitHub に置いておいた