64 lines
1.7 KiB
Go
64 lines
1.7 KiB
Go
package storage
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"local/dndex/config"
|
|
"local/dndex/storage/entity"
|
|
"sync"
|
|
|
|
"golang.org/x/time/rate"
|
|
)
|
|
|
|
type RateLimitedGraph struct {
|
|
g Graph
|
|
rps int
|
|
limiters *sync.Map
|
|
}
|
|
|
|
func NewRateLimitedGraph(path ...string) RateLimitedGraph {
|
|
return RateLimitedGraph{
|
|
g: NewGraph(path...),
|
|
rps: config.New().RPS,
|
|
limiters: &sync.Map{},
|
|
}
|
|
}
|
|
|
|
func (rlg RateLimitedGraph) limit(ctx context.Context, namespace string) error {
|
|
limiter, ok := rlg.limiters.Load(namespace)
|
|
if !ok {
|
|
config := config.New()
|
|
limiter = rate.NewLimiter(rate.Limit(config.RPS), config.RPS)
|
|
rlg.limiters.Store(namespace, limiter)
|
|
}
|
|
limit, ok := limiter.(*rate.Limiter)
|
|
if !ok {
|
|
return fmt.Errorf("rate limiter is of type %T", limiter)
|
|
}
|
|
return limit.Wait(ctx)
|
|
}
|
|
|
|
func (rlg RateLimitedGraph) Delete(ctx context.Context, namespace string, filter interface{}) error {
|
|
return rlg.g.Delete(ctx, namespace, filter)
|
|
}
|
|
|
|
func (rlg RateLimitedGraph) Insert(ctx context.Context, namespace string, one entity.One) error {
|
|
return rlg.g.Insert(ctx, namespace, one)
|
|
}
|
|
|
|
func (rlg RateLimitedGraph) List(ctx context.Context, namespace string, from ...string) ([]entity.One, error) {
|
|
return rlg.g.List(ctx, namespace, from...)
|
|
}
|
|
|
|
func (rlg RateLimitedGraph) ListCaseInsensitive(ctx context.Context, namespace string, from ...string) ([]entity.One, error) {
|
|
return rlg.g.ListCaseInsensitive(ctx, namespace, from...)
|
|
}
|
|
|
|
func (rlg RateLimitedGraph) Search(ctx context.Context, namespace string, nameContains string) ([]entity.One, error) {
|
|
return rlg.g.Search(ctx, namespace, nameContains)
|
|
}
|
|
|
|
func (rlg RateLimitedGraph) Update(ctx context.Context, namespace string, one entity.One, modify interface{}) error {
|
|
return rlg.g.Update(ctx, namespace, one, modify)
|
|
}
|