如何在Golang gRPC中存储请求中的用户请求计数变量

问题背景 #

需求是统计用户完成Auth接口调用后的请求次数，现有两个gRPC方法：StartTracking用于初始化计数器为0，AllPath用于更新计数器，但不清楚如何安全存储用户的当前计数值。

提供的Protobuf定义 #

service RecTask {
  rpc AllPath(UserToken) returns (UserToken) {
    option (google.api.http) = {
      get: "/*"
    };
  }
  rpc Auth(google.protobuf.Empty) returns (TokenRender) {
    option (google.api.http) = {
      get: "/auth"
    };
  }
  rpc StartTracking(UserToken) returns (UserToken) {
    option (google.api.http) = {
      get: "/auth/{user_token}"
    };
  }
}

message UserToken {
  string user_token = 1;
  int32 user_counter = 2;
}

现有代码片段 #

func (s *RecTask) StartTracking(ctx context.Context, req *pb.UserToken) (*pb.UserToken, error) {
    if contains(tokens, req.UserToken) {
        response := &pb.UserToken{
            UserCounter: 0,
        }
        return response, nil

func (s *RecTask) AllPath(ctx context.Context, req *pb.UserToken) (res *pb.UserToken, err error) {
    response := &pb.UserToken{
        UserCounter: +1,
    }
}

解决方案 #

由于gRPC服务是并发处理请求的，必须使用线程安全的存储结构来保存用户计数器，下面提供两种可行方案：

方案1：使用sync.Map（推荐，无需手动加锁） #

sync.Map是Go标准库提供的并发安全Map，适合高频读写场景。

首先修改服务结构体，添加存储字段：

import (
    "sync"
    pb "your/proto/path"
    "google.golang.org/grpc/codes"
    "google.golang.org/grpc/status"
)

type RecTask struct {
    pb.UnimplementedRecTaskServer
    userCounters sync.Map // key: 字符串类型的user_token, value: int32类型的计数器
}

修改StartTracking方法，初始化用户计数器：

func (s *RecTask) StartTracking(ctx context.Context, req *pb.UserToken) (*pb.UserToken, error) {
    // 验证用户token是否合法（假设contains是你实现的token校验函数）
    if contains(tokens, req.UserToken) {
        // 将用户计数器初始化为0并存储
        s.userCounters.Store(req.UserToken, int32(0))
        return &pb.UserToken{
            UserToken:   req.UserToken,
            UserCounter: 0,
        }, nil
    }
    // token不合法时返回错误
    return nil, status.Errorf(codes.Unauthenticated, "无效的用户token")
}

修改AllPath方法，更新并返回当前计数器：

func (s *RecTask) AllPath(ctx context.Context, req *pb.UserToken) (*pb.UserToken, error) {
    // 从存储中获取当前计数器
    val, ok := s.userCounters.Load(req.UserToken)
    if !ok {
        return nil, status.Errorf(codes.FailedPrecondition, "用户未开启追踪，请先调用StartTracking接口")
    }

    // 类型断言确保是int32类型
    counter, ok := val.(int32)
    if !ok {
        return nil, status.Errorf(codes.Internal, "计数器类型错误")
    }

    // 计数器加1并更新存储
    counter += 1
    s.userCounters.Store(req.UserToken, counter)

    return &pb.UserToken{
        UserToken:   req.UserToken,
        UserCounter: counter,
    }, nil
}

方案2：使用普通Map+sync.Mutex #

如果需要更细粒度的锁控制，可以用普通Map配合互斥锁：

修改服务结构体：

type RecTask struct {
    pb.UnimplementedRecTaskServer
    mu           sync.Mutex // 保护userCounters的互斥锁
    userCounters map[string]int32
}

// 初始化服务实例时创建Map
func NewRecTaskServer() *RecTask {
    return &RecTask{
        userCounters: make(map[string]int32),
    }
}

修改StartTracking方法：

func (s *RecTask) StartTracking(ctx context.Context, req *pb.UserToken) (*pb.UserToken, error) {
    if contains(tokens, req.UserToken) {
        s.mu.Lock()
        defer s.mu.Unlock() // 方法结束时自动释放锁
        s.userCounters[req.UserToken] = 0
        return &pb.UserToken{
            UserToken:   req.UserToken,
            UserCounter: 0,
        }, nil
    }
    return nil, status.Errorf(codes.Unauthenticated, "无效的用户token")
}

修改AllPath方法：

func (s *RecTask) AllPath(ctx context.Context, req *pb.UserToken) (*pb.UserToken, error) {
    s.mu.Lock()
    defer s.mu.Unlock()

    counter, ok := s.userCounters[req.UserToken]
    if !ok {
        return nil, status.Errorf(codes.FailedPrecondition, "用户未开启追踪，请先调用StartTracking接口")
    }

    counter += 1
    s.userCounters[req.UserToken] = counter

    return &pb.UserToken{
        UserToken:   req.UserToken,
        UserCounter: counter,
    }, nil
}

注意事项 #

• 上述方案都是内存存储，服务重启后计数器数据会丢失，如果需要持久化，可以将计数器存储到Redis、MySQL等数据库中。
• 确保contains函数能正确校验用户token的合法性，避免非法用户占用存储资源。

内容的提问来源于stack exchange，提问作者Kacper Łęczyński