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: oyblog@qq.com
* 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: oyblog@qq.com
* 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)。