rate_limit_round_tripper.go raw

   1  // Copyright 2021-2023 The sacloud/go-http authors
   2  //
   3  // Licensed under the Apache License, Version 2.0 (the "License");
   4  // you may not use this file except in compliance with the License.
   5  // You may obtain a copy of the License at
   6  //
   7  //      http://www.apache.org/licenses/LICENSE-2.0
   8  //
   9  // Unless required by applicable law or agreed to in writing, software
  10  // distributed under the License is distributed on an "AS IS" BASIS,
  11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12  // See the License for the specific language governing permissions and
  13  // limitations under the License.
  14  
  15  package http
  16  
  17  import (
  18  	"net/http"
  19  	"sync"
  20  
  21  	"go.uber.org/ratelimit"
  22  )
  23  
  24  // RateLimitRoundTripper 秒間アクセス数を制限するためのhttp.RoundTripper実装
  25  type RateLimitRoundTripper struct {
  26  	// Transport 親となるhttp.RoundTripper、nilの場合http.DefaultTransportが利用される
  27  	Transport http.RoundTripper
  28  	// RateLimitPerSec 秒あたりのリクエスト数
  29  	RateLimitPerSec int
  30  
  31  	once      sync.Once
  32  	rateLimit ratelimit.Limiter
  33  }
  34  
  35  // RoundTrip http.RoundTripperの実装
  36  func (r *RateLimitRoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
  37  	r.once.Do(func() {
  38  		r.rateLimit = ratelimit.New(r.RateLimitPerSec)
  39  	})
  40  	if r.Transport == nil {
  41  		r.Transport = http.DefaultTransport
  42  	}
  43  
  44  	r.rateLimit.Take()
  45  	return r.Transport.RoundTrip(req)
  46  }
  47