From 08df9258a532b653c243e077e82491dbe62ad854 Mon Sep 17 00:00:00 2001 From: Philipp Tanlak Date: Sat, 23 Sep 2023 17:41:57 +0200 Subject: refactor scraper into modules --- modules/ratelimit/ratelimit.go | 54 ++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) create mode 100644 modules/ratelimit/ratelimit.go (limited to 'modules/ratelimit/ratelimit.go') diff --git a/modules/ratelimit/ratelimit.go b/modules/ratelimit/ratelimit.go new file mode 100644 index 0000000..b02f5d5 --- /dev/null +++ b/modules/ratelimit/ratelimit.go @@ -0,0 +1,54 @@ +// This Source Code Form is subject to the terms of the Mozilla Public +// License, v. 2.0. If a copy of the MPL was not distributed with this +// file, You can obtain one at http://mozilla.org/MPL/2.0/. + +package ratelimit + +import ( + "time" + + "github.com/philippta/flyscrape" +) + +func init() { + flyscrape.RegisterModule(new(Module)) +} + +type Module struct { + Rate float64 `json:"rate"` + + ticker *time.Ticker + semaphore chan struct{} +} + +func (m *Module) ID() string { + return "ratelimit" +} + +func (m *Module) OnLoad(v flyscrape.Visitor) { + rate := time.Duration(float64(time.Second) / m.Rate) + + m.ticker = time.NewTicker(rate) + m.semaphore = make(chan struct{}, 1) + + go func() { + for range m.ticker.C { + m.semaphore <- struct{}{} + } + }() +} + +func (m *Module) OnRequest(_ *flyscrape.Request) { + <-m.semaphore +} + +func (m *Module) OnComplete() { + m.ticker.Stop() +} + +var ( + _ flyscrape.Module = (*Module)(nil) + _ flyscrape.OnRequest = (*Module)(nil) + _ flyscrape.OnLoad = (*Module)(nil) + _ flyscrape.OnComplete = (*Module)(nil) +) -- cgit v1.2.3