Go语言中如何在gRPC stats的HandleRPC合并不同事件输出

要实现调用结束时把响应Metadata和gRPC统计信息合并到同一个callResult对象里，你可以按以下步骤修改代码：

1. 扩展callResult结构体 #

首先给callResult新增字段来存储响应Metadata：

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

type callResult struct {
    Error     error
    Status    string
    Duration  time.Duration
    EndTime   time.Time
    Metadata  metadata.MD // 新增字段用于保存响应Metadata
}

2. 修改statsHandler结构体 #

在处理器结构体中添加字段，用来临时存储InHeader阶段捕获到的Metadata，同时用锁保证并发安全：

type statsHandler struct {
    lock    sync.RWMutex
    ignore  bool
    results chan<- *callResult
    hasLog  bool
    log     *zap.Logger
    id      string
    header  metadata.MD // 新增字段临时存储响应Metadata
}

3. 在InHeader事件中捕获Metadata #

当触发*stats.InHeader事件时，把响应Header存到statsHandler的header字段中：

case *stats.InHeader:
    if rs.Client {
        c.lock.Lock()
        c.header = rs.Header // 存储响应Metadata
        c.lock.Unlock()
    }

4. 在End事件中合并数据到callResult #

当调用结束触发*stats.End事件时，取出之前存储的Metadata，和统计信息一起放入callResult：

case *stats.End:
    ign := false
    c.lock.RLock()
    ign = c.ignore
    header := c.header // 取出存储的Metadata
    c.lock.RUnlock()

    if !ign {
        duration := rs.EndTime.Sub(rs.BeginTime)

        var st string
        s, ok := status.FromError(rs.Error)
        if ok {
            st = s.Code().String()
        }

        // 合并Metadata到callResult
        c.results <- &callResult{
            Error:    rs.Error,
            Status:   st,
            Duration: duration,
            EndTime:  rs.EndTime,
            Metadata: header,
        }

        if c.hasLog {
            c.log.Debugw("Received RPC Stats",
                "statsID", c.id, "code", st, "error", rs.Error,
                "duration", duration, "metadata", header, "stats", rs)
        }
    }

完整修改后的代码 #

import (
    "context"
    "time"

    "google.golang.org/grpc/metadata"
    "google.golang.org/grpc/stats"
    "google.golang.org/grpc/status"
    "go.uber.org/zap"
    "sync"
)

type callResult struct {
    Error     error
    Status    string
    Duration  time.Duration
    EndTime   time.Time
    Metadata  metadata.MD
}

type statsHandler struct {
    lock    sync.RWMutex
    ignore  bool
    results chan<- *callResult
    hasLog  bool
    log     *zap.Logger
    id      string
    header  metadata.MD
}

// HandleRPC implements per-RPC tracing and stats instrumentation.
func (c *statsHandler) HandleRPC(ctx context.Context, rs stats.RPCStats) {
    switch rs := rs.(type) {
    case *stats.InHeader:
        if rs.Client {
            c.lock.Lock()
            c.header = rs.Header
            c.lock.Unlock()
        }

    case *stats.End:
        ign := false
        c.lock.RLock()
        ign = c.ignore
        header := c.header
        c.lock.RUnlock()

        if !ign {
            duration := rs.EndTime.Sub(rs.BeginTime)

            var st string
            s, ok := status.FromError(rs.Error)
            if ok {
                st = s.Code().String()
            }

            c.results <- &callResult{
                Error:    rs.Error,
                Status:   st,
                Duration: duration,
                EndTime:  rs.EndTime,
                Metadata: header,
            }

            if c.hasLog {
                c.log.Debugw("Received RPC Stats",
                    "statsID", c.id, "code", st, "error", rs.Error,
                    "duration", duration, "metadata", header, "stats", rs)
            }
        }
    }
}

说明 #

• 同一个statsHandler实例对应单个RPC调用的生命周期，用结构体字段存储Metadata是安全的，加锁是为了避免潜在的并发访问问题
• 确保导入了google.golang.org/grpc/metadata包来处理gRPC Metadata

内容的提问来源于stack exchange，提问作者LST