1 package ratelimits 2 3 import ( 4 "context" 5 "fmt" 6 "sync" 7 "time" 8 ) 9 10 // ErrBucketNotFound indicates that the bucket was not found. 11 var ErrBucketNotFound = fmt.Errorf("bucket not found") 12 13 // source is an interface for creating and modifying TATs. 14 type source interface { 15 // Set stores the TAT at the specified bucketKey (formatted as 'name:id'). 16 // Implementations MUST ensure non-blocking operations by either: 17 // a) applying a deadline or timeout to the context WITHIN the method, or 18 // b) guaranteeing the operation will not block indefinitely (e.g. via 19 // the underlying storage client implementation). 20 Set(ctx context.Context, bucketKey string, tat time.Time) error 21 22 // Get retrieves the TAT associated with the specified bucketKey (formatted 23 // as 'name:id'). Implementations MUST ensure non-blocking operations by 24 // either: 25 // a) applying a deadline or timeout to the context WITHIN the method, or 26 // b) guaranteeing the operation will not block indefinitely (e.g. via 27 // the underlying storage client implementation). 28 Get(ctx context.Context, bucketKey string) (time.Time, error) 29 30 // Delete removes the TAT associated with the specified bucketKey (formatted 31 // as 'name:id'). Implementations MUST ensure non-blocking operations by 32 // either: 33 // a) applying a deadline or timeout to the context WITHIN the method, or 34 // b) guaranteeing the operation will not block indefinitely (e.g. via 35 // the underlying storage client implementation). 36 Delete(ctx context.Context, bucketKey string) error 37 } 38 39 // inmem is an in-memory implementation of the source interface used for 40 // testing. 41 type inmem struct { 42 sync.RWMutex 43 m map[string]time.Time 44 } 45 46 func newInmem() *inmem { 47 return &inmem{m: make(map[string]time.Time)} 48 } 49 50 func (in *inmem) Set(_ context.Context, bucketKey string, tat time.Time) error { 51 in.Lock() 52 defer in.Unlock() 53 in.m[bucketKey] = tat 54 return nil 55 } 56 57 func (in *inmem) Get(_ context.Context, bucketKey string) (time.Time, error) { 58 in.RLock() 59 defer in.RUnlock() 60 tat, ok := in.m[bucketKey] 61 if !ok { 62 return time.Time{}, ErrBucketNotFound 63 } 64 return tat, nil 65 } 66 67 func (in *inmem) Delete(_ context.Context, bucketKey string) error { 68 in.Lock() 69 defer in.Unlock() 70 delete(in.m, bucketKey) 71 return nil 72 } 73