Repository: corona-warn-app/cwa-server
File: services/download/src/main/java/app/coronawarn/server/services/download/FederationBatchProcessor.java (Java)

Code before:
      FederationBatchInfo currentBatchInfo = unprocessedBatches.remove();
      seenBatches.add(currentBatchInfo.getBatchTag());
      processBatchAndReturnNextBatchId(currentBatchInfo, ERROR)
          .ifPresent(nextBatchTag -> {
            if (isEfgsEnforceDateBasedDownloadAndNotSeen(nextBatchTag)) {
              unprocessedBatches.add(new FederationBatchInfo(nextBatchTag, currentBatchInfo.getDate()));
            }
          });
    }
  }

  private boolean isEfgsEnforceDateBasedDownloadAndNotSeen(String batchTag) {
    return config.getEfgsEnforceDateBasedDownload() && !seenBatches.contains(batchTag);
  }

  private Optional<String> processBatchAndReturnNextBatchId(
      FederationBatchInfo batchInfo, FederationBatchStatus errorStatus) throws FatalFederationGatewayException {
    LocalDate date = batchInfo.getDate();
    String batchTag = batchInfo.getBatchTag();
    logger.info("Processing batch for date {} and batchTag {}.", date, batchTag);
    AtomicReference<Optional<String>> nextBatchTag = new AtomicReference<>(Optional.empty());
    try {
      BatchDownloadResponse response = federationGatewayDownloadService.downloadBatch(batchTag, date);
      AtomicBoolean batchContainsInvalidKeys = new AtomicBoolean(false);
      if (config.isBatchAuditEnabled()) {
        federationGatewayDownloadService.auditBatch(batchTag, date);
      }
      nextBatchTag.set(response.getNextBatchTag());
      response.getDiagnosisKeyBatch().ifPresentOrElse(batch -> {
        logger.info("Downloaded {} keys for date {} and batchTag {}.", batch.getKeysCount(), date, batchTag);
        List<DiagnosisKey> validDiagnosisKeys = extractValidDiagnosisKeysFromBatch(batch);
        int numOfInvalidKeys = batch.getKeysCount() - validDiagnosisKeys.size();
        if (numOfInvalidKeys > 0) {
          batchContainsInvalidKeys.set(true);
          logger.info("{} keys failed validation and were skipped.", numOfInvalidKeys);
        }
        int insertedKeys = diagnosisKeyService.saveDiagnosisKeys(validDiagnosisKeys);
        logger.info("Successfully inserted {} keys for date {} and batchTag {}.", insertedKeys, date, batchTag);
      }, () -> logger.info("Batch for date {} and batchTag {} did not contain any keys.", date, batchTag));
      batchInfoService.updateStatus(batchInfo, batchContainsInvalidKeys.get() ? PROCESSED_WITH_ERROR : PROCESSED);
      return nextBatchTag.get();
    } catch (FatalFederationGatewayException e) {
      throw e;
    } catch (Exception e) {
      logger.error(
          "Federation batch processing for date " + date + " and batchTag " + batchTag + " failed. Status set to "
              + errorStatus.name() + ".", e);
      batchInfoService.updateStatus(batchInfo, errorStatus);
      return nextBatchTag.get();
    }
  }

Review comment (bug):
what happens, if the auditing fails? is an exception thrown? then the nextBatchTag won't be set, right?

Code after:
   */
  public void processUnprocessedFederationBatches() throws FatalFederationGatewayException {
    Deque<FederationBatchInfo> unprocessedBatches = new LinkedList<>(batchInfoService.findByStatus(UNPROCESSED));
    logger.info("{} unprocessed federation batches found.", unprocessedBatches.size());

    while (!unprocessedBatches.isEmpty()) {
      FederationBatchInfo currentBatchInfo = unprocessedBatches.remove();
      seenBatches.add(currentBatchInfo.getBatchTag());
      processBatchAndReturnNextBatchId(currentBatchInfo, ERROR)
          .ifPresent(nextBatchTag -> {
            if (isEfgsEnforceDateBasedDownloadAndNotSeen(nextBatchTag)) {
              unprocessedBatches.add(new FederationBatchInfo(nextBatchTag, currentBatchInfo.getDate()));
            }
          });
    }
  }

  private boolean isEfgsEnforceDateBasedDownloadAndNotSeen(String batchTag) {
    return config.getEfgsEnforceDateBasedDownload() && !seenBatches.contains(batchTag);
  }

  private Optional<String> processBatchAndReturnNextBatchId(
      FederationBatchInfo batchInfo, FederationBatchStatus errorStatus) throws FatalFederationGatewayException {
    LocalDate date = batchInfo.getDate();
    String batchTag = batchInfo.getBatchTag();
    logger.info("Processing batch for date {} and batchTag {}.", date, batchTag);
    AtomicReference<Optional<String>> nextBatchTag = new AtomicReference<>(Optional.empty());
    try {
      BatchDownloadResponse response = federationGatewayDownloadService.downloadBatch(batchTag, date);
      AtomicBoolean batchContainsInvalidKeys = new AtomicBoolean(false);
      nextBatchTag.set(response.getNextBatchTag());
      response.getDiagnosisKeyBatch().ifPresentOrElse(batch -> {
        logger.info("Downloaded {} keys for date {} and batchTag {}.", batch.getKeysCount(), date, batchTag);
        if (config.isBatchAuditEnabled()) {
          federationGatewayDownloadService.auditBatch(batchTag, date);
        }
        List<DiagnosisKey> validDiagnosisKeys = extractValidDiagnosisKeysFromBatch(batch);
        int numOfInvalidKeys = batch.getKeysCount() - validDiagnosisKeys.size();
        if (numOfInvalidKeys > 0) {
          batchContainsInvalidKeys.set(true);
          logger.info("{} keys failed validation and were skipped.", numOfInvalidKeys);
        }
        int insertedKeys = diagnosisKeyService.saveDiagnosisKeys(validDiagnosisKeys);
        logger.info("Successfully inserted {} keys for date {} and batchTag {}.", insertedKeys, date, batchTag);
      }, () -> logger.info("Batch for date {} and batchTag {} did not contain any keys.", date, batchTag));
      batchInfoService.updateStatus(batchInfo, batchContainsInvalidKeys.get() ? PROCESSED_WITH_ERROR : PROCESSED);
      return nextBatchTag.get();
    } catch (FatalFederationGatewayException e) {
      throw e;
    } catch (Exception e) {
      logger.error(