通过http库控制请求超时来一窥context的使用

相信很多人对于golang中的context 都有一定的了解，都知道它是一个上下文管理，有一个根context， 再通过这个根context 创建出有更多功能的context， 如具有cancel 功能的，有timeout功能的。 我们搜索的时候，我们可能看到最多的就是以下的代码

package main

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

func work(ctx context.Context, wg *sync.WaitGroup) {
	defer wg.Done()
	for {
		select {
		case <-ctx.Done():
			fmt.Println("ctx 被取消了")
			return
		default:
			fmt.Println("I am working")
			time.Sleep(1 * time.Second)
		}

	}
}

func main() {
	wg := sync.WaitGroup{}
	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)
	wg.Add(1)
	go work(ctx, &wg)
	time.Sleep(3 * time.Second)
	cancel()
	wg.Wait()
	fmt.Println("main process end!")

}

创建一个带有取消功能的ctx，然后用于控制协程里的执行过程。当协程里的select 不停的检查这个ctx是否被取消了，如果取消了，则协程也结束。

可是我一直在想一个问题，这样的玩具代码虽然可以说明ctx的作用，但是真实的项目中我们真的是这么用ctx的吗？ 在介绍ctx时，都会说ctx可以穿梭于协程之间， 一个协程创建了n个子协程，n个子协程又会创建m个孙协程，我们可以将一个带有cancel 功能的ctx挂在每个协程上，这样只要调用这个ctx则所有的协程都会结束。

可是道理都懂，该如何写出具有这种效果的代码呢？

比如我们写了一个http的服务

package main

import (
	"fmt"
	"net/http"
)

func SayHello(w http.ResponseWriter, r *http.Request) {
	fmt.Println("visit SayHello")
	n, err:=w.Write([]byte(string("hello world")))
	fmt.Printf("send %d bytes\n", n)
	if err!=nil{
		fmt.Println(err)
	}
}

func main() {
	http.HandleFunc("/", SayHello)
	http.ListenAndServe("127.0.0.1:8000", nil)

}

代码很简单，当访问 http://127.0.0.1:8000/时会返回 hello world ，服务端打印出

visit SayHello
send 11 bytes

但是我们现在考虑一个问题，这个SayHello 方法处理的太快的，客户端很快的就能得到结果，但是如果它是一个响应非常慢的函数，如需要处理5秒钟，如访问数据库，在做一些计算什么的，此时如果客户端等不及了，关掉浏览器了，这时会发生什么？ 我们改下代码

func SayHello(w http.ResponseWriter, r *http.Request) {
	fmt.Println("visit SayHello")
	time.Sleep(5 * time.Second) // 模拟长时间操作
	n, err := w.Write([]byte(string("hello world")))
	fmt.Printf("send %d bytes\n", n)
	if err != nil {
		fmt.Println(err)
	}
}

这时我们再次访问该地址，但是我们不等它返回结果，就关闭浏览器，或者使用curl 访问时，按ctrl+c 结束，这时在服务端依然还是可以得到上面的结果，也就说明，当调用SayHello 函数时，这个函数并没有意识到客户端已经关闭了，还在傻傻的进行处理。

作为服务端，当然是希望可以意识到客户端已经关闭了，这样它就不用再傻傻地进行接下来的计算工作。

这时我们就可以利用context 来检查客户端是否关闭了，http.Request结构体中有个ctx，这个就可以感知客户端是否已经关闭。 我们再改下代码

func SayHello(w http.ResponseWriter, r *http.Request) {
	fmt.Println("visit SayHello")
	ctx := r.Context()
	select {
	case <-ctx.Done():
		fmt.Println("client closed....")
	case <-time.After(5 * time.Second):
		n, err := w.Write([]byte(string("hello world")))
		fmt.Printf("send %d bytes\n", n)
		if err != nil {
			fmt.Println(err)
		}
	}
}

经过上面的代码，当客户端关闭了以后，在select流程中就会走到case <-ctx.Done() 流程，这样就不会走time.After流程，这样我们就做到了服务端可以意识到客户端关闭的效果。

