通过速率限制保护 API

对于我们要添加到服务器的最后一个拦截器，我们将使用一个速率限制器。更具体地说，我们将使用 golang.org/x/time/rate 包提供的令牌桶速率限制器实现。在本节中，我们不会深入讨论速率限制器是什么或者如何构建一个——这超出了本书的范围，不过，您将看到如何在 gRPC 上下文中使用速率限制器（无论是现成的实现还是自定义的）。

我们需要做的第一件事是获取速率限制器的依赖：

$ go get golang.org/x/time/rate

如果您没有获取过前面的 go-grpc-middleware 依赖，则需要运行以下命令。如果您是按照章节逐步操作的，通常不需要再运行它。

然后，我们获取拦截器的依赖：

$ go get github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/ratelimit

现在，我们将创建一个名为 limit.go 的文件，其中包含我们的逻辑和对 rate.Limiter 的包装。我们创建这样一个包装器，因为我们稍后使用的拦截器要求 Limiter 实现一个名为 Limit 的函数，该函数接受一个 context 作为参数，而 rate.Limiter 本身并没有这样的函数：

package main

import (
	"context"
	"fmt"

	"golang.org/x/time/rate"
)

type simpleLimiter struct {
	limiter *rate.Limiter
}

func (l *simpleLimiter) Limit(_ context.Context) error {
	if !l.limiter.Allow() {
		return fmt.Errorf("reached Rate-Limiting %v", l.limiter.Limit())
	}

	return nil
}

请注意，我们只是检查速率限制器是否允许（或不允许）调用通过。如果不允许，我们返回一个错误；否则，我们返回 nil。

最后，我们需要将 simpleLimiter 注册到拦截器中。我们将创建一个类型为 rate.Limiter 的实例，设置为每秒 2 个令牌（称为 r），以及 4 的突发大小（称为 b）。如果您不清楚这些参数的含义，建议您阅读 Limiter 的文档（ https://pkg.go.dev/golang.org/x/time/rate#Limiter ）：

import (
	//...
	"github.com/grpc-ecosystem/go-grpc-middleware/v2/interceptors/ratelimit"
)

func newGrpcServer(lis net.Listener, srvMetrics *grpcprom.ServerMetrics) (*grpc.Server, error) {
	//...
	limiter := &simpleLimiter{
		limiter: rate.NewLimiter(2, 4),
	}
	opts := []grpc.ServerOption{
		//...
		grpc.ChainUnaryInterceptor(
			ratelimit.UnaryServerInterceptor(limiter),
			//...
		),
		grpc.ChainStreamInterceptor(
			ratelimit.StreamServerInterceptor(limiter),
			//...
		),
		//...
	}
	//...
}

到此为止，我们已经为我们的 API 启用了速率限制。现在，我们可以运行我们的服务器：

$ go run ./server 0.0.0.0:50051 0.0.0.0:50052

metrics server listening at 0.0.0.0:50052
gRPC server listening at 0.0.0.0:50051

然后，我们可以尝试每秒执行超过两次的调用。这个过程应该不难，实际上，您通常可以运行客户端一次，应该就会失败。但是为了确保它失败，您可以多次运行客户端。在 Linux 和 Mac 上，您可以运行以下命令：

$ for i in {1..10}; do go run ./client 0.0.0.0:50051; done

在 Windows（PowerShell）上，您可以运行这个：

$ foreach ($item in 1..10) { go run ./client 0.0.0.0:50051 }

您应该看到一些查询返回响应，然后很快，您将看到以下消息：

rpc error: code = ResourceExhausted desc = /todo.v2.TodoService/UpdateTasks is rejected by grpc_ratelimit middleware, please retry later. reached Rate-Limiting 2

显然，我们的速率设置得非常低，这在生产环境中并不实用。我们选择如此低的速率是为了向您展示如何进行速率限制。在生产环境中，您将有特定的业务需求需要遵循。您需要调整我们展示的代码以匹配这些需求。

Bazel

为了使用 Bazel 运行这个示例，我们需要更新仓库并运行 Gazelle 以将新依赖（golang.org/x/time/rate）导入到我们的库中：

$ bazel run //:gazelle-update-repos
$ bazel run //:gazelle

之后，您应该可以像下面这样运行服务器：

$ bazel run //server:server 0.0.0.0:50051 0.0.0.0:50052

总结来说，我们可以在我们的 gRPC 服务器中集成一个速率限制器。go-grpc-middleware 提供的速率限制拦截器使得添加现成的实现或自定义实现变得非常简单。