Accept Task List (#854)

* Accept Task List

* oops unlimited headroom isn't zero

* lint

* oops

* reduce workload to the accepted IDs

* not every task has storage

* limit maxAcceptable to storage partial

* order in tasks

* cleanup

* fix sql

* nil fix

* fun with yb-not-pg type strictness

* ffi update and harmonytask fn

* fix merge issues

* pull

Author: snadrus

alertmanager/task_alert.go

@@ -164,9 +164,8 @@ func (a *AlertTask) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done
 
 }
 
-func (a *AlertTask) CanAccept(ids []harmonytask.TaskID, engine *harmonytask.TaskEngine) (*harmonytask.TaskID, error) {
-	id := ids[0]
-	return &id, nil
+func (a *AlertTask) CanAccept(ids []harmonytask.TaskID, engine *harmonytask.TaskEngine) ([]harmonytask.TaskID, error) {
+	return ids, nil
 }
 
 func (a *AlertTask) TypeDetails() harmonytask.TaskTypeDetails {

extern/filecoin-ffi

@@ -309,19 +309,11 @@ func ensureSchemaExists(connString, schema string) error {
 		return xerrors.New("schema must be of the form " + schemaREString + "\n Got: " + schema)
 	}
 
-	retryWait := InitialSerializationErrorRetryWait
+	_, err = backoffForSerializationError(func() (pgconn.CommandTag, error) {
+		return p.Exec(context.Background(), "CREATE SCHEMA IF NOT EXISTS "+schema)
+	})
 
-retry:
-	_, err = p.Exec(context.Background(), "CREATE SCHEMA IF NOT EXISTS "+schema)
-	if err != nil && IsErrSerialization(err) {
-		time.Sleep(retryWait)
-		retryWait *= 2
-		goto retry
-	}
-	if err != nil {
-		return xerrors.Errorf("cannot create schema: %w", err)
-	}
-	return nil
+	return err
 }
 
 //go:embed sql

harmony/harmonydb/harmonydb.go

@@ -12,11 +12,25 @@ import (
 	"github.com/samber/lo"
 	"github.com/yugabyte/pgx/v5"
 	"github.com/yugabyte/pgx/v5/pgconn"
+	"golang.org/x/xerrors"
 )
 
 var errTx = errors.New("cannot use a non-transaction func in a transaction")
 
-const InitialSerializationErrorRetryWait = 5 * time.Second
+var backoffs = lo.Map([]int{200, 400, 800, 1600, 2400, 5000, 8000}, func(item int, _ int) time.Duration {
+	return time.Duration(item) * time.Millisecond
+})
+
+func backoffForSerializationError[T any](f func() (T, error)) (whatever T, err error) {
+	for _, backoff := range backoffs {
+		res, err := f()
+		if !IsErrSerialization(err) {
+			return res, nil
+		}
+		time.Sleep(backoff)
+	}
+	return whatever, xerrors.Errorf("failed to execute function: %w", err)
+}
 
 // rawStringOnly is _intentionally_private_ to force only basic strings in SQL queries.
 // In any package, raw strings will satisfy compilation.  Ex:
@@ -34,17 +48,14 @@ func (db *DB) Exec(ctx context.Context, sql rawStringOnly, arguments ...any) (co
 		return 0, errTx
 	}
 
-	retryWait := InitialSerializationErrorRetryWait
-
-retry:
-	res, err := db.pgx.Exec(ctx, string(sql), arguments...)
-	if err != nil && IsErrSerialization(err) {
-		time.Sleep(retryWait)
-		retryWait *= 2
-		goto retry
+	res, err := backoffForSerializationError(func() (pgconn.CommandTag, error) {
+		return db.pgx.Exec(ctx, string(sql), arguments...)
+	})
+	if err != nil {
+		return 0, err
 	}
 
-	return int(res.RowsAffected()), err
+	return int(res.RowsAffected()), nil
 }
 
 type Qry interface {
@@ -79,7 +90,9 @@ func (db *DB) Query(ctx context.Context, sql rawStringOnly, arguments ...any) (*
 	if db.usedInTransaction() {
 		return &Query{}, errTx
 	}
-	q, err := db.pgx.Query(ctx, string(sql), arguments...)
+	q, err := backoffForSerializationError(func() (pgx.Rows, error) {
+		return db.pgx.Query(ctx, string(sql), arguments...)
+	})
 	return &Query{q}, err
 }
 
@@ -148,7 +161,9 @@ func (db *DB) Select(ctx context.Context, sliceOfStructPtr any, sql rawStringOnl
 	if db.usedInTransaction() {
 		return errTx
 	}
-	rows, err := db.pgx.Query(ctx, string(sql), arguments...)
+	rows, err := backoffForSerializationError(func() (pgx.Rows, error) {
+		return db.pgx.Query(ctx, string(sql), arguments...)
+	})
 	if err != nil {
 		return err
 	}
@@ -171,8 +186,7 @@ func (db *DB) usedInTransaction() bool {
 }
 
 type TransactionOptions struct {
-	RetrySerializationError            bool
-	InitialSerializationErrorRetryWait time.Duration
+	RetrySerializationError bool
 }
 
 type TransactionOption func(*TransactionOptions)
@@ -183,12 +197,6 @@ func OptionRetry() TransactionOption {
 	}
 }
 
-func OptionSerialRetryTime(d time.Duration) TransactionOption {
-	return func(o *TransactionOptions) {
-		o.InitialSerializationErrorRetryWait = d
-	}
-}
-
 // BeginTransaction is how you can access transactions using this library.
 // The entire transaction happens in the function passed in.
 // The return must be true or a rollback will occur.
@@ -208,23 +216,19 @@ func (db *DB) BeginTransaction(ctx context.Context, f func(*Tx) (commit bool, er
 	}
 
 	opts := TransactionOptions{
-		RetrySerializationError:            false,
-		InitialSerializationErrorRetryWait: InitialSerializationErrorRetryWait,
+		RetrySerializationError: false,
 	}
 
 	for _, o := range opt {
 		o(&opts)
 	}
 
-retry:
-	comm, err := db.transactionInner(ctx, f)
-	if err != nil && opts.RetrySerializationError && IsErrSerialization(err) {
-		time.Sleep(opts.InitialSerializationErrorRetryWait)
-		opts.InitialSerializationErrorRetryWait *= 2
-		goto retry
+	if opts.RetrySerializationError {
+		return backoffForSerializationError(func() (bool, error) {
+			return db.transactionInner(ctx, f)
+		})
 	}
-
-	return comm, err
+	return db.transactionInner(ctx, f)
 }
 
 func (db *DB) transactionInner(ctx context.Context, f func(*Tx) (commit bool, err error)) (didCommit bool, retErr error) {
@@ -256,8 +260,13 @@ func (db *DB) transactionInner(ctx context.Context, f func(*Tx) (commit bool, er
 
 // Exec in a transaction.
 func (t *Tx) Exec(sql rawStringOnly, arguments ...any) (count int, err error) {
-	res, err := t.Tx.Exec(t.ctx, string(sql), arguments...)
-	return int(res.RowsAffected()), err
+	res, err := backoffForSerializationError(func() (pgconn.CommandTag, error) {
+		return t.Tx.Exec(t.ctx, string(sql), arguments...)
+	})
+	if err != nil {
+		return 0, err
+	}
+	return int(res.RowsAffected()), nil
 }
 
 // Query in a transaction.

harmony/harmonydb/userfuncs.go

@@ -86,7 +86,7 @@ type TaskInterface interface {
 	// return null if the task type is not allowed on this machine.
 	// It should select the task it most wants to accomplish.
 	// It is also responsible for determining & reserving disk space (including scratch).
-	CanAccept([]TaskID, *TaskEngine) (*TaskID, error)
+	CanAccept([]TaskID, *TaskEngine) ([]TaskID, error)
 
 	// TypeDetails() returns static details about how this task behaves and
 	// how this machine will run it. Read once at the beginning.
@@ -217,7 +217,7 @@ func New(
 			// edge-case: if old assignments are not available tasks, unlock them.
 			h := e.taskMap[w.Name]
 			if h == nil || !h.considerWork(WorkSourceRecover, []TaskID{TaskID(w.ID)}) {
-				_, err := db.Exec(e.ctx, `UPDATE harmony_task SET owner_id=NULL WHERE id=$1`, w.ID)
+				_, err := db.Exec(e.ctx, `UPDATE harmony_task SET owner_id=NULL WHERE id=$1 AND owner_id=$2`, w.ID, e.ownerID)
 				if err != nil {
 					log.Errorw("Cannot remove self from owner field", "error", err)
 					continue // not really fatal, but not great
@@ -420,7 +420,7 @@ func (e *TaskEngine) pollerTryAllWork(schedulable bool) bool {
 			}
 		}
 
-		if err := v.AssertMachineHasCapacity(); err != nil {
+		if _, err := v.AssertMachineHasCapacity(); err != nil {
 			log.Debugf("skipped scheduling %s type tasks on due to %s", v.Name, err.Error())
 			continue
 		}
@@ -468,7 +468,7 @@ func (e *TaskEngine) pollerTryAllWork(schedulable bool) bool {
 	// if no work was accepted, are we bored? Then find work in priority order.
 	for _, v := range e.handlers {
 		v := v
-		if v.AssertMachineHasCapacity() != nil {
+		if _, err := v.AssertMachineHasCapacity(); err != nil {
 			continue
 		}
 		if v.IAmBored != nil {
@@ -579,3 +579,7 @@ func Reg(t TaskInterface) bool {
 
 	return true
 }
+
+func (e *TaskEngine) RunningCount(name string) int {
+	return int(e.taskMap[name].Max.Active())
+}