Lennard Brinkhaus
31bd0cf639
Co-authored-by: Fenpaws <soxx@fenpa.ws> Reviewed-on: anthrove/e621-to-graph#13 Reviewed-by: SoXX <fenpaws@noreply.localhost> Co-authored-by: Lennard Brinkhaus <lennard.brinkhaus@noreply.localhost> Co-committed-by: Lennard Brinkhaus <lennard.brinkhaus@noreply.localhost>
62 lines
1.2 KiB
Go
62 lines
1.2 KiB
Go
package e621
|
|
|
|
import (
|
|
"context"
|
|
"e621_to_neo4j/utils"
|
|
log "github.com/sirupsen/logrus"
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
type Scheduler struct {
|
|
queue utils.Queue
|
|
limiter *rate.Limiter
|
|
}
|
|
|
|
func NewScheduler() *Scheduler {
|
|
scheduler := &Scheduler{
|
|
queue: utils.NewQueue(),
|
|
limiter: nil,
|
|
}
|
|
go scheduler.StartExecutionHandler()
|
|
return scheduler
|
|
}
|
|
|
|
func Schedule[T utils.DataType](s *Scheduler, t utils.Task[T], username string, apiKey string) func() (T, error) {
|
|
channel := make(chan utils.DataResponse[T])
|
|
schedulerTask := NewSchedulerTaskImpl[T](t, channel, username, apiKey)
|
|
log.Debug("Psuh task")
|
|
err := s.queue.Push(schedulerTask)
|
|
log.Debug("Element pushed")
|
|
|
|
if err != nil {
|
|
return func() (T, error) {
|
|
var nil T
|
|
return nil, err
|
|
}
|
|
}
|
|
return func() (T, error) {
|
|
data := <-channel
|
|
return data.Data, data.Error
|
|
}
|
|
}
|
|
|
|
func (s *Scheduler) SetLimiter(limiter *rate.Limiter) {
|
|
s.limiter = limiter
|
|
}
|
|
|
|
func (s *Scheduler) StartExecutionHandler() {
|
|
for {
|
|
if s.limiter != nil {
|
|
s.limiter.Wait(context.Background())
|
|
}
|
|
s.queue.WaitForElement()
|
|
log.Debug("element found")
|
|
task, err := s.queue.Pop()
|
|
if err != nil {
|
|
continue
|
|
}
|
|
ExecuteGetAPIRequest(task)
|
|
}
|
|
|
|
}
|