| package unthrottle |
| |
| import ( |
| "context" |
| |
| "cloud.google.com/go/datastore" |
| "go.skia.org/infra/go/ds" |
| ) |
| |
| // Throttle tracks whether a given roller should be throttled. |
| // TODO(borenet): This should include throttle-enabling as well as disabling, |
| // and it should maintain history like modes and strategies. |
| type Throttle interface { |
| // Unthrottle the given roller. |
| Unthrottle(ctx context.Context, roller string) error |
| // Reset the should-unthrottle status of the roller, allowing it to become |
| // throttled again if necessary. |
| Reset(ctx context.Context, roller string) error |
| // Get determines whether the given roller should be unthrottled. |
| Get(ctx context.Context, roller string) (bool, error) |
| } |
| |
| // entry is a struct representing the should-unthrottle status for a roller. |
| type entry struct { |
| ShouldUnthrottle bool `datastore:"shouldUnthrottle,noindex"` |
| } |
| |
| // Fake ancestor we supply for unthrottling, to force consistency. |
| // We lose some performance this way but it keeps our tests from |
| // flaking. |
| func fakeAncestor() *datastore.Key { |
| rv := ds.NewKey(ds.KIND_AUTOROLL_UNTHROTTLE_ANCESTOR) |
| rv.ID = 13 // Bogus ID. |
| return rv |
| } |
| |
| // Return a datastore key for the given roller. |
| func key(roller string) *datastore.Key { |
| k := ds.NewKey(ds.KIND_AUTOROLL_UNTHROTTLE) |
| k.Parent = fakeAncestor() |
| k.Name = roller + "_unthrottle" |
| return k |
| } |
| |
| // DatastoreThrottle is an implementation of Throttle which uses Datastore. |
| type DatastoreThrottle struct{} |
| |
| // NewDatastore returns an implementation of Throttle which uses Datastore. |
| func NewDatastore(ctx context.Context) *DatastoreThrottle { |
| return &DatastoreThrottle{} |
| } |
| |
| // Unthrottle implements the Throttle interface. |
| func (t *DatastoreThrottle) Unthrottle(ctx context.Context, roller string) error { |
| return set(ctx, roller, true) |
| } |
| |
| // Reset implements the Throttle interface. |
| func (t *DatastoreThrottle) Reset(ctx context.Context, roller string) error { |
| return set(ctx, roller, false) |
| } |
| |
| // Set whether the given roller should be unthrottled. |
| func set(ctx context.Context, roller string, shouldUnthrottle bool) error { |
| e := &entry{ |
| ShouldUnthrottle: shouldUnthrottle, |
| } |
| _, err := ds.DS.RunInTransaction(ctx, func(tx *datastore.Transaction) error { |
| _, err := tx.Put(key(roller), e) |
| return err |
| }) |
| return err |
| } |
| |
| // Get implements the Throttle interface. |
| func (t *DatastoreThrottle) Get(ctx context.Context, roller string) (bool, error) { |
| var e entry |
| _, err := ds.DS.RunInTransaction(ctx, func(tx *datastore.Transaction) error { |
| return tx.Get(key(roller), &e) |
| }) |
| if err != nil { |
| if err == datastore.ErrNoSuchEntity { |
| return false, nil |
| } else { |
| return false, err |
| } |
| } |
| return e.ShouldUnthrottle, nil |
| } |
| |
| var _ Throttle = &DatastoreThrottle{} |