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)
|
||
|
|
}
|
||
|
|
|
||
|
|
}
|