diff --git a/README.md b/README.md index 9f0ff98..2cce014 100644 --- a/README.md +++ b/README.md @@ -1,2 +1,46 @@ # rltransport The RoundTripper which limits the number of concurrent requests. + +## examples + +```golang +package main + +import ( + "fmt" + "net/http" + "time" + + "github.com/HatsuneMiku3939/rltransport" + + "golang.org/x/time/rate" +) + +const ( + // TestBurstSize is the default value for the rate limiter's burst size. + TestBurstSize = 10 + // TestRefillRate is the default value for the rate limiter's refill rate. + TestRefillRate = 1.0 + // TestURL is the URL to use for testing. + TestHost = "http://localhost:8080/" +) + +func main() { + // Create a "tocket bucket" limiter with a burst size of 10 and a refill rate of 1.0/sec. + limiter := rate.NewLimiter(TestRefillRate, TestBurstSize) + + // Create a new http.Client with the limiter. + client := &http.Client{ + Transport: &rltransport.RoundTripper{ + Limiter: limiter, + }, + } + + // Make a request to the server. + // First 10 requests will be sented immadiately, after that it will be sented by 1.0 req/sec. + for i := 0; i < 20; i++ { + res, _ := client.Get(TestHost) + fmt.Printf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05"), res.Status) + } +} +``` diff --git a/example/go.mod b/example/go.mod new file mode 100644 index 0000000..29d5271 --- /dev/null +++ b/example/go.mod @@ -0,0 +1,10 @@ +module github.com/HatsuneMiku3939/rltransport/example + +go 1.17 + +replace github.com/HatsuneMiku3939/rltransport => ../ + +require ( + github.com/HatsuneMiku3939/rltransport v0.0.0-00010101000000-000000000000 + golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 +) diff --git a/example/go.sum b/example/go.sum new file mode 100644 index 0000000..e3832da --- /dev/null +++ b/example/go.sum @@ -0,0 +1,2 @@ +golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 h1:M73Iuj3xbbb9Uk1DYhzydthsj6oOd6l9bpuFcNoUvTs= +golang.org/x/time v0.0.0-20220224211638-0e9765cccd65/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= diff --git a/example/main.go b/example/main.go new file mode 100644 index 0000000..0b208ed --- /dev/null +++ b/example/main.go @@ -0,0 +1,61 @@ +package main + +import ( + "fmt" + "net/http" + "time" + + "github.com/HatsuneMiku3939/rltransport" + + "golang.org/x/time/rate" +) + +const ( + // TestBurstSize is the default value for the rate limiter's burst size. + TestBurstSize = 10 + // TestRefillRate is the default value for the rate limiter's refill rate. + TestRefillRate = 1.0 + // TestURL is the URL to use for testing. + TestHost = "http://localhost:8080/" +) + +func main() { + // Create a "tocket bucket" limiter with a burst size of 10 and a refill rate of 1.0/sec. + limiter := rate.NewLimiter(TestRefillRate, TestBurstSize) + + // Create a new http.Client with the limiter. + client := &http.Client{ + Transport: &rltransport.RoundTripper{ + Limiter: limiter, + }, + } + + // Make a request to the server. + // First 10 requests will be sented immadiately, after that it will be sented by 1.0 req/sec. + for i := 0; i < 20; i++ { + res, _ := client.Get(TestHost) + fmt.Printf("[%s] %s\n", time.Now().Format("2006-01-02 15:04:05"), res.Status) + } + + // Will be printed: + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:09] 200 OK + // [2022-04-06 20:11:10] 200 OK ## <-- First 10 requests will be sented immadiately. + // [2022-04-06 20:11:11] 200 OK + // [2022-04-06 20:11:12] 200 OK + // [2022-04-06 20:11:13] 200 OK + // [2022-04-06 20:11:14] 200 OK + // [2022-04-06 20:11:15] 200 OK + // [2022-04-06 20:11:16] 200 OK + // [2022-04-06 20:11:17] 200 OK + // [2022-04-06 20:11:18] 200 OK + // [2022-04-06 20:11:19] 200 OK +} diff --git a/go.mod b/go.mod new file mode 100644 index 0000000..3cfe595 --- /dev/null +++ b/go.mod @@ -0,0 +1,5 @@ +module github.com/HatsuneMiku3939/rltransport + +go 1.17 + +require golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 diff --git a/go.sum b/go.sum new file mode 100644 index 0000000..e3832da --- /dev/null +++ b/go.sum @@ -0,0 +1,2 @@ +golang.org/x/time v0.0.0-20220224211638-0e9765cccd65 h1:M73Iuj3xbbb9Uk1DYhzydthsj6oOd6l9bpuFcNoUvTs= +golang.org/x/time v0.0.0-20220224211638-0e9765cccd65/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ= diff --git a/limiter.go b/limiter.go new file mode 100644 index 0000000..3f800ee --- /dev/null +++ b/limiter.go @@ -0,0 +1,19 @@ +package rltransport + +import ( + "context" +) + +// Limiter is an interface that rate limiter must implement. +type Limiter interface { + // Wait blocks until the request can be sent. + Wait(ctx context.Context) error +} + +// unlimitedLimiter is a limiter that always allows requests to be sent. +type unlimitedLimiter struct{} + +// Wait always success. +func (l *unlimitedLimiter) Wait(ctx context.Context) error { + return nil +} diff --git a/roundtripper.go b/roundtripper.go new file mode 100644 index 0000000..83529e0 --- /dev/null +++ b/roundtripper.go @@ -0,0 +1,48 @@ +package rltransport + +import ( + "net/http" + "sync" +) + +// RoundTripper implements the http.RoundTripper interface. +type RoundTripper struct { + // once ensures that the logic to initialize the default client runs at + // most once, in a single thread. + once sync.Once + + // Limiter is used to rate limit the number of requests that can be made + // to the underlying client. + Limiter Limiter + + // Transport is the underlying RoundTripper that will be used to make + // the actual HTTP requests. + Transport http.RoundTripper +} + +// RoundTrip satisfies the http.RoundTripper interface. +func (rt *RoundTripper) RoundTrip(req *http.Request) (*http.Response, error) { + // Ensure that the Transport is initialized. + rt.once.Do(rt.init) + + // Wait for the rate limiter within context of the request (if limiter is given). + if err := rt.Limiter.Wait(req.Context()); err != nil { + return nil, err + } + + // Execute the request. + resp, err := rt.Transport.RoundTrip(req) + + return resp, err +} + +// init initializes the underlying transport. +func (rt *RoundTripper) init() { + if rt.Transport == nil { + rt.Transport = http.DefaultTransport + } + + if rt.Limiter == nil { + rt.Limiter = &unlimitedLimiter{} + } +}