add generic and auth rate limiting to ntg

master
Bel LaPointe 2022-01-11 22:56:03 -05:00
parent e546034c26
commit b54be3c32c
2 changed files with 8 additions and 2 deletions

View File

@ -6,7 +6,11 @@ import (
"golang.org/x/time/rate" "golang.org/x/time/rate"
) )
var limiter = rate.NewLimiter(rate.Limit(0.3), 1) // once per minute
var authlimiter = rate.NewLimiter(rate.Limit(1.0/60.0), 1)
// thrice per minute
var limiter = rate.NewLimiter(rate.Limit(1.0/20.0), 1)
type Broker interface { type Broker interface {
Search([]config.State) ([]Job, error) Search([]config.State) ([]Job, error)

View File

@ -2,6 +2,7 @@ package broker
import ( import (
"bytes" "bytes"
"context"
"encoding/json" "encoding/json"
"errors" "errors"
"fmt" "fmt"
@ -100,7 +101,6 @@ func (ntg NTGVision) search(states []config.State) (io.ReadCloser, error) {
} }
func (ntg NTGVision) refreshAuth() error { func (ntg NTGVision) refreshAuth() error {
time.Sleep(time.Minute * 2) // TODO
b, _ := json.Marshal(map[string]string{ b, _ := json.Marshal(map[string]string{
"username": config.Get().Brokers.NTG.Username, "username": config.Get().Brokers.NTG.Username,
"password": config.Get().Brokers.NTG.Password, "password": config.Get().Brokers.NTG.Password,
@ -110,6 +110,7 @@ func (ntg NTGVision) refreshAuth() error {
return err return err
} }
setNTGHeaders(request) setNTGHeaders(request)
authlimiter.Wait(context.Background())
resp, err := http.DefaultClient.Do(request) resp, err := http.DefaultClient.Do(request)
if err != nil { if err != nil {
return err return err
@ -140,6 +141,7 @@ func (ntg NTGVision) _search(states []config.State) (io.ReadCloser, error) {
return nil, err return nil, err
} }
c := http.Client{Timeout: time.Minute} c := http.Client{Timeout: time.Minute}
limiter.Wait(context.Backgroud())
resp, err := c.Do(request) resp, err := c.Do(request)
if err != nil { if err != nil {
return nil, err return nil, err