在使用 context 的時候遇到了開協(xié)程處理任務(wù)的情況，但是直接在協(xié)程里使用主線程的 context 會導(dǎo)致當(dāng)主線程返回時協(xié)程任務(wù)也會因為 context cancel 而失敗，本文提供了兩種辦法可以取消掉 context 里的 timeout 和 deadline，再設(shè)置一個新的 timeout 上去

golang 去除 context 的 deadline

背景

在使用 context 的時候遇到了開協(xié)程處理任務(wù)的情況，但是直接在協(xié)程里使用主線程的 context 會導(dǎo)致當(dāng)主線程返回時協(xié)程任務(wù)也會因為 context cancel 而失敗。
本文提供了兩種辦法可以取消掉 context 里的 timeout 和 deadline，再設(shè)置一個新的 timeout 上去。

方法一，創(chuàng)建一個新 context

最簡單的方案是通過創(chuàng)建一個新的 context 來解決這個問題。

func main() {
	ctx := context.WithValue(context.Background(), "trace", "123")
	ctx, cancel := context.WithTimeout(ctx, time.Second)
	defer cancel()

	go func(ctx context.Context) {
		newCtx := context.WithValue(context.Background(), "trace", ctx.Value("trace"))
		// do something with newCtx
	}(ctx)

	fmt.Println("main finished")
}

但是這種方案有一個缺點，當(dāng)生成一個新的 context 的時候，需要手動把老 context 中的 value 手動拿出來，再設(shè)置到新的里面去。但是在很多情況下，這個 context 是上游傳過來的，并不知道 value 里面有哪些具體的 key?；蛘呃锩娴?value 過多，寫起來很麻煩。

方法二，使用自定義結(jié)構(gòu)體

通過看 context 的源碼，其實可以發(fā)現(xiàn) context 是一個 interface，這就給了我們操作的空間。

type Context interface {
	Deadline() (deadline time.Time, ok bool)
	Done() <-chan struct{}
	Err() error
	Value(key any) any
}

Context 是通過 Deadline() 這個函數(shù)控制整個 ctx 是否超時了的。那么我們就可以通過重寫這個函數(shù)來規(guī)避超時。

// contextWithoutDeadline 偽造了一個沒有 deadline 的 context
type contextWithoutDeadline struct {
	ctx context.Context
}

func (*contextWithoutDeadline) Deadline() (time.Time, bool) {
	return time.Time{}, false
}

func (*contextWithoutDeadline) Done() <-chan struct{} {
	return nil
}

func (*contextWithoutDeadline) Err() error {
	return nil
}

func (l *contextWithoutDeadline) Value(key interface{}) interface{} {
	return l.ctx.Value(key)
}

// DetachDeadline 從 context 剝離 deadline
func DetachDeadline(ctx context.Context) context.Context {
	return &contextWithoutDeadline{ctx: ctx}
}

// SetNewTimeout 設(shè)置新的超時時間
func SetNewTimeout(ctx context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
	return context.WithTimeout(DetachDeadline(ctx), timeout)
}

可以通過 DetachDeadline() 方法來將原 ctx 的 deadline 取消掉，或者直接通過 SetNewTimeout 的方法設(shè)置一個新的超時時間上去。

演示

func main() {
	ctx, _ := context.WithTimeout(context.Background(), time.Duration(10)*time.Second)
	fmt.Println(ctx.Deadline())

	newCtx := contextwarp.DetachDeadline(ctx)
	fmt.Println(newCtx.Deadline())

	newCtx2, _ := contextwarp.SetNewTimeout(ctx, time.Duration(1)*time.Second)
	fmt.Println(newCtx2.Deadline())
}