go/src/infra/appengine/cr-rev/backend/gitiles/throttling_client.go

// Copyright 2020 The Chromium Authors
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

package gitiles

import (
	"context"

	"golang.org/x/time/rate"
	"google.golang.org/grpc"

	gitilesProto "go.chromium.org/luci/common/proto/gitiles"
)

// NewThrottlingClient creates REST Gitiles client and consumes limiter quota
// on each API call to Gitiles. If there is no quota left, it blocks until
// there is.
func NewThrottlingClient(c Client, limiter *rate.Limiter) Client {
	return &throttlingClient{
		client:  c,
		limiter: limiter,
	}
}

// throttlingClient is a wrapper around Gitiles client that throttles API calls.
type throttlingClient struct {
	client  Client
	limiter *rate.Limiter
}

func (t *throttlingClient) Log(ctx context.Context, in *gitilesProto.LogRequest, opts ...grpc.CallOption) (*gitilesProto.LogResponse, error) {
	if err := t.limiter.Wait(ctx); err != nil {
		return nil, err
	}
	return t.client.Log(ctx, in, opts...)
}

func (t *throttlingClient) Refs(ctx context.Context, in *gitilesProto.RefsRequest, opts ...grpc.CallOption) (*gitilesProto.RefsResponse, error) {
	if err := t.limiter.Wait(ctx); err != nil {
		return nil, err
	}
	return t.client.Refs(ctx, in, opts...)
}

func (t *throttlingClient) Archive(ctx context.Context, in *gitilesProto.ArchiveRequest, opts ...grpc.CallOption) (*gitilesProto.ArchiveResponse, error) {
	if err := t.limiter.Wait(ctx); err != nil {
		return nil, err
	}
	return t.client.Archive(ctx, in, opts...)
}

func (t *throttlingClient) DownloadFile(ctx context.Context, in *gitilesProto.DownloadFileRequest, opts ...grpc.CallOption) (*gitilesProto.DownloadFileResponse, error) {
	if err := t.limiter.Wait(ctx); err != nil {
		return nil, err
	}
	return t.client.DownloadFile(ctx, in, opts...)
}

func (t *throttlingClient) Projects(ctx context.Context, in *gitilesProto.ProjectsRequest, opts ...grpc.CallOption) (*gitilesProto.ProjectsResponse, error) {
	if err := t.limiter.Wait(ctx); err != nil {
		return nil, err
	}
	return t.client.Projects(ctx, in, opts...)
}