autoroll/go/unthrottle/unthrottle.go

package unthrottle

import (
	"context"

	"cloud.google.com/go/datastore"
	"go.skia.org/infra/go/ds"
)

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

// Unthrottle the given roller.
func Unthrottle(ctx context.Context, roller string) error {
	return set(ctx, roller, true)
}

// Reset the should-unthrottle status of the roller.
func 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
}

// Determine whether the given roller should be unthrottled.
func 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
}