content: /*
   * Licensed to the Apache Software Foundation (ASF) under one
   * or more contributor license agreements.  See the NOTICE file
   * distributed with this work for additional information
   * regarding copyright ownership.  The ASF licenses this file
   * to you under the Apache License, Version 2.0 (the
   * "License"); you may not use this file except in compliance
   * with the License.  You may obtain a copy of the License at
   *
   *   http://www.apache.org/licenses/LICENSE-2.0
   *
   * Unless required by applicable law or agreed to in writing,
   * software distributed under the License is distributed on an
   * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   * KIND, either express or implied.  See the License for the
   * specific language governing permissions and limitations
   * under the License.
   */
  
  package org.apache.bookkeeper.bookie.datainteg;
  
  import com.google.common.collect.ImmutableSortedMap;
  import io.reactivex.rxjava3.core.Flowable;
  import io.reactivex.rxjava3.core.Maybe;
  import io.reactivex.rxjava3.core.Scheduler;
  import io.reactivex.rxjava3.core.Single;
  import io.reactivex.rxjava3.disposables.Disposable;
  import java.io.IOException;
  import java.util.Comparator;
  import java.util.List;
  import java.util.Map;
  import java.util.NavigableMap;
  import java.util.Optional;
  import java.util.Set;
  import java.util.UUID;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.ConcurrentSkipListMap;
  import java.util.concurrent.TimeUnit;
  import java.util.concurrent.atomic.AtomicInteger;
  import java.util.concurrent.atomic.AtomicLong;
  import java.util.concurrent.atomic.AtomicReference;
  import java.util.stream.Collectors;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.bookie.BookieException;
  import org.apache.bookkeeper.bookie.LedgerStorage;
  import org.apache.bookkeeper.bookie.LedgerStorage.StorageState;
  import org.apache.bookkeeper.client.BKException;
  import org.apache.bookkeeper.client.BookKeeperAdmin;
  import org.apache.bookkeeper.client.api.LedgerMetadata;
  import org.apache.bookkeeper.common.concurrent.FutureUtils;
  import org.apache.bookkeeper.meta.LedgerManager;
  import org.apache.bookkeeper.net.BookieId;
  
  /**
   * An implementation of the DataIntegrityCheck interface.
   */
  @Slf4j
  public class DataIntegrityCheckImpl implements DataIntegrityCheck {
      private static final int MAX_INFLIGHT = 300;
      private static final int MAX_ENTRIES_INFLIGHT = 3000;
      private static final int ZK_TIMEOUT_S = 30;
      private final BookieId bookieId;
      private final LedgerManager ledgerManager;
      private final LedgerStorage ledgerStorage;
      private final EntryCopier entryCopier;
      private final BookKeeperAdmin admin;
      private final Scheduler scheduler;
      private final AtomicReference<Map<Long, LedgerMetadata>> ledgersCacheRef =
          new AtomicReference<>(null);
      private CompletableFuture<Void> preBootFuture;
  
      public DataIntegrityCheckImpl(BookieId bookieId,
                                    LedgerManager ledgerManager,
                                    LedgerStorage ledgerStorage,
                                    EntryCopier entryCopier,
                                    BookKeeperAdmin admin,
                                    Scheduler scheduler) {
          this.bookieId = bookieId;
          this.ledgerManager = ledgerManager;
          this.ledgerStorage = ledgerStorage;
          this.entryCopier = entryCopier;
          this.admin = admin;
          this.scheduler = scheduler;
      }
  
      @Override
      public synchronized CompletableFuture<Void> runPreBootCheck(String reason) {
          // we only run this once, it could be kicked off by different checks
          if (preBootFuture == null) {
              preBootFuture = runPreBootSequence(reason);
          }
          return preBootFuture;
  
      }
  
      private CompletableFuture<Void> runPreBootSequence(String reason) {
          String runId = UUID.randomUUID().toString();
          log.info("Event: {}, RunId: {}, Reason: {}", Events.PREBOOT_START, runId, reason);
          try {
              this.ledgerStorage.setStorageStateFlag(StorageState.NEEDS_INTEGRITY_CHECK);
          } catch (IOException ioe) {
              log.error("Event: {}, RunId: {}", Events.PREBOOT_ERROR, runId, ioe);
              return FutureUtils.exception(ioe);
          }
  
          MetadataAsyncIterator iter = new MetadataAsyncIterator(scheduler,
                  ledgerManager, MAX_INFLIGHT, ZK_TIMEOUT_S, TimeUnit.SECONDS);
          CompletableFuture<Void> promise = new CompletableFuture<>();
          Map<Long, LedgerMetadata> ledgersCache =
              new ConcurrentSkipListMap<>(Comparator.<Long>naturalOrder().reversed());
          iter.forEach((ledgerId, metadata) -> {
                  if (ensemblesContainBookie(metadata, bookieId)) {
                      ledgersCache.put(ledgerId, metadata);
                      try {
                          if (!ledgerStorage.ledgerExists(ledgerId)) {
                              ledgerStorage.setMasterKey(ledgerId, new byte[0]);
                          }
                      } catch (IOException ioe) {
                          log.error("Event: {}, RunId: {}, LedgerId: {}",
                                  Events.ENSURE_LEDGER_ERROR, runId, ledgerId, ioe);
                          return FutureUtils.exception(ioe);
                      }
                  }
                  return processPreBoot(ledgerId, metadata, runId);
              })
              .whenComplete((ignore, exception) -> {
                      if (exception != null) {
                          log.error("Event: {}, runId: {}", Events.PREBOOT_ERROR, runId, exception);
                          promise.completeExceptionally(exception);
                      } else {
                          try {
                              this.ledgerStorage.flush();
  
                              updateMetadataCache(ledgersCache);
  
                              log.info("Event: {}, runId: {}, processed: {}",
                                      Events.PREBOOT_END, runId, ledgersCache.size());
                              promise.complete(null);
                          } catch (Throwable t) {
                              log.error("Event: {}, runId: {}", Events.PREBOOT_ERROR, runId, t);
                              promise.completeExceptionally(t);
                          }
                      }
                  });
          return promise;
      }
  
      @Override
      public boolean needsFullCheck() throws IOException {
          return this.ledgerStorage.getStorageStateFlags()
              .contains(StorageState.NEEDS_INTEGRITY_CHECK);
      }
  
      @Override
      public CompletableFuture<Void> runFullCheck() {
          String runId = UUID.randomUUID().toString();
  
          log.info("Event: {}, runId: {}", Events.FULL_CHECK_INIT, runId);
          return getCachedOrReadMetadata(runId)
              .thenCompose(
                      (ledgers) -> {
                          log.info("Event: {}, runId: {}, ledgerCount: {}",
                                  Events.FULL_CHECK_START, runId, ledgers.size());
                          return checkAndRecoverLedgers(ledgers, runId).thenApply((resolved) -> {
                                  for (LedgerResult r : resolved) {
                                      if (r.isMissing() || r.isOK()) {
                                          ledgers.remove(r.getLedgerId());
                                      } else if (r.isError()) {
                                          // if there was an error, make sure we have the latest
                                          // metadata for the next iteration
                                          ledgers.put(r.getLedgerId(), r.getMetadata());
                                      }
                                  }
                                  Optional<Throwable> firstError = resolved.stream().filter(r -> r.isError())
                                      .map(r -> r.getThrowable()).findFirst();
  
                                  if (firstError.isPresent()) {
                                      log.error("Event: {}, runId: {}, ok: {}"
                                              + ", error: {}, missing: {}, ledgersToRetry: {}",
                                              Events.FULL_CHECK_END, runId,
                                              resolved.stream().filter(r -> r.isOK()).count(),
                                              resolved.stream().filter(r -> r.isError()).count(),
                                              resolved.stream().filter(r -> r.isMissing()).count(),
                                              ledgers.size(), firstError.get());
                                  } else {
                                      log.info("Event: {}, runId: {}, ok: {}, error: 0, missing: {}, ledgersToRetry: {}",
                                              Events.FULL_CHECK_END, runId,
                                              resolved.stream().filter(r -> r.isOK()).count(),
                                              resolved.stream().filter(r -> r.isMissing()).count(),
                                              ledgers.size());
                                  }
                                  return ledgers;
                              });
                      })
              .thenCompose(
                      (ledgers) -> {
                          CompletableFuture<Void> promise = new CompletableFuture<>();
                          try {
                              this.ledgerStorage.flush();
                              if (ledgers.isEmpty()) {
                                  log.info("Event: {}, runId: {}", Events.CLEAR_INTEGCHECK_FLAG, runId);
                                  this.ledgerStorage.clearStorageStateFlag(
                                          StorageState.NEEDS_INTEGRITY_CHECK);
                              }
                              // not really needed as we are modifying the map in place
                              updateMetadataCache(ledgers);
                              log.info("Event: {}, runId: {}", Events.FULL_CHECK_COMPLETE, runId);
                              promise.complete(null);
                          } catch (IOException ioe) {
                              log.error("Event: {}, runId: {}", Events.FULL_CHECK_ERROR, runId, ioe);
                              promise.completeExceptionally(ioe);
                          }
                          return promise;
                      });
      }
  
      void updateMetadataCache(Map<Long, LedgerMetadata> ledgers) {
          ledgersCacheRef.set(ledgers);
      }
  
      CompletableFuture<Map<Long, LedgerMetadata>> getCachedOrReadMetadata(String runId) {
          Map<Long, LedgerMetadata> map = ledgersCacheRef.get();
          if (map != null) {
              log.info("Event: {}, runId: {}, ledgerCount: {}", Events.USE_CACHED_METADATA, runId,
                      map.size());
              return CompletableFuture.completedFuture(map);
          } else {
              log.info("Event: {}, runId: {}", Events.REFRESH_METADATA, runId);
              MetadataAsyncIterator iter = new MetadataAsyncIterator(scheduler,
                      ledgerManager, MAX_INFLIGHT, ZK_TIMEOUT_S, TimeUnit.SECONDS);
              Map<Long, LedgerMetadata> ledgersCache =
                  new ConcurrentSkipListMap<>(Comparator.<Long>naturalOrder().reversed());
              return iter.forEach((ledgerId, metadata) -> {
                      if (ensemblesContainBookie(metadata, bookieId)) {
                          ledgersCache.put(ledgerId, metadata);
                      }
                      return CompletableFuture.completedFuture(null);
                  })
                  .thenApply(ignore -> {
                          updateMetadataCache(ledgersCache);
                          return ledgersCache;
                      });
          }
      }
  
      /**
       * Check whether the current bookie exists in the last ensemble of the bookie.
       * If it does, and the ledger is not closed, then this bookie may have accepted a fencing
       * request or an entry which it no longer contains. The only way to resolve this is to
       * open/recover the ledger. This bookie should not take part in the recovery, so the bookie
       * must be marked as in limbo. This will stop the bookie from responding to read requests for
       * that ledger, so clients will not be able to take into account the response of the bookie
       * during recovery. Effectively we are telling the client that we don't know whether we had
       * certain entries or not, so go look elsewhere.
       * We also fence all ledgers with this bookie in the last segment, to prevent any new writes,
       * so that after the limbo state is cleared, we won't accept any new writes.
  
       * We only need to consider final ensembles in non-closed ledgers at the moment of time that
       * the preboot check commences. If this bookie is added to a new ensemble after that point in
       * time, we know that we haven't received any entries for that segment, nor have we received
       * a fencing request, because we are still in the preboot sequence.
       */
      private CompletableFuture<Void> processPreBoot(long ledgerId, LedgerMetadata metadata,
                                                     String runId) {
          Map.Entry<Long, ? extends List<BookieId>> lastEnsemble = metadata.getAllEnsembles().lastEntry();
          CompletableFuture<Void> promise = new CompletableFuture<>();
          if (lastEnsemble == null) {
              log.error("Event: {}, runId: {}, metadata: {}, ledger: {}",
                      Events.INVALID_METADATA, runId, metadata, ledgerId);
              promise.completeExceptionally(
                      new IllegalStateException(
                              String.format("All metadata must have at least one ensemble, %d does not", ledgerId)));
              return promise;
          }
  
  
          if (!metadata.isClosed() && lastEnsemble.getValue().contains(bookieId)) {
              try {
                  log.info("Event: {}, runId: {}, metadata: {}, ledger: {}",
                          Events.MARK_LIMBO, runId, metadata, ledgerId);
                  ledgerStorage.setLimboState(ledgerId);
                  ledgerStorage.setFenced(ledgerId);
                  promise.complete(null);
              } catch (IOException ioe) {
                  log.info("Event: {}, runId: {}, metadata: {}, ledger: {}",
                          Events.LIMBO_OR_FENCE_ERROR, runId, metadata, ledgerId, ioe);
                  promise.completeExceptionally(ioe);
              }
          } else {
              promise.complete(null);
          }
          return promise;
      }
  
      static class LedgerResult {
          enum State {
              MISSING, ERROR, OK
          }
  
          static LedgerResult missing(long ledgerId) {
              return new LedgerResult(State.MISSING, ledgerId, null, null);
          }
  
          static LedgerResult ok(long ledgerId, LedgerMetadata metadata) {
              return new LedgerResult(State.OK, ledgerId, metadata, null);
          }
  
          static LedgerResult error(long ledgerId, LedgerMetadata metadata, Throwable t) {
              return new LedgerResult(State.ERROR, ledgerId, metadata, t);
          }
  
          private final State state;
          private final long ledgerId;
          private final LedgerMetadata metadata;
          private final Throwable throwable;
  
          private LedgerResult(State state, long ledgerId,
                               LedgerMetadata metadata, Throwable throwable) {
              this.state = state;
              this.ledgerId = ledgerId;
              this.metadata = metadata;
              this.throwable = throwable;
          }
  
          boolean isMissing() {
              return state == State.MISSING;
          }
          boolean isOK() {
              return state == State.OK;
          }
          boolean isError() {
              return state == State.ERROR;
          }
          long getLedgerId() {
              return ledgerId;
          }
          LedgerMetadata getMetadata() {
              return metadata;
          }
          Throwable getThrowable() {
              return throwable;
          }
      }
  
      /**
       * Check each ledger passed.
       * If the ledger is in limbo, recover it.
       * Check that the bookie has all entries that it is expected to have.
       * Copy any entries that are missing.
       * @return The set of results for all ledgers passed. A result can be OK, Missing or Error.
       *         OK and missing ledgers do not need to be looked at again. Error should be retried.
       */
      CompletableFuture<Set<LedgerResult>> checkAndRecoverLedgers(Map<Long, LedgerMetadata> ledgers,
                                                                  String runId) {
          CompletableFuture<Set<LedgerResult>> promise = new CompletableFuture<>();
          final Disposable disposable = Flowable.fromIterable(ledgers.entrySet())
                  .subscribeOn(scheduler, false)
                  .flatMapSingle((mapEntry) -> {
                              long ledgerId = mapEntry.getKey();
                              LedgerMetadata originalMetadata = mapEntry.getValue();
                              return recoverLedgerIfInLimbo(ledgerId, mapEntry.getValue(), runId)
                                      .map(newMetadata -> LedgerResult.ok(ledgerId, newMetadata))
                                      .onErrorReturn(t -> LedgerResult.error(ledgerId, originalMetadata, t))
                                      .defaultIfEmpty(LedgerResult.missing(ledgerId))
                                      .flatMap((res) -> {
                                          try {
                                              if (res.isOK()) {
                                                  this.ledgerStorage.clearLimboState(ledgerId);
                                              }
                                              return Single.just(res);
                                          } catch (IOException ioe) {
                                              return Single.just(LedgerResult.error(res.getLedgerId(),
                                                      res.getMetadata(), ioe));
                                          }
                                      });
                          },
                          true /* delayErrors */,
                          MAX_INFLIGHT)
                  .flatMapSingle((res) -> {
                              if (res.isOK()) {
                                  return checkAndRecoverLedgerEntries(res.getLedgerId(),
                                          res.getMetadata(), runId)
                                          .map(ignore -> LedgerResult.ok(res.getLedgerId(),
                                                  res.getMetadata()))
                                          .onErrorReturn(t -> LedgerResult.error(res.getLedgerId(),
                                                  res.getMetadata(), t));
                              } else {
                                  return Single.just(res);
                              }
                          },
                          true /* delayErrors */,
                          1 /* copy 1 ledger at a time to keep entries together in entrylog */)
                  .collect(Collectors.toSet())
                  .subscribe(resolved -> promise.complete(resolved),
                          throwable -> promise.completeExceptionally(throwable));
          promise.whenComplete((result, ex) -> disposable.dispose());
          return promise;
      }
  
      /**
       * Run ledger recovery on all a ledger if it has been marked as in limbo.
       * @return a maybe with the most up-to-date metadata we have for the ledger.
       *         If the ledger has been deleted, returns empty.
       */
      Maybe<LedgerMetadata> recoverLedgerIfInLimbo(long ledgerId, LedgerMetadata origMetadata,
                                                   String runId) {
          try {
              if (!this.ledgerStorage.ledgerExists(ledgerId)) {
                  this.ledgerStorage.setMasterKey(ledgerId, new byte[0]);
              }
              if (this.ledgerStorage.hasLimboState(ledgerId)) {
                  log.info("Event: {}, runId: {}, metadata: {}, ledger: {}",
                          Events.RECOVER_LIMBO_LEDGER, runId, origMetadata, ledgerId);
                  return recoverLedger(ledgerId, runId)
                      .toMaybe()
                      .onErrorResumeNext(t -> {
                              if (t instanceof BKException.BKNoSuchLedgerExistsOnMetadataServerException) {
                                  log.info("Event: {}, runId: {}, metadata: {}, ledger: {}",
                                          Events.RECOVER_LIMBO_LEDGER_MISSING, runId, origMetadata, ledgerId);
                                  return Maybe.empty();
                              } else {
                                  log.info("Event: {}, runId: {}, metadata: {}, ledger: {}",
                                          Events.RECOVER_LIMBO_LEDGER_ERROR, runId, origMetadata, ledgerId);
                                  return Maybe.error(t);
                              }
                          });
              } else {
                  return Maybe.just(origMetadata);
              }
          } catch (IOException ioe) {
              return Maybe.error(ioe);
          }
      }
  
      Single<LedgerMetadata> recoverLedger(long ledgerId, String runId) {
          return Single.create((emitter) ->
                  admin.asyncOpenLedger(ledgerId, (rc, handle, ctx) -> {
                          if (rc != BKException.Code.OK) {
                              emitter.onError(BKException.create(rc));
                          } else {
                              LedgerMetadata metadata = handle.getLedgerMetadata();
                              handle.closeAsync().whenComplete((ignore, exception) -> {
                                      if (exception != null) {
                                          log.warn("Event: {}, runId: {}, ledger: {}",
                                                  Events.RECOVER_LIMBO_LEDGER_CLOSE_ERROR, runId, ledgerId, exception);
                                      }
                                  });
                              emitter.onSuccess(metadata);
                          }
                      }, null));
  
      }
  
      /**
       * Check whether the local storage has all the entries as specified in the metadata.
       * If not, copy them from other available nodes.
  
       * Returns a single value which is the ledgerId or an error if any entry failed to copy
       * should throw error if any entry failed to copy.
       */
      Single<Long> checkAndRecoverLedgerEntries(long ledgerId, LedgerMetadata metadata,
                                                String runId) {
          WriteSets writeSets = new WriteSets(metadata.getEnsembleSize(),
                                              metadata.getWriteQuorumSize());
  
          NavigableMap<Long, Integer> bookieIndices = metadata.getAllEnsembles()
              .entrySet().stream()
              .collect(ImmutableSortedMap.toImmutableSortedMap(Comparator.naturalOrder(),
                                                               e -> e.getKey(),
                                                               e -> e.getValue().indexOf(bookieId)));
  
          long lastKnownEntry;
          if (metadata.isClosed()) {
              lastKnownEntry = metadata.getLastEntryId();
          } else {
              // if ledger is not closed, last known entry is the last entry of
              // the penultimate ensemble
              lastKnownEntry = metadata.getAllEnsembles().lastEntry().getKey() - 1;
          }
          if (lastKnownEntry < 0) {
              return Single.just(ledgerId);
          }
  
          EntryCopier.Batch batch;
          try {
              batch = entryCopier.newBatch(ledgerId, metadata);
          } catch (IOException ioe) {
              return Single.error(ioe);
          }
          AtomicLong byteCount = new AtomicLong(0);
          AtomicInteger count = new AtomicInteger(0);
          AtomicInteger errorCount = new AtomicInteger(0);
          AtomicReference<Throwable> firstError = new AtomicReference<>(null);
          log.info("Event: {}, runId: {}, metadata: {}, ledger: {}",
                  Events.LEDGER_CHECK_AND_COPY_START, runId, metadata, ledgerId);
          return Flowable.rangeLong(0, lastKnownEntry + 1)
              .subscribeOn(scheduler, false)
              .flatMapMaybe((entryId) -> {
                      return maybeCopyEntry(writeSets, bookieIndices, ledgerId, entryId, batch)
                          .doOnError((t) -> {
                                  firstError.compareAndSet(null, t);
                                  errorCount.incrementAndGet();
                              });
                  }, true /* delayErrors */, MAX_ENTRIES_INFLIGHT)
              .doOnNext((bytes) -> {
                      byteCount.addAndGet(bytes);
                      count.incrementAndGet();
                  })
              .count() // do nothing with result, but gives a single even if empty
              .doOnTerminate(() -> {
                      if (firstError.get() != null) {
                          log.warn("Event: {}, runId: {}, metadata: {}, ledger: {}, entries: {}, bytes: {}, errors: {}",
                                  Events.LEDGER_CHECK_AND_COPY_END, runId,
                                  metadata, ledgerId, count.get(), byteCount.get(), firstError.get());
                      } else {
                          log.info("Event: {}, runId: {}, metadata: {}, ledger: {}, entries: {}, bytes: {}, errors: 0",
                                  Events.LEDGER_CHECK_AND_COPY_END, runId,
                                  metadata, ledgerId, count.get(), byteCount.get());
                      }
                  })
              .map(ignore -> ledgerId);
      }
  
      /**
       * @return the number of bytes copied.
       */
      Maybe<Long> maybeCopyEntry(WriteSets writeSets, NavigableMap<Long, Integer> bookieIndices,
                                       long ledgerId, long entryId, EntryCopier.Batch batch) {
          try {
              if (isEntryMissing(writeSets, bookieIndices, ledgerId, entryId)) {
                  return Maybe.fromCompletionStage(batch.copyFromAvailable(entryId));
              } else {
                  return Maybe.empty();
              }
          } catch (BookieException | IOException ioe) {
              return Maybe.error(ioe);
          }
      }
  
      boolean isEntryMissing(WriteSets writeSets, NavigableMap<Long, Integer> bookieIndices,
                             long ledgerId, long entryId) throws IOException, BookieException {
          int bookieIndexForEntry = bookieIndices.floorEntry(entryId).getValue();
          if (bookieIndexForEntry < 0) {
              return false;
          }
  
          return writeSets.getForEntry(entryId).contains(bookieIndexForEntry)
              && !ledgerStorage.entryExists(ledgerId, entryId);
      }
  
      static boolean ensemblesContainBookie(LedgerMetadata metadata, BookieId bookieId) {
          return metadata.getAllEnsembles().values().stream()
              .anyMatch(ensemble -> ensemble.contains(bookieId));
      }
  }
  