golang-grpc拦截器

proto/default.proto

syntax="proto3";

// import "google/protobuf/timestamp.proto";

option go_package="./pb";

service ToDoService{
    rpc DoWork(TodoRequest) returns (TodoResponse);
}

enum Week{
    Sunday = 0;
    Monday = 1;
    Tuesday = 2;
    Wednesday = 3;
    Thursday = 4;
    Friday = 5;
    Saturday = 6;
}

message TodoRequest{
    string todo = 1;
    Week week = 2;
    map<string,string> bookMap = 3;
    // google.protobuf.Timestamp doneTime  = 4;
}

message TodoResponse{
    bool done = 1;
}

// go install google.golang.org/protobuf/cmd/protoc-gen-go@v1.28
// go install google.golang.org/grpc/cmd/protoc-gen-go-grpc@v1.2
// export PATH="$PATH:$(go env GOPATH)/bin"
// protoc --go_out=. --go-grpc_out=. proto/default.proto

server.go

package main

import (
	"context"
	"fmt"
	"net"
	"study/pb"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

type ToDo struct {
	pb.UnimplementedToDoServiceServer
}

func (t *ToDo) DoWork(ctx context.Context, req *pb.TodoRequest) (*pb.TodoResponse, error) {
	fmt.Println("开始处理客户端请求")
	md, ok := metadata.FromIncomingContext(ctx)
	if !ok {
		fmt.Println("没有metadata")
	}
	fmt.Println("开始遍历metadata")
	for k, v := range md {
		fmt.Printf("server recv %s:%s\n", k, v)
	}
	fmt.Println("结束遍历metadata")
	fmt.Println("md.name:", md.Get("name"))
	fmt.Println("md.name-bin:", md.Get("name-bin"))
	time.Sleep(2 * time.Second)
	fmt.Println(req.Todo + "已完成！")
	return &pb.TodoResponse{Done: true}, nil
}

func MyInterceptor(ctx context.Context, req any, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
	fmt.Println("服务端进入MyInterceptor")
	now := time.Now()
	fmt.Println("服务端开始执行DoWork方法")
	resp, err = handler(ctx, req)
	d := time.Since(now)
	fmt.Printf("服务端执行时间：%d\n", d.Milliseconds())
	fmt.Println("服务端离开MyInterceptor")
	return
}

func main() {
	serverOption := grpc.UnaryInterceptor(MyInterceptor)
	server := grpc.NewServer(serverOption)
	pb.RegisterToDoServiceServer(server, &ToDo{})
	listen, err := net.Listen("tcp", ":8080")
	if err != nil {
		panic(err)
	}
	err = server.Serve(listen)
	if err != nil {
		panic(err)
	}
}

client.go

package main

import (
	"context"
	"fmt"
	"log"
	"study/pb"
	"time"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/metadata"
)

func main() {
	clientInterceptor := func(ctx context.Context, method string, req, reply any, cc *grpc.ClientConn, invoker grpc.UnaryInvoker, opts ...grpc.CallOption) error {
		fmt.Println("客户端进入clientInterceptor")
		now := time.Now()
		fmt.Println("客户端开始执行DoWork方法")
		err := invoker(ctx, method, req, reply, cc, opts...)
		d := time.Since(now)
		fmt.Printf("客户端执行时间：%d\n", d.Milliseconds())
		fmt.Println("客户端离开clientInterceptor")
		return err
	}
	opt := grpc.WithUnaryInterceptor(clientInterceptor)
	conn, err := grpc.Dial("127.0.0.1:8080", grpc.WithTransportCredentials(insecure.NewCredentials()), opt)
	if err != nil {
		log.Fatalf("did not connect: %v", err)
	}
	defer conn.Close()

	client := pb.NewToDoServiceClient(conn)

	md1 := metadata.New(map[string]string{
		"name": "golang",
	})
	// md1 := metadata.Pairs("name-bin", "golang语言")

	ctx := metadata.NewOutgoingContext(context.Background(), md1)
	res, err := client.DoWork(ctx, &pb.TodoRequest{Todo: "复习"})
	if err != nil {
		log.Fatalf("could not DoWork: %v", err)
		return
	}
	log.Printf("Server Response: %v", res.Done)
}