golang使用context控制協程超時時間

背景

　　最近項目需要在實現一個視頻加工的功能主要是用的ffmpeg命令行工具后面會出文章講一講，這里面有用到協程，部門老大review代碼后把我屌 了😢，問我怎么沒對協程設置超時時間。我當時是用的WaitGroup包，去等待協程結果的，這樣會有一個問題就是如果協程處理時間太長就會出現協程堆積的情況爆cup、爆內存，這個問題在我們目前的生產環境是存在的並且有點嚴重，因為一直都有開發任務所以一直沒去處理。

一、基本原理

　　Context.Done()方法的返回值是個<-chan struct{}，當上下文超時或者手動取消上下文時，會自動關閉該channel，利用無緩沖channel會阻塞協程的特點我們可以阻塞協程等到協程執行完畢或或者超時再判斷結果。

二、使用Done方法阻塞協程，等待執行結果

package main

import (
	"context"
	"testing"
	"time"
)

func TestContext(t *testing.T)  {
	//context.WithTimeout 需要傳入一個上下文父上下文，這里只有一個協程所以用context.Background()聲明一個上下文即可
	ctx, cancelFunc := context.WithTimeout(context.Background(), time.Second*2) //定義一個帶有超時時間的上下文
	go func() {
		defer cancelFunc()  //執行完畢后就手動取消上下文，避免傻傻的等待超時，浪費時間
		time.Sleep(time.Second * 5) //默認耗時操作
	}()
	<-ctx.Done() //等待協程執行完成

	
	//判斷執行結果，ctx.Err()，可以拿到執行錯誤，可以判斷是否有超時
	//err := ctx.Err()
	//if err!=nil {
	//	if err.Error() == "context canceled" {
	//		fmt.Println("協程執行完畢")
	//	}
	//	if err.Error() == "context deadline exceeded" {
	//		fmt.Println("上下文超時")
	//	}
	//}

}

=== RUN   TestContext
--- PASS: TestContext (2.00s)
PASS

1、說明

1.context.WithTimeout(parent Context, timeout time.Duration) (Context, CancelFunc)

　　方法需要傳入一個父上下文parent，和一個超時時間timeout，返回的是一個新的上下文Context，和一個取消上下文函數（參數）CancelFunc。

2、Err() error

　　1、在手動調用CancelFunc方法后調用該方法返回的是context canceled err；

　　2、在上下文超時后調用該方法返回的是context deadline exceeded err;

　　3、當上下文沒有超時或者沒有調用CancelFunc方法時調用返回的是nil;

　　4、源碼備注

// If Done is not yet closed, Err returns nil.
// If Done is closed, Err returns a non-nil error explaining why:
// Canceled if the context was canceled
// or DeadlineExceeded if the context's deadline passed.
// After Err returns a non-nil error, successive calls to Err return the same error.
Err() error

3、Done() <-chan struct{}

　　1、當調用 CancelFunc方法或者上下文超時時會關閉channel，阻塞結束

三、context+select 形式

package main

import (
	"context"
	"fmt"
	"testing"
	"time"
)

func TestContext(t *testing.T)  {
	//context.WithTimeout 需要傳入一個上下文父上下文，這里只有一個協程所以用context.Background()聲明一個上下文即可
	ctx, cancelFunc := context.WithTimeout(context.Background(), time.Second*2) //定義一個帶有超時時間的上下文
	go func() {
		defer cancelFunc() //執行完畢后就手動取消上下文，避免傻傻的等待超時，浪費時間
		time.Sleep(time.Second * 1)  //模擬耗時操作
	}()
	for  {
		select {
		case <-ctx.Done():
			fmt.Println("協程序執行完了")
			fmt.Println(ctx.Err())
			return
		default:  //去掉default的話，select就會阻塞主協程，導致后面的代碼無法執行

		}
		fmt.Println("協程還沒執行完！")
		time.Sleep(time.Millisecond*500) //只是為了讓輸出少點
	}
}

=== RUN   TestContext
協程還沒執行完！
協程還沒執行完！
協程序執行完了
context canceled
--- PASS: TestContext (1.01s)
PASS

四、多協程超時控制(協程執行完畢就退出)

/*
* email: [e-mail]
* Author:  oy
* Date:    2021/6/23 下午5:09
* 文章禁止轉載
 */
package main

import (
	"context"
	"testing"
	"time"
)

func Test1(t *testing.T) {
	withTimeout, _ := context.WithTimeout(context.Background(), time.Second*10)
	go func() { //協程1
		time.Sleep(time.Second)
	}()

	go func() { //協程2
		time.Sleep(time.Second * 2)
	}()

	<-withTimeout.Done() //part1  會一直阻塞，直到上下文超時或者上下文被取消

}

=== RUN   Test1
--- PASS: Test1 (10.00s)
PASS

　　以上例子中，協程1會等待1秒，協程2會等待2秒，因為是使用協程並發處理的，兩個函數實際只需要兩秒就可以執行完畢，但是因為我們沒有手動取消上下文導致程序一直阻塞在part1，我們現在想要的是：
　　1、能控制子協程超時時間，比如子協程執行了10s還沒執行完我就不等了，直接進行下一步操作
　　2、就是在子協程都執行完了我就主動退出，不等上下文超時
　　問題1：上下文能知道子協程是否有超時這個不用我們管
　　問題2：要只知道子協程是否都執行完了，這個可以通過go自帶的sync.WaitGroup模塊得到
　　將兩者結合並通過select接受協程信號就可以實現我們想要的效果

/*
* email: [e-mail]
* Author:  oy
* Date:    2021/6/23 下午5:09
* 文章禁止轉載
 */
package UnitTest

import (
	"context"
	"sync"
	"testing"
	"time"
)

func Test1(t *testing.T) {
	withTimeout, cancelFunc := context.WithTimeout(context.Background(), time.Second*10)
	waitGroup := sync.WaitGroup{}
	waitGroup.Add(2)
	go func() { //協程1
		time.Sleep(time.Second * 1)
		waitGroup.Done()
	}()

	go func() { //協程2
		time.Sleep(time.Second * 2)
		waitGroup.Done()
	}()

	go func() { //協程3  監聽協程1、協程2是否完成
		select {
		case <-withTimeout.Done(): //part1
			return //結束監聽協程
		default: //part2 等待協程1、協程2執行完畢，執行完畢后就手動取消上下文，停止阻塞
			waitGroup.Wait()
			cancelFunc()
			return //結束監聽協程
		}
	}()
	<-withTimeout.Done()
    //todo something
}

=== RUN   Test1
--- PASS: Test1 (2.00s)
PASS

　　我們可以定義一個協程3用於監聽子協程，當協程1、協程2執行完畢后就手動取消上下文(part2)，並退出協程3。考慮到在實際生產中我們會存在手動取消上下文的情況，比如在協程1里面已經執行失敗了那么我就沒必要等待其它協程2的結果，於是手動取消上下文，這樣會存在一個問題就是waitGroup.Done()會存在沒有執行到的情況，如果協程3里面只有part2部分，協程3就有可能會變成一個孤兒進程，所以協程3這里設置兩個退出條件即子協程完成退出(part1)和上下文超時也退出(part2)。