private static CompletableFuture<Map<TaskName, Checkpoint>> restoreAllTaskInstances(
      Map<TaskName, Map<String, TaskRestoreManager>> taskRestoreManagers, Map<TaskName, Checkpoint> taskCheckpoints,
      Map<TaskName, Map<String, Set<String>>> taskBackendFactoryToStoreNames, JobContext jobContext,
      ContainerModel containerModel, SamzaContainerMetrics samzaContainerMetrics, CheckpointManager checkpointManager,
      Config config, Map<TaskName, TaskInstanceMetrics> taskInstanceMetrics, ExecutorService executor,
      File loggedStoreDir, Set<String> forceRestoreTask) {

    Map<TaskName, Checkpoint> newTaskCheckpoints = new ConcurrentHashMap<>();
    List<Future<Void>> restoreAndCleanupFutures = new ArrayList<>();

    // Submit restore callable for each taskInstance
    taskRestoreManagers.forEach((taskInstanceName, restoreManagersMap) -> {
      // Submit for each restore factory
      restoreManagersMap.forEach((factoryName, taskRestoreManager) -> {
        long startTime = System.currentTimeMillis();
        String taskName = taskInstanceName.getTaskName();
        LOG.info("Starting restore for state for task: {}", taskName);

        CompletableFuture<Void> restoreFuture;
        if (forceRestoreTask.contains(taskName) && taskRestoreManager instanceof BlobStoreRestoreManager) {
          // If init was retried with getDeleted, force restore with getDeleted as well, since init only inits the
          // restoreManager with deleted SnapshotIndex but does not retry to recover the deleted blobs and delegates it
          // to restore().
          // Create an empty future that fails immediately with DeletedException to force retry in restore.
          restoreFuture = new CompletableFuture<>();
          restoreFuture.completeExceptionally(new SamzaException(new DeletedException()));
        } else {
          restoreFuture = taskRestoreManager.restore();
        }

        // Order of the following async operations is critical. They are chained as follows:
        // 1. check if restore succeeded in the first try -> if it failed with DeletedException, flag restore to be retried.
        // 2. retry restore with getDeleted if previous step completes and returns true.
        // 3. Close taskRestoreManager after 1 (and 2, if there is a retry) completes

        CompletableFuture<Boolean> retryRestoreFuture =
            restoreFuture.handle(
              (res, t) -> shouldRetryFailedRestore(t, startTime, samzaContainerMetrics, taskInstanceName));

        CompletableFuture<Checkpoint> taskCheckpointFuture =
            retryRestoreFuture.thenCompose(
              shouldRetryRestore -> retryFailedRestore(shouldRetryRestore, taskInstanceName, taskCheckpoints,
                  checkpointManager, taskRestoreManager, config, taskInstanceMetrics, executor, loggedStoreDir,
                  jobContext, containerModel, factoryName, taskBackendFactoryToStoreNames));

        CompletableFuture<Void> cleanUpFuture =
            taskCheckpointFuture.handle(
              (checkpoint, t) -> cleanUpResources(checkpoint, t, startTime, samzaContainerMetrics, taskInstanceName,
                  taskRestoreManager, newTaskCheckpoints));

        restoreAndCleanupFutures.add(cleanUpFuture);
      });
    });

    return CompletableFuture.allOf(restoreAndCleanupFutures.toArray(new CompletableFuture[0]))
        .thenApply(aVoid -> newTaskCheckpoints);
  }
Restores all TaskInstances and returns a future for each TaskInstance restore. Note: In case of
{@link BlobStoreRestoreManager}, this method restore with getDeleted flag if it receives a
{@link DeletedException}. This will create a new Checkpoint.
samza-core/src/main/scala/org/apache/samza/storage/ContainerStorageManagerRestoreUtil.java