summaryrefslogtreecommitdiff
path: root/modules/ratelimit
diff options
context:
space:
mode:
authorPhilipp Tanlak <philipp.tanlak@gmail.com>2023-10-11 20:20:30 +0200
committerPhilipp Tanlak <philipp.tanlak@gmail.com>2023-10-11 20:20:30 +0200
commitfb84ca746e92e371161f1e1de3b01a048a9ae979 (patch)
tree5bb8fbb7fd654b241b389697cc46bad00ce2f8b7 /modules/ratelimit
parentcd40ab75f44e9f6ac86beca576a934fd790fc9fb (diff)
Implement file based caching
Diffstat (limited to 'modules/ratelimit')
-rw-r--r--modules/ratelimit/ratelimit.go16
-rw-r--r--modules/ratelimit/ratelimit_test.go8
2 files changed, 14 insertions, 10 deletions
diff --git a/modules/ratelimit/ratelimit.go b/modules/ratelimit/ratelimit.go
index 9588db3..b23cd7a 100644
--- a/modules/ratelimit/ratelimit.go
+++ b/modules/ratelimit/ratelimit.go
@@ -5,6 +5,7 @@
package ratelimit
import (
+ "net/http"
"time"
"github.com/philippta/flyscrape"
@@ -45,11 +46,14 @@ func (m *Module) Provision(v flyscrape.Context) {
}()
}
-func (m *Module) BuildRequest(_ *flyscrape.Request) {
+func (m *Module) AdaptTransport(t http.RoundTripper) http.RoundTripper {
if m.disabled() {
- return
+ return t
}
- <-m.semaphore
+ return flyscrape.RoundTripFunc(func(r *http.Request) (*http.Response, error) {
+ <-m.semaphore
+ return t.RoundTrip(r)
+ })
}
func (m *Module) Finalize() {
@@ -64,7 +68,7 @@ func (m *Module) disabled() bool {
}
var (
- _ flyscrape.RequestBuilder = (*Module)(nil)
- _ flyscrape.Provisioner = (*Module)(nil)
- _ flyscrape.Finalizer = (*Module)(nil)
+ _ flyscrape.TransportAdapter = (*Module)(nil)
+ _ flyscrape.Provisioner = (*Module)(nil)
+ _ flyscrape.Finalizer = (*Module)(nil)
)
diff --git a/modules/ratelimit/ratelimit_test.go b/modules/ratelimit/ratelimit_test.go
index ffd061c..1fe22b1 100644
--- a/modules/ratelimit/ratelimit_test.go
+++ b/modules/ratelimit/ratelimit_test.go
@@ -23,17 +23,17 @@ func TestRatelimit(t *testing.T) {
scraper := flyscrape.NewScraper()
scraper.LoadModule(&starturl.Module{URL: "http://www.example.com"})
scraper.LoadModule(&followlinks.Module{})
- scraper.LoadModule(&ratelimit.Module{
- Rate: 100,
- })
scraper.LoadModule(hook.Module{
AdaptTransportFn: func(rt http.RoundTripper) http.RoundTripper {
return flyscrape.MockTransport(200, `<a href="foo">foo</a>`)
},
- BuildRequestFn: func(r *flyscrape.Request) {
+ ReceiveResponseFn: func(r *flyscrape.Response) {
times = append(times, time.Now())
},
})
+ scraper.LoadModule(&ratelimit.Module{
+ Rate: 100,
+ })
start := time.Now()
scraper.Run()