fix: yield in PSClientPoll correctly (#601)

* fix: yield in PSClientPoll correctly

* fix build

* pspoll: Don't timeout the context

Author: magik6k

cmd/curio/debug-proofsvc.go

@@ -202,7 +202,7 @@ var proofsvcStatusCmd = &cli.Command{
 		if err != nil {
 			return err
 		}
-		resp, err := proofsvc.GetProofStatus(rcid)
+		resp, err := proofsvc.GetProofStatus(context.Background(), rcid)
 		if err != nil {
 			return err
 		}

lib/proofsvc/clientctl.go

@@ -249,10 +249,10 @@ func RequestProof(request common.ProofRequest) (bool, error) {
 }
 
 // GetProofStatus checks the status of a proof request by ID
-func GetProofStatus(requestCid cid.Cid) (common.ProofResponse, error) {
+func GetProofStatus(ctx context.Context, requestCid cid.Cid) (common.ProofResponse, error) {
 	start := time.Now()
 	defer recordClientctlDuration("GetProofStatus", start)
-	ctx, cancel := context.WithTimeout(context.Background(), MaxRetryTime)
+	ctx, cancel := context.WithTimeout(ctx, MaxRetryTime)
 	defer cancel()
 
 	return retryWithBackoff(ctx, func() (common.ProofResponse, error) {
@@ -296,7 +296,7 @@ func WaitForProof(request common.ProofRequest) ([]byte, error) {
 	start := time.Now()
 	defer recordClientctlDuration("WaitForProof", start)
 	// Wait for the proof
-	proofResp, err := GetProofStatus(request.Data)
+	proofResp, err := GetProofStatus(context.Background(), request.Data)
 	if err != nil {
 		return nil, xerrors.Errorf("failed to get proof: %w", err)
 	}

tasks/proofshare/task_client_poll.go

@@ -127,7 +127,8 @@ func (t *TaskClientPoll) CanAccept(ids []harmonytask.TaskID, engine *harmonytask
 
 // Do implements harmonytask.TaskInterface.
 func (t *TaskClientPoll) Do(taskID harmonytask.TaskID, stillOwned func() bool) (done bool, err error) {
-	ctx := context.Background()
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
 
 	var clientRequest ClientRequest
 	err = t.db.QueryRow(ctx, `
@@ -148,10 +149,28 @@ func (t *TaskClientPoll) Do(taskID harmonytask.TaskID, stillOwned func() bool) (
 		return false, nil
 	}
 
+	pollCtx, ownedCancel := context.WithCancel(ctx)
+	go func() {
+		const pollInterval = 10 * time.Second
+		defer ownedCancel()
+
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(pollInterval):
+				if !stillOwned() {
+					// close the owned context
+					return
+				}
+			}
+		}
+	}()
+
 	var proof []byte
 	for {
 		var stateChanged bool
-		stateChanged, proof, err = pollForProof(ctx, t.db, taskID, &clientRequest)
+		stateChanged, proof, err = pollForProof(pollCtx, t.db, taskID, &clientRequest)
 		if err != nil {
 			return false, xerrors.Errorf("failed to poll for proof: %w", err)
 		}
@@ -224,7 +243,7 @@ func pollForProof(ctx context.Context, db *harmonydb.DB, taskID harmonytask.Task
 	}
 
 	// Get proof status by CID
-	proofResp, err := proofsvc.GetProofStatus(requestCid)
+	proofResp, err := proofsvc.GetProofStatus(ctx, requestCid)
 	if err != nil || proofResp.Proof == nil {
 		log.Infow("proof not ready", "taskID", taskID, "spID", clientRequest.SpID, "sectorNumber", clientRequest.SectorNumber)
 		// Not ready yet, continue polling