golang_context

部分初学者使用web框架时会发现每个handler都会传入一个context，但并没有深究他的作用，甚至直接传入一个context.Background()草草了事。本文结合我在工作中的实际场景，盘点下context的使用

1. 控制goroutine

package main

import (
	"sync"
)

var wg sync.WaitGroup

func main() {
	wg.Add(1)
	go func() {
		defer wg.Done()
		daemon1()
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()
		daemon2()
	}()

	wg.Wait()
}

func daemon1() {
	// 开启8080端口tcp监听，并处理连接
}

func daemon2() {
	// 开启8080端口udp监听，并处理连接
}

如上，我们在main函数中启动了两个守护协程，并使用WaitGroup等待goroutine结束，但由于daemon中监听并处理，一般有for {}结构，无法退出。当我们退出主协程，由于等待，会导致我们的程序卡死

1.1 ctx.Done

我们的诉求是在主协程退出时，子协程需要一起退出。代码做出如下修改

package main

import (
	"context"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

var wg sync.WaitGroup

func main() {
	ctx, cancel := context.WithCancel(context.Background())
	// 监听系统中断信号
	exitCh := make(chan os.Signal, 1)
	signal.Notify(exitCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-exitCh
		cancel()
	}()

	wg.Add(1)
	go func(ctx context.Context) {
		defer wg.Done()
		daemon1(ctx)
	}(ctx)

	wg.Add(1)
	go func(ctx context.Context) {
		defer wg.Done()
		daemon2(ctx)
	}(ctx)

	wg.Wait()
	log.Println("main 结束...")
}

func daemon1(ctx context.Context) {
	// TODO: 开启8080端口tcp监听，并处理连接

	for {
		select {
		case <-ctx.Done():
			// 关闭连接和listener
			log.Println("daemon1 结束...")
			return
		default:
		}
		// 处理连接
	}
}

func daemon2(ctx context.Context) {
	// TODO: 开启8080端口udp监听，并处理连接
	for {
		select {
		case <-ctx.Done():
			// 关闭连接和listener
			log.Println("daemon2 结束...")
			return
		default:
		}
		// 处理连接
	}
}

Tips: 在使用context.WithTimeout和context.withDeadline时，在超时后虽然ctx.Done会触发，但是这个channel并不会关闭，我们需要手动调用cancel()关闭这个ctx.Done()这个channel

2. 携带参数

package main

import (
	"context"
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"
)

var wg sync.WaitGroup

func main() {
	ctx, cancel := context.WithCancel(context.Background())

	ctx = context.WithValue(ctx, "TCPHandlerName", "tcp-daemon")
	ctx = context.WithValue(ctx, "UDPHandlerName", "udp-daemon")

	// 监听系统中断信号
	exitCh := make(chan os.Signal, 1)
	signal.Notify(exitCh, syscall.SIGINT, syscall.SIGTERM)
	go func() {
		<-exitCh
		cancel()
	}()

	wg.Add(1)
	go func(ctx context.Context) {
		defer wg.Done()
		daemon1(ctx)
	}(ctx)

	wg.Add(1)
	go func(ctx context.Context) {
		defer wg.Done()
		daemon2(ctx)
	}(ctx)

	wg.Wait()
	log.Println("main 结束...")
}

func daemon1(ctx context.Context) {
	// TODO: 开启8080端口tcp监听，并处理连接

	for {
		select {
		case <-ctx.Done():
			// 关闭连接和listener
			log.Println(ctx.Value("TCPHandlerName"), " 结束...")
			return
		default:
		}
		// 处理连接
	}
}

func daemon2(ctx context.Context) {
	// TODO: 开启8080端口udp监听，并处理连接
	for {
		select {
		case <-ctx.Done():
			// 关闭连接和listener
			log.Println(ctx.Value("UDPHandlerName"), " 结束...")
			return
		default:
		}
		// 处理连接
	}
}

Tips: context的valule只能向子ctx传递，并且在网络中无法传递，如果您正在使用grpc，请考虑使用metadata