update harmonytask Do() to require context

Author: LexLuthr

alertmanager/task_alert.go

@@ -100,9 +100,7 @@ func NewAlertTask(
 	}
 }
 
-func (a *AlertTask) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
-	ctx := context.Background()
-
+func (a *AlertTask) Do(ctx context.Context, taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
 	alMap := make(map[string]*alertOut)
 
 	altrs := &alerts{

harmony/harmonytask/harmonytask.go

@@ -83,7 +83,7 @@ type TaskInterface interface {
 	// ONLY be called by harmonytask.
 	// Indicate if the task no-longer needs scheduling with done=true including
 	// cases where it's past the deadline.
-	Do(taskID TaskID, stillOwned func() bool) (done bool, err error)
+	Do(ctx context.Context, taskID TaskID, stillOwned func() bool) (done bool, err error)
 
 	// CanAccept should return if the task can run on this machine. It should
 	// return null if the task type is not allowed on this machine.
@@ -242,7 +242,7 @@ func New(
 // passing a deadline will ignore those still running (to be picked-up later).
 func (e *TaskEngine) GracefullyTerminate() {
 
-	// call the cancel func to avoid picking up any new tasks. Running tasks have context.Background()
+	// call the cancel func to avoid picking up any new tasks. Running tasks now inherit e.ctx.
 	// Call shutdown to stop posting heartbeat to DB.
 	e.grace()
 	e.reg.Shutdown()

harmony/harmonytask/task_type_handler.go

@@ -270,7 +270,10 @@ func (h *taskTypeHandler) considerWork(from string, ids []TaskID) (workAccepted
 				}
 			}()
 
-			done, doErr = h.Do(tID, func() bool {
+			taskCtx, taskCancel := context.WithCancel(h.TaskEngine.ctx)
+			defer taskCancel()
+
+			done, doErr = h.Do(taskCtx, tID, func() bool {
 				if taskhelp.IsBackgroundTask(h.Name) || h.CanYield {
 					if h.TaskEngine.yieldBackground.Load() {
 						log.Infow("yielding background task", "name", h.Name, "id", tID)
@@ -279,8 +282,8 @@ func (h *taskTypeHandler) considerWork(from string, ids []TaskID) (workAccepted
 				}
 
 				var owner int
-				// Background here because we don't want GracefulRestart to block this save.
-				err := h.TaskEngine.db.QueryRow(context.Background(),
+				// Keep ownership checks tied to the same task context used by Do().
+				err := h.TaskEngine.db.QueryRow(taskCtx,
 					`SELECT owner_id FROM harmony_task WHERE id=$1`, tID).Scan(&owner)
 				if err != nil {
 					log.Error("Cannot determine ownership: ", err)

itests/alertnow_test.go

@@ -1,6 +1,7 @@
 package itests
 
 import (
+	"context"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -31,7 +32,7 @@ func TestAlertNow(t *testing.T) {
 	alertmanager.AlertFuncs = []alertmanager.AlertFunc{alertmanager.NowCheck}
 	// Create a new alert task
 	at := alertmanager.NewAlertTask(nil, db, config.CurioAlertingConfig{}, as)
-	done, err := at.Do(123, func() bool { return true })
+	done, err := at.Do(context.Background(), 123, func() bool { return true })
 	require.NoError(t, err)
 	require.True(t, done)
 	require.Equal(t, "Machine alertNowMachine: testMessage", tp.output)

tasks/balancemgr/task_balancemgr.go

@@ -176,8 +176,7 @@ func (b *BalanceMgrTask) CanAccept(ids []harmonytask.TaskID, engine *harmonytask
 }
 
 // Do implements harmonytask.TaskInterface.
-func (b *BalanceMgrTask) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
-	ctx := context.Background()
+func (b *BalanceMgrTask) Do(ctx context.Context, taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
 
 	// select task info
 	var id int64

tasks/expmgr/task_expmgr.go

@@ -38,8 +38,7 @@ func NewExpMgrTask(db *harmonydb.DB, chain api.FullNode, pcs *chainsched.CurioCh
 	}
 }
 
-func (e *ExpMgrTask) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
-	ctx := context.Background()
+func (e *ExpMgrTask) Do(ctx context.Context, taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
 
 	log.Infow("starting expiration manager task", "task_id", taskID)
 

tasks/f3/f3_task.go

@@ -66,9 +66,9 @@ func NewF3Task(db *harmonydb.DB, api F3ParticipationAPI, actors *config.Dynamic[
 	}
 }
 
-func (f *F3Task) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
+func (f *F3Task) Do(ctx context.Context, taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
 	// Ensure that all chain calls are made on the same node (the first call will determine the node)
-	ctx := deps.OnSingleNode(context.Background())
+	ctx = deps.OnSingleNode(ctx)
 
 	var spID int64
 	err = f.db.QueryRow(ctx, "SELECT sp_id FROM f3_tasks WHERE task_id = $1", taskID).Scan(&spID)

tasks/gc/pipeline_meta_gc.go

@@ -24,7 +24,7 @@ func NewPipelineGC(db *harmonydb.DB) *PipelineGC {
 	}
 }
 
-func (s *PipelineGC) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
+func (s *PipelineGC) Do(ctx context.Context, taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
 	if err := s.cleanupSealed(); err != nil {
 		return false, xerrors.Errorf("cleanupSealed: %w", err)
 	}

tasks/gc/storage_endpoint_gc.go

@@ -42,7 +42,7 @@ func NewStorageEndpointGC(si paths.SectorIndex, remote *paths.Remote, db *harmon
 	}
 }
 
-func (s *StorageEndpointGC) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
+func (s *StorageEndpointGC) Do(ctx context.Context, taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
 	/*
 		1. Get all storage paths + urls (endpoints)
 		2. Ping each url, record results
@@ -52,8 +52,6 @@ func (s *StorageEndpointGC) Do(taskID harmonytask.TaskID, stillOwned func() bool
 		4.2.1 in the same transaction remove sector refs to the dead path
 	*/
 
-	ctx := context.Background()
-
 	var pathRefs []struct {
 		StorageID     storiface.ID `db:"storage_id"`
 		Urls          string       `db:"urls"`

tasks/gc/storage_gc_mark.go

@@ -57,8 +57,7 @@ func NewStorageGCMark(si paths.SectorIndex, remote *paths.Remote, db *harmonydb.
 	}
 }
 
-func (s *StorageGCMark) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
-	ctx := context.Background()
+func (s *StorageGCMark) Do(ctx context.Context, taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
 
 	/*
 		CREATE TABLE storage_removal_marks (