[fix][ml] Fix races in ledger initialization timeout handling

managed-ledger/src/main/java/org/apache/bookkeeper/mledger/impl/ManagedLedgerFactoryImpl.java

@@ -152,12 +152,50 @@ private static class PendingInitializeManagedLedger {
 
  private final ManagedLedgerImpl ledger;
  private final long createTimeMs;
+ private boolean done;
 
  PendingInitializeManagedLedger(ManagedLedgerImpl ledger) {
  this.ledger = ledger;
  this.createTimeMs = System.currentTimeMillis();
  }
 
+ public synchronized boolean done() {
+ if (done) {
+ return false;
+ }
+ done = true;
+ return true;
+ }
+
+ public synchronized boolean close() {
+ return close(null);
+ }
+
+ public synchronized boolean close(Runnable closeCallback) {
+ if (done) {
+ return false;
+ }
+ done = true;
+ ledger.asyncClose(new CloseCallback() {
+ @Override
+ public void closeComplete(Object ctx) {
+ if (closeCallback != null) {
+ closeCallback.run();
+ }
+ }
+
+ @Override
+ public void closeFailed(ManagedLedgerException exception, Object ctx) {
+ log.warn("[{}] Failed to a pending initialization managed ledger", ledger.getName(),
+ exception);
+ if (closeCallback != null) {
+ closeCallback.run();
+ }
+ }
+ }, null);
+ return true;
+ }
+
  }
 
  public ManagedLedgerFactoryImpl(MetadataStoreExtended metadataStore, ClientConfiguration bkClientConfiguration)
@@ -391,14 +429,17 @@ public void asyncOpen(final String name, final ManagedLedgerConfig config, final
  PendingInitializeManagedLedger pendingLedger = pendingInitializeLedgers.get(name);
  if (null != pendingLedger) {
  long pendingMs = System.currentTimeMillis() - pendingLedger.createTimeMs;
- if (pendingMs > TimeUnit.SECONDS.toMillis(config.getMetadataOperationsTimeoutSeconds())) {
+ if (pendingMs > TimeUnit.SECONDS.toMillis(config.getMetadataOperationsTimeoutSeconds())
+ && pendingInitializeLedgers.remove(name, pendingLedger)) {
  log.warn("[{}] Managed ledger has been pending in initialize state more than {} milliseconds,"
- + " remove it from cache to retry ...", name, pendingMs);
- ledgers.remove(name, existingFuture);
- pendingInitializeLedgers.remove(name, pendingLedger);
+ + " remove it from cache to retry ...", name, pendingMs);
+ pendingLedger.close(() -> {
+ if (ledgers.remove(name, existingFuture)) {
+ entryCacheManager.removeEntryCache(name);
+ }
+ });
  }
  }
-
  }
  }
 
@@ -420,15 +461,20 @@ public void asyncOpen(final String name, final ManagedLedgerConfig config, final
  newledger.initialize(new ManagedLedgerInitializeLedgerCallback() {
  @Override
  public void initializeComplete() {
- log.info("[{}] Successfully initialize managed ledger", name);
  pendingInitializeLedgers.remove(name, pendingLedger);
- future.complete(newledger);
-
- // May need to update the cursor position
- newledger.maybeUpdateCursorBeforeTrimmingConsumedLedger();
- // May need to trigger offloading
- if (config.isTriggerOffloadOnTopicLoad()) {
- newledger.maybeOffloadInBackground(NULL_OFFLOAD_PROMISE);
+ if (pendingLedger.done()) {
+ log.info("[{}] Successfully initialize managed ledger", name);
+ future.complete(newledger);
+
+ // May need to update the cursor position
+ newledger.maybeUpdateCursorBeforeTrimmingConsumedLedger();
+ // May need to trigger offloading
+ if (config.isTriggerOffloadOnTopicLoad()) {
+ newledger.maybeOffloadInBackground(NULL_OFFLOAD_PROMISE);
+ }
+ } else {
+ future.completeExceptionally(
+ new ManagedLedgerException("Managed ledger initialization timed out"));
  }
  }
 
@@ -437,26 +483,14 @@ public void initializeFailed(ManagedLedgerException e) {
  if (config.isCreateIfMissing()) {
  log.error("[{}] Failed to initialize managed ledger: {}", name, e.getMessage());
  }
-
  // Clean the map if initialization fails
- ledgers.remove(name, future);
- entryCacheManager.removeEntryCache(name);
-
  if (pendingInitializeLedgers.remove(name, pendingLedger)) {
- pendingLedger.ledger.asyncClose(new CloseCallback() {
- @Override
- public void closeComplete(Object ctx) {
- // no-op
- }
-
- @Override
- public void closeFailed(ManagedLedgerException exception, Object ctx) {
- log.warn("[{}] Failed to a pending initialization managed ledger", name,
- exception);
+ pendingLedger.close(() -> {
+ if (ledgers.remove(name, existingFuture)) {
+ entryCacheManager.removeEntryCache(name);
  }
- }, null);
+ });
  }
-
  future.completeExceptionally(e);
  }
  }, null);
@@ -612,8 +646,8 @@ public void closeFailed(ManagedLedgerException exception, Object ctx) {
 
  }, scheduledExecutor.chooseThread());
  //close pendingInitializeManagedLedger directly to make sure all callbacks is called.
- PendingInitializeManagedLedger pendingLedger = pendingInitializeLedgers.get(ledgerName);
- if (pendingLedger != null && !ledgerFuture.isDone()) {
+ PendingInitializeManagedLedger pendingLedger = pendingInitializeLedgers.remove(ledgerName);
+ if (pendingLedger != null && pendingLedger.close()) {
  ledgerFuture.completeExceptionally(new ManagedLedgerException.ManagedLedgerFactoryClosedException());
  }
  }