以为这样就完了吗？ 显然太简单了! time.After()也只是一个玩具代码，它其实什么都没有做，只是瞎等了5秒钟，但是这个函数其实也是执行了，也是返回了，只是在流程上并没有执行它的结果而已，我们可以将time.After想象一个实际的操作，如果数据库查询，文件操作等耗时操作.

type longtask struct {
	C chan string
}

func (l longtask) longTask() {
	fmt.Println("long task start")
	time.Sleep(5 * time.Second)
	fmt.Println("long task stop")
	l.C <- "hello wrold"
}

func SayHello(w http.ResponseWriter, r *http.Request) {
	fmt.Println("visit SayHello")
	ctx := r.Context()
	lt := longtask{C: make(chan string)}
	go lt.longTask()
	select {
	case <-ctx.Done():
		fmt.Println("client closed....")
	case <-time.After(6 * time.Second):
		fmt.Println("req timeout")
	case s := <-lt.C:
		n, err := w.Write([]byte(s))
		fmt.Printf("send %d bytes\n", n)
		if err != nil {
			fmt.Println(err)
		}
	}
}

上面的代码，将之前的time.After改为一个具体的具体的函数longTask, 此时如果因为是通过 go longTask()方式调用，相当于又启了一个子协程，现在服务端打印出

visit SayHello
long task start
client closed....
long task stop

也就是longTask中的函数还是全部执行了，虽然sayHello 函数中可以实现客户端关闭它就结束，但是通过sayHello 启的协程还是意识不到。 现在的调用链为

其中，http.Request中的ctx 只在sayHello 中，并没有传到longTask中，所以如果想要longTask中也跟着客户端的关闭而关闭，则需要将ctx传到longTask中，并且使用selct 进行监听 修改一下longTask的实现, 在main函数调用的时候，把ctx也传过去。

func (l longtask) longTask(ctx context.Context) {
	fmt.Println("long task start")
	select {
	case <-ctx.Done():
		fmt.Println("in longtask receive stop")
	case <-time.After(5 * time.Second):
		fmt.Println("long task stop")
		l.C <- "hello wrold"
	}
}

func main(){
    ...
    ...
    go lt.longTask(ctx)
    ...
    ...
    
}

这时longTask 函数就可以自动关闭了。

细心的你可能又发现了， 你这种实现这不又回到最开始的time.Sleep 了吗？ 是的，现在其实本质上又回到了最开始的状态了。

上面的代码我们可以将其简化，其实只需要最后一调用的那个函数进行select判断即可

package main

import (
	"context"
	"fmt"
	"net/http"
	"time"
)

type longtask struct {
	C chan string
}

func (l longtask) longTask(ctx context.Context) string {
	fmt.Println("long task start")
	select {
	case <-ctx.Done():
		fmt.Println("long task receive stop msg")
		return ""
	case <-time.After(5 * time.Second):
		fmt.Println("long task stop")
		return "hello wrold"
	}
}

func SayHello(w http.ResponseWriter, r *http.Request) {
	fmt.Println("visit SayHello")
	ctx := r.Context()
	lt := longtask{C: make(chan string)}
	result := lt.longTask(ctx)
	if ctx.Err() != nil {
		fmt.Println(ctx.Err())
	} else {
		w.Write([]byte(result))
	}
    // 以下的select 就可以不用了
	// select {
	// case <-ctx.Done():
	// 	fmt.Println("client closed....")
	// case <-time.After(6 * time.Second):
	// 	fmt.Println("req timeout")
	// case s := <-lt.C:
	// 	n, err := w.Write([]byte(s))
	// 	fmt.Printf("send %d bytes\n", n)
	// 	if err != nil {
	// 		fmt.Println(err)
	// 	}
	// }
}

func main() {
	rdb.Ping(context.Background())
	http.HandleFunc("/", SayHello)
	http.ListenAndServe("127.0.0.1:8000", nil)

}

这也是很多第三方库第一个参数是个ctx的原因， 有了这个ctx， 我们只需要最后判断一个该ctx.Err 是否为空即可，在ctx传递的过程中，只需要最后调用的函数进行select 即可。

要使用ctx 需要满足以后几点

1. ctx 需要在协程和其子协程中传递
2. 在调用链的最后需要使用select 进行判断