Golang - client 优雅关闭 grpc server 流

#grpc

2025-02-18 09:40:00

最近遇到一个问题，节点 A 和节点 B 之间通过 grpc stream 流来同步信息，且有多条 stream 流，期望其中某条 stream 流断开时，将其余的 stream 流也断开，于是有了今天的故事。

client 主动断开长连接

proto 文件如下：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18

syntax = "proto3";

package hello;
option go_package="./hello";

message HelloReq {
  string name = 1;
}

message HelloResp {
  string message = 1;
}

service Hello {
  rpc sayHello(HelloReq) returns (HelloResp);

  rpc sayHelloStream(HelloReq) returns (stream HelloResp);
}

sayHelloStream 为 server 流，对应的 client 如下：

1
2
3
4

type HelloClient interface {
	SayHello(ctx context.Context, in *HelloReq, opts ...grpc.CallOption) (*HelloResp, error)
	SayHelloStream(ctx context.Context, in *HelloReq, opts ...grpc.CallOption) (Hello_SayHelloStreamClient, error)
}

这里，ctx 可用于优雅关闭，举个例子：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

func sayHelloStream() {
	conn, _ := createConn()
	if conn == nil {
		return
	}

	ctx := context.Background()
	ctx, cancel := context.WithCancel(ctx)

	defer func() {
		log.Default().Println("close conn in defer")
		_ = conn.Close()
		// cancel 可以多次调用，内部有锁，仅第一次调用有效，后续调用 return
		cancel()
	}()

	helloClient := hello.NewHelloClient(conn)
	client, err := helloClient.SayHelloStream(ctx, &hello.HelloReq{Name: "world"})
	if err != nil {
		log.Default().Println(err.Error())
		return
	}

	go func() {
		time.Sleep(5 * time.Second)
		log.Default().Println("call cancel...")
		// cancel 可以多次调用，内部有锁，仅第一次调用有效，后续调用 return
		cancel()
	}()

	for {
		resp, err2 := client.Recv()
		if err2 != nil {
			if ctx.Err() == nil || !errors.Is(ctx.Err(), context.Canceled) {
				log.Default().Println(err2.Error())
			}
			log.Default().Println("exit receive")
			break
		}
		log.Default().Printf("received %+v", resp)
	}
	log.Default().Println("exit...")
}

我们开启了一个协程，5 秒后调用了 ctx 的 cancel 方法，此时原本阻塞的 client.Recv() 会返回，且 err2 为 context.Canceled，而 server 端，ctx 的 done 会被触发；

server 感知 client 退出 

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43

func (HelloServer) SayHelloStream(req *hello.HelloReq, server hello.Hello_SayHelloStreamServer) error {
	data := make(chan *hello.HelloResp, 10)
	closed := make(chan struct{}, 1)

	go func() {
		data <- &hello.HelloResp{Message: "Hello World"}
		ticker := time.NewTicker(30 * time.Second)
		defer ticker.Stop()
		for {
			select {
			case <-ticker.C:
				data <- &hello.HelloResp{Message: "Hello World"}
			case <-closed:
				return
			}
		}
	}()

	ctx := server.Context()

	defer func() {
		closed <- struct{}{}
	}()

	for {
		select {
		case <-ctx.Done():
			log.Default().Println("done...")
			return nil
		case resp := <-data:
			err := server.Send(resp)
			if err != nil {
				ctxErr := ctx.Err()
				if ctxErr != nil && errors.Is(ctxErr, context.Canceled) {
					log.Default().Println("exit stream")
					return nil
				}
				log.Default().Println(err.Error())
				return err
			}
		}
	}
}

server 端代码如上，最后的 for 循环监听 ctx.Done 以及 data 两个 Chanel，若 ctx.Done 发生则返回 nil，若有 data 则发送给 client，若发送失败，且 ctx.Err 为 context.Canceled 则返回 nil，否则返回 err，server 在退出时（defer 里面）向 closed Chanel 发送了消息，前面的协程原本是用于制造数据，但在收到 closed 消息后也会退出，从而避免该协程泄漏；

比起 client 直接关闭 conn，client 使用 context 的 cancel 方法会更平滑，当然，server 监听 client 的消息，client 调用 SendAndClose 也 OK，但这 server 需要额外的协程去处理消息接收，若 server 本身不接收 client 消息，则浪费，若 server 本身也接收 client 消息则 OK；

server 如需主动关闭，只需要退出循环即可，比如这里退出 for 循环，并返回 nil 即可；

最后更新于
再聊 sshhugo 主题 hextra 定制记录