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 io.reactivex.rxjava3.core.Completable;
  import io.reactivex.rxjava3.core.Flowable;
  import io.reactivex.rxjava3.core.Scheduler;
  import io.reactivex.rxjava3.disposables.Disposable;
  import java.io.IOException;
  import java.util.Iterator;
  import java.util.concurrent.CompletableFuture;
  import java.util.concurrent.CompletionException;
  import java.util.concurrent.ExecutionException;
  import java.util.concurrent.TimeUnit;
  import java.util.function.BiFunction;
  import lombok.extern.slf4j.Slf4j;
  import org.apache.bookkeeper.client.BKException;
  import org.apache.bookkeeper.client.api.LedgerMetadata;
  import org.apache.bookkeeper.meta.LedgerManager;
  import org.apache.bookkeeper.versioning.Versioned;
  
  /**
   * An rxjava ledger metadata iterator.
   */
  @Slf4j
  public class MetadataAsyncIterator {
      private final Scheduler scheduler;
      private final LedgerManager ledgerManager;
      private final long zkTimeoutMs;
      private final int maxInFlight;
  
      MetadataAsyncIterator(Scheduler scheduler,
                            LedgerManager ledgerManager, int maxInFlight,
                            int zkTimeout, TimeUnit zkTimeoutUnit) {
          this.scheduler = scheduler;
          this.ledgerManager = ledgerManager;
          this.maxInFlight = maxInFlight;
          this.zkTimeoutMs = zkTimeoutUnit.toMillis(zkTimeout);
      }
  
  
      private static class FlatIterator {
          final LedgerManager.LedgerRangeIterator ranges;
          Iterator<Long> range = null;
          FlatIterator(LedgerManager.LedgerRangeIterator ranges) {
              this.ranges = ranges;
          }
          boolean hasNext() throws IOException {
              if (range == null || !range.hasNext()) {
                  if (ranges.hasNext()) {
                      range = ranges.next().getLedgers().iterator();
                  }
              }
              return range != null && range.hasNext();
          }
          Long next() throws IOException {
              return range.next();
          }
      }
  
      public CompletableFuture<Void> forEach(BiFunction<Long, LedgerMetadata, CompletableFuture<Void>> consumer) {
          CompletableFuture<Void> promise = new CompletableFuture<>();
          final Disposable disposable = Flowable.<Long, FlatIterator>generate(
                          () -> new FlatIterator(ledgerManager.getLedgerRanges(zkTimeoutMs)),
                          (iter, emitter) -> {
                              try {
                                  if (iter.hasNext()) {
                                      emitter.onNext(iter.next());
                                  } else {
                                      emitter.onComplete();
                                  }
                              } catch (Exception e) {
                                  emitter.onError(e);
                              }
                          })
                  .subscribeOn(scheduler)
                  .flatMapCompletable((ledgerId) -> Completable.fromCompletionStage(processOne(ledgerId, consumer)),
                          false /* delayErrors */,
                          maxInFlight)
                  .subscribe(() -> promise.complete(null),
                          t -> promise.completeExceptionally(unwrap(t)));
          promise.whenComplete((result, ex) -> disposable.dispose());
          return promise;
      }
  
      private CompletableFuture<Void> processOne(long ledgerId,
                                                 BiFunction<Long, LedgerMetadata, CompletableFuture<Void>> consumer) {
          return ledgerManager.readLedgerMetadata(ledgerId)
              .thenApply(Versioned::getValue)
              .thenCompose((metadata) -> consumer.apply(ledgerId, metadata))
              .exceptionally((e) -> {
                      Throwable realException = unwrap(e);
                      log.warn("Got exception processing ledger {}", ledgerId, realException);
                      if (realException instanceof BKException.BKNoSuchLedgerExistsOnMetadataServerException) {
                          return null;
                      } else {
                          throw new CompletionException(realException);
                      }
                  });
      }
  
      static Throwable unwrap(Throwable e) {
          if (e instanceof CompletionException || e instanceof ExecutionException) {
              return unwrap(e.getCause());
          }
          return e;
      }
  }
  