fix: fill remaining replicas from ready spare slots

Author: strahe

internal/worker/manager.go

@@ -161,7 +161,7 @@ func (m *Manager) reconcileUnavailableDataSets(ctx context.Context) {
 		}
 		for i := range bindings {
 			binding := &bindings[i]
-			if err := ensureReplicaRepairTask(ctx, m.repos, binding, m.uploadMaxRetries); err != nil {
+			if _, err := ensureReplicaRepairTask(ctx, m.repos, binding, m.uploadMaxRetries); err != nil {
 				m.logger.Error("failed to ensure unavailable data set repair", "dataSetID", binding.ID, "error", err)
 			}
 			afterID = binding.ID
@@ -370,7 +370,7 @@ func (m *Manager) reconcileIngressUpload(ctx context.Context, version model.Obje
 		return
 	}
 	if binding.Status == model.StorageDataSetStatusUnavailable {
-		if err := ensureReplicaRepairTask(ctx, m.repos, binding, m.uploadMaxRetries); err != nil {
+		if _, err := ensureReplicaRepairTask(ctx, m.repos, binding, m.uploadMaxRetries); err != nil {
 			m.logger.Error("failed to ensure recovered ingress repair", "dataSetID", binding.ID, "error", err)
 		}
 	}
@@ -458,7 +458,7 @@ func (m *Manager) reconcileReplicatingUpload(ctx context.Context, version model.
 			continue
 		}
 		if binding != nil && binding.Status == model.StorageDataSetStatusUnavailable {
-			if err := ensureReplicaRepairTask(ctx, m.repos, binding, m.uploadMaxRetries); err != nil {
+			if _, err := ensureReplicaRepairTask(ctx, m.repos, binding, m.uploadMaxRetries); err != nil {
 				m.logger.Error("failed to ensure recovered peer repair", "dataSetID", binding.ID, "error", err)
 			}
 			continue

internal/worker/uploader.go

@@ -534,10 +534,19 @@ func (u *Uploader) processDurableUploadTask(
 		u.ensureUploadDataSet(ctx, task, version, bucket, uploadID, copyIndex, logger)
 		return
 	}
-	if err := u.ensureReplicaRepairTask(ctx, binding, task.MaxRetries); err != nil {
+	if binding.Status == model.StorageDataSetStatusReady {
+		u.processStagedTask(ctx, task, version, bucket, stage, logger)
+		return
+	}
+	queued, err := ensureReplicaRepairTask(ctx, u.repos, binding, task.MaxRetries)
+	if err != nil {
 		u.handleTaskFailure(ctx, task, logger, "handoff durable upload repair", err)
 		return
 	}
+	if !queued {
+		u.waitForStorageDependency(ctx, task, logger, "Waiting for in-place replica recovery")
+		return
+	}
 	completeWorkerTask(ctx, u.repos, task, "uploader", logger)
 }
 
@@ -989,23 +998,40 @@ func (u *Uploader) ensureBucketProviderBindings(ctx context.Context, bucket *mod
 	targetCopies = boundedTargetCopies(targetCopies)
 	selected := make([]model.StorageDataSet, 0, targetCopies)
 	excluded := make([]sdktypes.BigInt, 0, len(bindings))
-	for _, binding := range bindings {
+	inFlight := 0
+	for i := range bindings {
+		binding := bindings[i]
 		existing[binding.CopyIndex] = binding
 		excluded = append(excluded, binding.ProviderID.SDK())
-	}
-	missingIndexes := make([]int, 0, targetCopies)
-	for copyIndex := range targetCopies {
-		if binding, ok := existing[copyIndex]; ok {
-			selected = append(selected, binding)
-		} else {
-			missingIndexes = append(missingIndexes, copyIndex)
+		if dataSetBindingCanOccupyUploadSlot(uploadID, &binding) {
+			if len(selected) < targetCopies {
+				selected = append(selected, binding)
+			}
+			continue
+		}
+		if dataSetBindingInFlightForOtherUpload(uploadID, &binding) {
+			inFlight++
 		}
 	}
-	if len(missingIndexes) == 0 {
+	if len(selected) >= targetCopies {
 		plan := newBucketBindingPlan(selected, uploadID)
 		plan.complete = true
 		return plan, nil
 	}
+	if inFlight > 0 {
+		return newBucketBindingPlan(selected, uploadID), nil
+	}
+	need := targetCopies - len(selected)
+	missingIndexes := make([]int, 0, need)
+	for copyIndex := 0; len(missingIndexes) < need && copyIndex < model.StorageCopiesMax; copyIndex++ {
+		if _, occupied := existing[copyIndex]; occupied {
+			continue
+		}
+		missingIndexes = append(missingIndexes, copyIndex)
+	}
+	if len(missingIndexes) == 0 {
+		return newBucketBindingPlan(selected, uploadID), nil
+	}
 	contexts, err := u.storage.CreateContexts(ctx, &storage.CreateContextsOptions{
 		Copies:             len(missingIndexes),
 		ExcludeProviderIDs: excluded,
@@ -1313,23 +1339,6 @@ func (u *Uploader) ensureUploadDataSet(ctx context.Context, task *model.Task, ve
 			}
 		}
 	}
-	if durableObjectState(version.State) {
-		bindingID := binding.ID
-		binding, err = u.repos.Uploads.GetDataSetBindingByID(ctx, bindingID)
-		if err != nil || binding == nil {
-			if err == nil {
-				err = fmt.Errorf("dataset binding %d not found", bindingID)
-			}
-			u.handleTaskFailure(ctx, task, logger, "reload durable upload data set", err)
-			return
-		}
-		if err := u.ensureReplicaRepairTask(ctx, binding, task.MaxRetries); err != nil {
-			u.handleTaskFailure(ctx, task, logger, "handoff durable upload repair", err)
-			return
-		}
-		completeWorkerTask(ctx, u.repos, task, "uploader", logger)
-		return
-	}
 	nextStage := uploadStagePeerPull
 	if copyRow.TransferMethod == model.StorageCopyTransferMethodIngress {
 		nextStage = uploadStageIngressStore
@@ -2002,6 +2011,18 @@ func (u *Uploader) finishPeerCopy(ctx context.Context, task *model.Task, version
 	if !u.repairReadableBinding(ctx, task, version, uploadID, logger, "repair readable binding") {
 		return
 	}
+	ref := repository.ObjectVersionRef{ObjectID: version.ObjectID, VersionID: version.VersionID}
+	_, needsPreparation, err := u.scheduleRemainingPeerCopies(ctx, ref, version.BucketID, uploadID, task.MaxRetries)
+	if err != nil {
+		u.handleTaskFailure(ctx, task, logger, "schedule remaining upload copies", err)
+		return
+	}
+	if needsPreparation {
+		if err := u.enqueueRepairUploadForVersion(ctx, ref, task.MaxRetries, uploadID); err != nil {
+			u.handleTaskFailure(ctx, task, logger, "schedule upload preparation", err)
+			return
+		}
+	}
 	if _, _, err := u.repos.Uploads.FinalizeUploadIfTargetCopiesMet(ctx, u.finalizeUploadInput(uploadID)); err != nil {
 		u.handleTaskFailure(ctx, task, logger, "finalize upload", err)
 		return
@@ -2820,6 +2841,34 @@ func uploadCanUseDataSetBinding(uploadID int64, binding *model.StorageDataSet) b
 		dataSetBindingCanEnsureWrite(binding)
 }
 
+func dataSetBindingCanOccupyUploadSlot(uploadID int64, binding *model.StorageDataSet) bool {
+	if binding == nil {
+		return false
+	}
+	switch binding.Status {
+	case model.StorageDataSetStatusReady:
+		return true
+	case model.StorageDataSetStatusUnavailable:
+		return dataSetBindingEstablished(binding)
+	case model.StorageDataSetStatusPending, model.StorageDataSetStatusCreating, model.StorageDataSetStatusFailed:
+		return uploadCanUseDataSetBinding(uploadID, binding)
+	default:
+		return false
+	}
+}
+
+func dataSetBindingInFlightForOtherUpload(uploadID int64, binding *model.StorageDataSet) bool {
+	if binding == nil || dataSetBindingCanOccupyUploadSlot(uploadID, binding) {
+		return false
+	}
+	switch binding.Status {
+	case model.StorageDataSetStatusPending, model.StorageDataSetStatusCreating:
+		return true
+	default:
+		return false
+	}
+}
+
 func uploadTracksDataSetBinding(uploadID int64, binding *model.StorageDataSet) bool {
 	if binding == nil {
 		return false

internal/worker/uploader_replica_repair.go

@@ -60,35 +60,36 @@ func parseReplicaRepairPayload(task *model.Task) (replicaRepairPayload, error) {
 }
 
 func (u *Uploader) ensureReplicaRepairTask(ctx context.Context, binding *model.StorageDataSet, maxRetries int) error {
-	return ensureReplicaRepairTask(ctx, u.repos, binding, maxRetries)
+	_, err := ensureReplicaRepairTask(ctx, u.repos, binding, maxRetries)
+	return err
 }
 
-func ensureReplicaRepairTask(ctx context.Context, repos *repository.Repositories, binding *model.StorageDataSet, maxRetries int) error {
+func ensureReplicaRepairTask(ctx context.Context, repos *repository.Repositories, binding *model.StorageDataSet, maxRetries int) (bool, error) {
 	if binding == nil || binding.ID <= 0 || !dataSetBindingEstablished(binding) {
-		return nil
+		return false, nil
 	}
 	if binding.Status != model.StorageDataSetStatusUnavailable && binding.Status != model.StorageDataSetStatusReady {
-		return nil
+		return false, nil
 	}
 	copyRow, err := repos.Uploads.NextFinalizableCopyForDataSet(ctx, binding.ID)
 	if err != nil {
-		return fmt.Errorf("select replica finalization copy for data set %d: %w", binding.ID, err)
+		return false, fmt.Errorf("select replica finalization copy for data set %d: %w", binding.ID, err)
 	}
 	if copyRow == nil {
 		copyRow, err = repos.Uploads.NextIncompleteCopyForDataSet(ctx, binding.ID)
 		if err != nil {
-			return fmt.Errorf("select replica repair copy for data set %d: %w", binding.ID, err)
+			return false, fmt.Errorf("select replica repair copy for data set %d: %w", binding.ID, err)
 		}
 	}
 	if copyRow == nil {
-		return nil
+		return false, nil
 	}
 	upload, err := repos.Uploads.GetByID(ctx, copyRow.UploadID)
 	if err != nil {
-		return fmt.Errorf("load replica repair upload %d: %w", copyRow.UploadID, err)
+		return false, fmt.Errorf("load replica repair upload %d: %w", copyRow.UploadID, err)
 	}
 	if upload == nil || upload.SourceVersionID == "" {
-		return fmt.Errorf("load replica repair upload %d: %w", copyRow.UploadID, repository.ErrNotFound)
+		return false, fmt.Errorf("load replica repair upload %d: %w", copyRow.UploadID, repository.ErrNotFound)
 	}
 	stage := uploadStageRepairReplica
 	task := &model.Task{
@@ -104,9 +105,9 @@ func ensureReplicaRepairTask(ctx context.Context, repos *repository.Repositories
 		ScheduledAt:    time.Now(),
 	}
 	if _, err := repos.Tasks.EnsureRecurring(ctx, task); err != nil {
-		return fmt.Errorf("ensure replica repair task for data set %d: %w", binding.ID, err)
+		return false, fmt.Errorf("ensure replica repair task for data set %d: %w", binding.ID, err)
 	}
-	return nil
+	return true, nil
 }
 
 func (u *Uploader) processReplicaRepairTask(ctx context.Context, task *model.Task, logger *slog.Logger) {