用Context WithCancel取消所有正在运行的goroutines

95 阅读2分钟

在这个例子中,我们正在运行4个goroutines并等待它们完成工作。然而,如果其中任何一个需要超过80毫秒,我们将返回一个错误。如果出现这种情况,所有其他的goroutines将被取消,而不需要等待它们完成它们的工作。为此,我们将使用context.WithCancel功能。

例子

package main

import (
	"context"
	"fmt"
	"math/rand"
	"sync"
	"time"
)

var wg sync.WaitGroup

func main() {
	fmt.Println(">>>>> BEGIN")

	// Prevent picking up the same random number all the time for sleeping.
	rand.Seed(time.Now().UnixNano())

	ctx, cancel := context.WithCancel(context.Background())
	// Cancel even if everything goes fine without an error to release resources.
	defer cancel()

	for i := 1; i <= 4; i++ {
		wg.Add(1)
		go func(i int) {
			fmt.Println(i, ": STARTED")
			err := doSomething(ctx, &wg, i)
			if err != nil {
				fmt.Println(err)
				cancel()
			} else {
				fmt.Println(i, ": FINISHED")
			}
		}(i)
	}
	wg.Wait()

	fmt.Println(">>>>> END")
}

func doSomething(ctx context.Context, wg *sync.WaitGroup, id int) error  {
	defer wg.Done()

	// Pick a random number to simulate time it takes to finish the job.
	delay := rand.Intn(100)

	select {
	case <-ctx.Done():
		return fmt.Errorf("%d: CANCELLED (%dms)", id, delay)
	default:
		// Prevent from blocking.
	}

	if delay > 80 {
		return fmt.Errorf("%d: FAILED (%dms)", id, delay)
	}
	time.Sleep(time.Duration(delay) * time.Millisecond)

	return nil
}

测试

所有的人都设法在不超过80毫秒的时间内完成他们的工作:

>>>>> BEGIN
1 : STARTED
2 : STARTED
3 : STARTED
4 : STARTED
2 : FINISHED
4 : FINISHED
3 : FINISHED
1 : FINISHED
>>>>> END

只有1个成功完成了。虽然3号和4号的时间要短于80毫秒,但它们被取消了,因为2号在84毫秒时失败了:

>>>>> BEGIN
1 : STARTED
2 : STARTED
2: FAILED (84ms)
3 : STARTED
3: CANCELLED (20ms)
4 : STARTED
4: CANCELLED (14ms)
1 : FINISHED
>>>>> END

由于1号在一开始就以99毫秒的成绩失败了,所以其他所有的人都被取消了,尽管他们的目的是要比80毫秒短:

>>>>> BEGIN
1 : STARTED
1: FAILED (99ms)
2 : STARTED
2: CANCELLED (73ms)
3 : STARTED
3: CANCELLED (20ms)
4 : STARTED
4: CANCELLED (16ms)
>>>>> END

由于1号在一开始就以87毫秒的速度失败,所以其他所有的都被取消了。虽然3号也被设定为以84毫秒失败,但它没有机会了。相反,它也被取消了:

>>>>> BEGIN
1 : STARTED
1: FAILED (87ms)
2 : STARTED
2: CANCELLED (64ms)
3 : STARTED
3: CANCELLED (84ms)
4 : STARTED
4: CANCELLED (77ms)
>>>>> END
>>>>> BEGIN
1 : STARTED
2 : STARTED
4 : STARTED
3 : STARTED
3: FAILED (85ms)
2 : FINISHED
1 : FINISHED
4 : FINISHED
>>>>> END