46 lines
1.2 KiB
Go
46 lines
1.2 KiB
Go
package grpc_conn
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"google.golang.org/grpc/stats"
|
|
)
|
|
|
|
// 1. 实现 stats.Handler 接口
|
|
type StatsHandler struct{}
|
|
|
|
func (h *StatsHandler) TagRPC(ctx context.Context, info *stats.RPCTagInfo) context.Context {
|
|
// 给 RPC 调用打标签（例如记录方法名）
|
|
return context.WithValue(ctx, "rpc_method", info.FullMethodName)
|
|
}
|
|
|
|
func (h *StatsHandler) HandleRPC(ctx context.Context, s stats.RPCStats) {
|
|
// 处理 RPC 统计信息
|
|
switch t := s.(type) {
|
|
case *stats.Begin:
|
|
fmt.Printf("RPC started: %s\n", ctx.Value("rpc_method"))
|
|
case *stats.End:
|
|
fmt.Printf("RPC finished: %s (duration: %v)\n",
|
|
ctx.Value("rpc_method"), t.EndTime.Sub(t.BeginTime))
|
|
case *stats.InPayload:
|
|
fmt.Printf("Received %d bytes\n", t.WireLength)
|
|
case *stats.OutPayload:
|
|
fmt.Printf("Sent %d bytes\n", t.WireLength)
|
|
}
|
|
}
|
|
|
|
func (h *StatsHandler) TagConn(ctx context.Context, info *stats.ConnTagInfo) context.Context {
|
|
// 给连接打标签
|
|
return ctx
|
|
}
|
|
|
|
func (h *StatsHandler) HandleConn(ctx context.Context, s stats.ConnStats) {
|
|
// 处理连接事件
|
|
switch s.(type) {
|
|
case *stats.ConnBegin:
|
|
fmt.Println("Connection established")
|
|
case *stats.ConnEnd:
|
|
fmt.Println("Connection closed")
|
|
}
|
|
}
|