fix(coordinator): task dispatch fixes — priority order, proofs-ready guard, attempts refund

coordinator/internal/controller/api/get_task.go

@@ -3,7 +3,7 @@ package api
 import (
 	"errors"
 	"fmt"
-	"math/rand"
+	"sort"
 
 	"github.com/gin-gonic/gin"
 	"github.com/prometheus/client_golang/prometheus"
@@ -79,7 +79,9 @@ func (ptc *GetTaskController) incGetTaskAccessCounter(ctx *gin.Context) error {
 	return nil
 }
 
-// GetTasks get assigned chunk/batch task
+// GetTasks get assigned chunk/batch/bundle task, trying proof types in priority
+// order: Bundle > Batch > Chunk. This lets the pipeline finalize the earliest
+// ready bundle before proving unrelated chunks/batches further ahead.
 func (ptc *GetTaskController) GetTasks(ctx *gin.Context) {
 	var getTaskParameter coordinatorType.GetTaskParameter
 	if err := ctx.ShouldBind(&getTaskParameter); err != nil {
@@ -99,35 +101,36 @@ func (ptc *GetTaskController) GetTasks(ctx *gin.Context) {
 		}
 	}
 
-	proofType := ptc.proofType(&getTaskParameter)
-	proverTask, isExist := ptc.proverTasks[proofType]
-	if !isExist {
-		nerr := fmt.Errorf("parameter wrong proof type:%v", proofType)
-		types.RenderFailure(ctx, types.ErrCoordinatorParameterInvalidNo, nerr)
-		return
-	}
+	proofTypes := ptc.prioritizedProofTypes(&getTaskParameter)
 
 	if err := ptc.incGetTaskAccessCounter(ctx); err != nil {
 		log.Warn("get_task access counter inc failed", "error", err.Error())
 	}
 
-	result, err := proverTask.Assign(ctx, &getTaskParameter)
-	if err != nil {
-		nerr := fmt.Errorf("return prover task err:%w", err)
-		types.RenderFailure(ctx, types.ErrCoordinatorGetTaskFailure, nerr)
-		return
-	}
+	for _, proofType := range proofTypes {
+		proverTask, isExist := ptc.proverTasks[proofType]
+		if !isExist {
+			continue
+		}
 
-	if result == nil {
-		nerr := errors.New("get empty prover task")
-		types.RenderFailure(ctx, types.ErrCoordinatorEmptyProofData, nerr)
-		return
+		result, err := proverTask.Assign(ctx, &getTaskParameter)
+		if err != nil {
+			nerr := fmt.Errorf("return prover task err:%w", err)
+			types.RenderFailure(ctx, types.ErrCoordinatorGetTaskFailure, nerr)
+			return
+		}
+
+		if result != nil {
+			types.RenderSuccess(ctx, result)
+			return
+		}
 	}
 
-	types.RenderSuccess(ctx, result)
+	nerr := errors.New("get empty prover task")
+	types.RenderFailure(ctx, types.ErrCoordinatorEmptyProofData, nerr)
 }
 
-func (ptc *GetTaskController) proofType(para *coordinatorType.GetTaskParameter) message.ProofType {
+func (ptc *GetTaskController) prioritizedProofTypes(para *coordinatorType.GetTaskParameter) []message.ProofType {
 	var proofTypes []message.ProofType
 	for _, proofType := range para.TaskTypes {
 		proofTypes = append(proofTypes, message.ProofType(proofType))
@@ -141,8 +144,9 @@ func (ptc *GetTaskController) proofType(para *coordinatorType.GetTaskParameter)
 		}
 	}
 
-	rand.Shuffle(len(proofTypes), func(i, j int) {
-		proofTypes[i], proofTypes[j] = proofTypes[j], proofTypes[i]
+	// Bundle (3) > Batch (2) > Chunk (1)
+	sort.Slice(proofTypes, func(i, j int) bool {
+		return proofTypes[i] > proofTypes[j]
 	})
-	return proofTypes[0]
+	return proofTypes
 }

coordinator/internal/logic/provertask/batch_prover_task.go

@@ -199,7 +199,7 @@ func (bp *BatchProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 
 	taskMsg, err := bp.formatProverTask(ctx.Copy(), proverTask, batchTask, hardForkName)
 	if err != nil {
-		bp.recoverActiveAttempts(ctx, batchTask)
+		bp.recoverAttempts(ctx, taskCtx, batchTask)
 		log.Error("format prover task failure", "task_id", batchTask.Hash, "err", err)
 		return nil, ErrCoordinatorInternalFailure
 	}
@@ -208,7 +208,7 @@ func (bp *BatchProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 
 		taskMsg, metadata, err = bp.applyUniversal(taskMsg)
 		if err != nil {
-			bp.recoverActiveAttempts(ctx, batchTask)
+			bp.recoverAttempts(ctx, taskCtx, batchTask)
 			log.Error("Generate universal prover task failure", "task_id", batchTask.Hash, "type", "batch", "err", err)
 			return nil, ErrCoordinatorInternalFailure
 		}
@@ -217,7 +217,8 @@ func (bp *BatchProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 		if isCompatibilityFixingVersion(taskCtx.ProverVersion) {
 			log.Info("Apply compatibility fixing for prover", "version", taskCtx.ProverVersion)
 			if err := fixCompatibility(taskMsg); err != nil {
-				log.Error("apply compatibility failure", "err", err)
+				bp.recoverAttempts(ctx, taskCtx, batchTask)
+				log.Error("apply compatibility failure", "task_id", batchTask.Hash, "err", err)
 				return nil, ErrCoordinatorInternalFailure
 			}
 		}
@@ -226,7 +227,7 @@ func (bp *BatchProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 	// Store session info.
 	if taskCtx.hasAssignedTask == nil {
 		if err = bp.proverTaskOrm.InsertProverTask(ctx.Copy(), proverTask); err != nil {
-			bp.recoverActiveAttempts(ctx, batchTask)
+			bp.recoverAttempts(ctx, taskCtx, batchTask)
 			log.Error("insert batch prover task info fail", "task_id", batchTask.Hash, "publicKey", taskCtx.PublicKey, "err", err)
 			return nil, ErrCoordinatorInternalFailure
 		}
@@ -288,7 +289,17 @@ func (bp *BatchProverTask) formatProverTask(ctx context.Context, task *orm.Prove
 	return taskMsg, nil
 }
 
-func (bp *BatchProverTask) recoverActiveAttempts(ctx *gin.Context, batchTask *orm.Batch) {
+// recoverAttempts rolls back the attempt charged by UpdateBatchAttempts when the coordinator
+// fails to dispatch a freshly assigned task (hasAssignedTask == nil): both counters are refunded
+// and proving_status is reset to unassigned. For a re-poll of an already assigned task nothing
+// was charged in this call, so only active_attempts is decremented as before.
+func (bp *BatchProverTask) recoverAttempts(ctx *gin.Context, taskCtx *proverTaskContext, batchTask *orm.Batch) {
+	if taskCtx.hasAssignedTask == nil {
+		if err := bp.batchOrm.RefundAttemptsByHash(ctx.Copy(), batchTask.Hash); err != nil {
+			log.Error("failed to refund batch attempts", "hash", batchTask.Hash, "error", err)
+		}
+		return
+	}
 	if err := bp.batchOrm.DecreaseActiveAttemptsByHash(ctx.Copy(), batchTask.Hash); err != nil {
 		log.Error("failed to recover batch active attempts", "hash", batchTask.Hash, "error", err)
 	}

coordinator/internal/logic/provertask/bundle_prover_task.go

@@ -196,15 +196,15 @@ func (bp *BundleProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinat
 
 	taskMsg, err := bp.formatProverTask(ctx.Copy(), proverTask, hardForkName)
 	if err != nil {
-		bp.recoverActiveAttempts(ctx, bundleTask)
+		bp.recoverAttempts(ctx, taskCtx, bundleTask)
 		log.Error("format bundle prover task failure", "task_id", bundleTask.Hash, "err", err)
 		return nil, ErrCoordinatorInternalFailure
 	}
 	if getTaskParameter.Universal {
 		var metadata []byte
 		taskMsg, metadata, err = bp.applyUniversal(taskMsg)
 		if err != nil {
-			bp.recoverActiveAttempts(ctx, bundleTask)
+			bp.recoverAttempts(ctx, taskCtx, bundleTask)
 			log.Error("Generate universal prover task failure", "task_id", bundleTask.Hash, "type", "bundle", "err", err)
 			return nil, ErrCoordinatorInternalFailure
 		}
@@ -215,7 +215,8 @@ func (bp *BundleProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinat
 		if isCompatibilityFixingVersion(taskCtx.ProverVersion) {
 			log.Info("Apply compatibility fixing for prover", "version", taskCtx.ProverVersion)
 			if err := fixCompatibility(taskMsg); err != nil {
-				log.Error("apply compatibility failure", "err", err)
+				bp.recoverAttempts(ctx, taskCtx, bundleTask)
+				log.Error("apply compatibility failure", "task_id", bundleTask.Hash, "err", err)
 				return nil, ErrCoordinatorInternalFailure
 			}
 		}
@@ -224,7 +225,7 @@ func (bp *BundleProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinat
 	// Store session info.
 	if taskCtx.hasAssignedTask == nil {
 		if err = bp.proverTaskOrm.InsertProverTask(ctx.Copy(), proverTask); err != nil {
-			bp.recoverActiveAttempts(ctx, bundleTask)
+			bp.recoverAttempts(ctx, taskCtx, bundleTask)
 			log.Error("insert bundle prover task info fail", "task_id", bundleTask.Hash, "publicKey", taskCtx.PublicKey, "err", err)
 			return nil, ErrCoordinatorInternalFailure
 		}
@@ -313,7 +314,17 @@ func (bp *BundleProverTask) formatProverTask(ctx context.Context, task *orm.Prov
 	return taskMsg, nil
 }
 
-func (bp *BundleProverTask) recoverActiveAttempts(ctx *gin.Context, bundleTask *orm.Bundle) {
+// recoverAttempts rolls back the attempt charged by UpdateBundleAttempts when the coordinator
+// fails to dispatch a freshly assigned task (hasAssignedTask == nil): both counters are refunded
+// and proving_status is reset to unassigned. For a re-poll of an already assigned task nothing
+// was charged in this call, so only active_attempts is decremented as before.
+func (bp *BundleProverTask) recoverAttempts(ctx *gin.Context, taskCtx *proverTaskContext, bundleTask *orm.Bundle) {
+	if taskCtx.hasAssignedTask == nil {
+		if err := bp.bundleOrm.RefundAttemptsByHash(ctx.Copy(), bundleTask.Hash); err != nil {
+			log.Error("failed to refund bundle attempts", "hash", bundleTask.Hash, "error", err)
+		}
+		return
+	}
 	if err := bp.bundleOrm.DecreaseActiveAttemptsByHash(ctx.Copy(), bundleTask.Hash); err != nil {
 		log.Error("failed to recover bundle active attempts", "hash", bundleTask.Hash, "error", err)
 	}

coordinator/internal/logic/provertask/chunk_prover_task.go

@@ -194,7 +194,7 @@ func (cp *ChunkProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 
 	taskMsg, err := cp.formatProverTask(ctx.Copy(), proverTask, chunkTask, hardForkName)
 	if err != nil {
-		cp.recoverActiveAttempts(ctx, chunkTask)
+		cp.recoverAttempts(ctx, taskCtx, chunkTask)
 		log.Error("format prover task failure", "task_id", chunkTask.Hash, "err", err)
 		return nil, ErrCoordinatorInternalFailure
 	}
@@ -203,7 +203,7 @@ func (cp *ChunkProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 		var metadata []byte
 		taskMsg, metadata, err = cp.applyUniversal(taskMsg)
 		if err != nil {
-			cp.recoverActiveAttempts(ctx, chunkTask)
+			cp.recoverAttempts(ctx, taskCtx, chunkTask)
 			log.Error("Generate universal prover task failure", "task_id", chunkTask.Hash, "type", "chunk", "err", err)
 			return nil, ErrCoordinatorInternalFailure
 		}
@@ -212,7 +212,7 @@ func (cp *ChunkProverTask) Assign(ctx *gin.Context, getTaskParameter *coordinato
 
 	if taskCtx.hasAssignedTask == nil {
 		if err = cp.proverTaskOrm.InsertProverTask(ctx.Copy(), proverTask); err != nil {
-			cp.recoverActiveAttempts(ctx, chunkTask)
+			cp.recoverAttempts(ctx, taskCtx, chunkTask)
 			log.Error("insert chunk prover task fail", "task_id", chunkTask.Hash, "publicKey", taskCtx.PublicKey, "err", err)
 			return nil, ErrCoordinatorInternalFailure
 		}
@@ -269,7 +269,17 @@ func (cp *ChunkProverTask) formatProverTask(ctx context.Context, task *orm.Prove
 	return proverTaskSchema, nil
 }
 
-func (cp *ChunkProverTask) recoverActiveAttempts(ctx *gin.Context, chunkTask *orm.Chunk) {
+// recoverAttempts rolls back the attempt charged by UpdateChunkAttempts when the coordinator
+// fails to dispatch a freshly assigned task (hasAssignedTask == nil): both counters are refunded
+// and proving_status is reset to unassigned. For a re-poll of an already assigned task nothing
+// was charged in this call, so only active_attempts is decremented as before.
+func (cp *ChunkProverTask) recoverAttempts(ctx *gin.Context, taskCtx *proverTaskContext, chunkTask *orm.Chunk) {
+	if taskCtx.hasAssignedTask == nil {
+		if err := cp.chunkOrm.RefundAttemptsByHash(ctx, chunkTask.Hash); err != nil {
+			log.Error("failed to refund chunk attempts", "hash", chunkTask.Hash, "error", err)
+		}
+		return
+	}
 	if err := cp.chunkOrm.DecreaseActiveAttemptsByHash(ctx, chunkTask.Hash); err != nil {
 		log.Error("failed to recover chunk active attempts", "hash", chunkTask.Hash, "error", err)
 	}

coordinator/internal/orm/batch.go

@@ -211,13 +211,25 @@ func (o *Batch) GetAttemptsByHash(ctx context.Context, hash string) (int16, int1
 func (o *Batch) CheckIfBundleBatchProofsAreReady(ctx context.Context, bundleHash string) (bool, error) {
 	db := o.db.WithContext(ctx)
 	db = db.Model(&Batch{})
+	db = db.Where("bundle_hash = ?", bundleHash)
+
+	var totalCount int64
+	if err := db.Count(&totalCount).Error; err != nil {
+		return false, fmt.Errorf("Batch.CheckIfBundleBatchProofsAreReady error: %w, bundle hash: %v", err, bundleHash)
+	}
+	if totalCount == 0 {
+		return false, nil
+	}
+
+	db = o.db.WithContext(ctx)
+	db = db.Model(&Batch{})
 	db = db.Where("bundle_hash = ? AND proving_status != ?", bundleHash, types.ProvingTaskVerified)
 
-	var count int64
-	if err := db.Count(&count).Error; err != nil {
-		return false, fmt.Errorf("Chunk.CheckIfBundleBatchProofsAreReady error: %w, bundle hash: %v", err, bundleHash)
+	var unreadyCount int64
+	if err := db.Count(&unreadyCount).Error; err != nil {
+		return false, fmt.Errorf("Batch.CheckIfBundleBatchProofsAreReady error: %w, bundle hash: %v", err, bundleHash)
 	}
-	return count == 0, nil
+	return unreadyCount == 0, nil
 }
 
 // GetBatchByHash retrieves the given batch.
@@ -459,3 +471,32 @@ func (o *Batch) DecreaseActiveAttemptsByHash(ctx context.Context, batchHash stri
 	}
 	return nil
 }
+
+// RefundAttemptsByHash refunds a full assignment attempt of a batch given its hash:
+// it decrements both total_attempts and active_attempts and resets proving_status to unassigned.
+// It is used to roll back UpdateBatchAttempts when the coordinator fails to dispatch the task
+// (e.g. task formatting or prover task insertion failure), so the attempt is not burned permanently.
+func (o *Batch) RefundAttemptsByHash(ctx context.Context, batchHash string, dbTX ...*gorm.DB) error {
+	db := o.db
+	if len(dbTX) > 0 && dbTX[0] != nil {
+		db = dbTX[0]
+	}
+	db = db.WithContext(ctx)
+	db = db.Model(&Batch{})
+	db = db.Where("hash = ?", batchHash)
+	db = db.Where("total_attempts > ?", 0)
+	db = db.Where("active_attempts > ?", 0)
+	db = db.Where("proving_status != ?", int(types.ProvingTaskVerified))
+	result := db.Updates(map[string]interface{}{
+		"total_attempts":  gorm.Expr("total_attempts - 1"),
+		"active_attempts": gorm.Expr("active_attempts - 1"),
+		"proving_status":  int(types.ProvingTaskUnassigned),
+	})
+	if result.Error != nil {
+		return fmt.Errorf("Batch.RefundAttemptsByHash error: %w, batch hash: %v", result.Error, batchHash)
+	}
+	if result.RowsAffected == 0 {
+		log.Warn("No rows were affected in RefundAttemptsByHash", "batch hash", batchHash)
+	}
+	return nil
+}

coordinator/internal/orm/bundle.go

@@ -243,3 +243,32 @@ func (o *Bundle) DecreaseActiveAttemptsByHash(ctx context.Context, bundleHash st
 	}
 	return nil
 }
+
+// RefundAttemptsByHash refunds a full assignment attempt of a bundle given its hash:
+// it decrements both total_attempts and active_attempts and resets proving_status to unassigned.
+// It is used to roll back UpdateBundleAttempts when the coordinator fails to dispatch the task
+// (e.g. task formatting or prover task insertion failure), so the attempt is not burned permanently.
+func (o *Bundle) RefundAttemptsByHash(ctx context.Context, bundleHash string, dbTX ...*gorm.DB) error {
+	db := o.db
+	if len(dbTX) > 0 && dbTX[0] != nil {
+		db = dbTX[0]
+	}
+	db = db.WithContext(ctx)
+	db = db.Model(&Bundle{})
+	db = db.Where("hash = ?", bundleHash)
+	db = db.Where("total_attempts > ?", 0)
+	db = db.Where("active_attempts > ?", 0)
+	db = db.Where("proving_status != ?", int(types.ProvingTaskVerified))
+	result := db.Updates(map[string]interface{}{
+		"total_attempts":  gorm.Expr("total_attempts - 1"),
+		"active_attempts": gorm.Expr("active_attempts - 1"),
+		"proving_status":  int(types.ProvingTaskUnassigned),
+	})
+	if result.Error != nil {
+		return fmt.Errorf("Bundle.RefundAttemptsByHash error: %w, bundle hash: %v", result.Error, bundleHash)
+	}
+	if result.RowsAffected == 0 {
+		log.Warn("No rows were affected in RefundAttemptsByHash", "bundle hash", bundleHash)
+	}
+	return nil
+}