You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pulsar.apache.org by lh...@apache.org on 2026/10/21 07:17:39 UTC
(pulsar) 01/02: [fix][meta] Fix repeated metadata recovery failures after a ZooKeeper session expires (#26845)
This is an automated email from the ASF dual-hosted git repository.
lhotari pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/pulsar.git
commit 233d73dcde5aa39dd239ef12169a007284006a89
Author: Radiance' 波波 <39...@users.noreply.github.com>
AuthorDate: Wed Oct 21 07:56:13 2026 +0800
[fix][meta] Fix repeated metadata recovery failures after a ZooKeeper session expires (#26845)
Co-authored-by: hechengbo <ra...@gmail.com>
(cherry picked from commit 296ef674bbffa263a6aec0570ba3e3b1c2c6f66d)
---
.../metadata/impl/PulsarZooKeeperClient.java | 391 +++++++++++---
.../pulsar/metadata/impl/ZKMetadataStore.java | 32 +-
.../pulsar/metadata/impl/ZKSessionWatcher.java | 355 +++++++++---
.../metadata/impl/PulsarZooKeeperClientTest.java | 278 ++++++++++
.../pulsar/metadata/impl/ZKSessionWatcherTest.java | 594 +++++++++++++++++++--
5 files changed, 1464 insertions(+), 186 deletions(-)
diff --git a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/PulsarZooKeeperClient.java b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/PulsarZooKeeperClient.java
index 17d3824f46ce..b9de3557bb96 100644
--- a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/PulsarZooKeeperClient.java
+++ b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/PulsarZooKeeperClient.java
@@ -34,6 +34,7 @@ import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
import lombok.extern.slf4j.Slf4j;
import org.apache.bookkeeper.stats.NullStatsLogger;
@@ -62,6 +63,7 @@ import org.apache.zookeeper.Watcher;
import org.apache.zookeeper.Watcher.Event.EventType;
import org.apache.zookeeper.Watcher.Event.KeeperState;
import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.ZooKeeper.States;
import org.apache.zookeeper.client.ZKClientConfig;
import org.apache.zookeeper.data.ACL;
import org.apache.zookeeper.data.Stat;
@@ -72,6 +74,8 @@ import org.apache.zookeeper.server.quorum.QuorumPeerConfig;
*/
@Slf4j
public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoCloseable {
+ private static final long INITIAL_HANDLE_GENERATION = 0;
+ private static final long NO_HANDLE_GENERATION = -1;
private static final int DEFAULT_RETRY_EXECUTOR_THREAD_COUNT = 1;
@@ -83,7 +87,15 @@ public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoClo
// state for the zookeeper client
private final AtomicReference<ZooKeeper> zk = new AtomicReference<ZooKeeper>();
private final AtomicBoolean closed = new AtomicBoolean(false);
+ private final AtomicBoolean reconnectScheduled = new AtomicBoolean(false);
+ // Only accessed by the single-threaded connectExecutor.
+ private long nextHandleGeneration;
+ private final AtomicLong currentHandleGeneration = new AtomicLong();
+ private final Object reconnectLock = new Object();
+ private long reconnectingHandleGeneration = NO_HANDLE_GENERATION;
+ private long pendingReconnectHandleGeneration = NO_HANDLE_GENERATION;
private final ZooKeeperWatcherBase watcherManager;
+ private final InitialHandleWatcher initialHandleWatcher;
private final ScheduledExecutorService retryExecutor;
private final ExecutorService connectExecutor;
@@ -115,64 +127,118 @@ public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoClo
@Override
public void run() {
- try {
- ZooWorker.syncCallWithRetries(null, new ZooWorker.ZooCallable<ZooKeeper>() {
-
- @Override
- public ZooKeeper call() throws KeeperException, InterruptedException {
- log.info("Reconnecting zookeeper {}.", connectString);
- // close the previous one
- closeZkHandle();
-
- // ZooKeeper can deliver SyncConnected after createZooKeeper() returns but before zk.set(newZk)
- // publishes the new instance. Hold these events until the new instance is published, so child
- // watchers never observe a new-session event while PulsarZooKeeperClient still points at the
- // old handle.
- CountDownLatch newZkSetLatch = new CountDownLatch(1);
- Watcher forwardEventsWatcher = event -> {
+ while (true) {
+ try {
+ ZooWorker.syncCallWithRetries(null, new ZooWorker.ZooCallable<ZooKeeper>() {
+
+ @Override
+ public ZooKeeper call() throws KeeperException, InterruptedException {
+ if (closed.get()) {
+ throw new InterruptedException("ZooKeeper client is closed");
+ }
+ log.info("Reconnecting zookeeper {}.", connectString);
+ // close the previous one
+ closeZkHandle();
+
+ // ZooKeeper can deliver SyncConnected after createZooKeeper() returns but before
+ // zk.set(newZk) publishes the new instance. Hold these events until the new instance is
+ // published, so child watchers never observe a new-session event while
+ // PulsarZooKeeperClient still points at the old handle.
+ CountDownLatch newZkSetLatch = new CountDownLatch(1);
+ CountDownLatch newZkConnectedLatch = new CountDownLatch(1);
+ AtomicBoolean newZkPublished = new AtomicBoolean(false);
+ AtomicReference<ZooKeeper> newZkRef = new AtomicReference<>();
+ long newZkGeneration = ++nextHandleGeneration;
+ Watcher forwardEventsWatcher = event -> {
+ try {
+ boolean awaited = newZkSetLatch.await(sessionTimeoutMs, TimeUnit.MILLISECONDS);
+ if (!awaited) {
+ log.warn("Timed out waiting for ZooKeeper instance to be published before "
+ + "forwarding event {}", event);
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ log.warn("Interrupted while waiting for ZooKeeper instance to be published, "
+ + "event {}", event, e);
+ return;
+ }
+ if (!newZkPublished.get()) {
+ return;
+ }
+ if (event.getState() == KeeperState.SyncConnected
+ || (allowReadOnlyMode
+ && event.getState() == KeeperState.ConnectedReadOnly)) {
+ // The ZooKeeper connection event is authoritative. Child watchers are observers
+ // and must not delay or prevent the connection attempt from completing.
+ newZkConnectedLatch.countDown();
+ }
+ ZooKeeper eventSource = newZkRef.get();
+ watcherManager.process(new HandleStateWatchedEvent(
+ event, newZkGeneration, eventSource.getSessionId(), eventSource.getState()));
+ };
+
+ ZooKeeper newZk;
try {
- boolean awaited = newZkSetLatch.await(sessionTimeoutMs, TimeUnit.MILLISECONDS);
- if (!awaited) {
- log.warn("Timed out waiting for ZooKeeper instance to be published before "
- + "forwarding event {}", event);
+ newZk = createZooKeeper(forwardEventsWatcher);
+ } catch (IOException | QuorumPeerConfig.ConfigException e) {
+ log.error("Failed to create zookeeper instance to {} with config path {}",
+ connectString, configPath, e);
+ throw KeeperException.create(KeeperException.Code.CONNECTIONLOSS);
+ }
+ newZkRef.set(newZk);
+
+ // Publish the new instance before releasing the forwarding watcher. Wait on this handle's
+ // connection latch rather than the watcher manager's latch, which may belong to an old
+ // handle.
+ synchronized (reconnectLock) {
+ if (!closed.get()) {
+ zk.set(newZk);
+ // Publish the generation after the handle. A reader that observes the new
+ // generation must also observe the corresponding handle.
+ currentHandleGeneration.set(newZkGeneration);
+ newZkPublished.set(true);
}
- } catch (InterruptedException e) {
- Thread.currentThread().interrupt();
- log.warn("Interrupted while waiting for ZooKeeper instance to be published, event {}",
- event, e);
- return;
}
- watcherManager.process(event);
- };
-
- ZooKeeper newZk;
- try {
- newZk = createZooKeeper(forwardEventsWatcher);
- } catch (IOException | QuorumPeerConfig.ConfigException e) {
- log.error("Failed to create zookeeper instance to {} with config path {}",
- connectString, configPath, e);
- throw KeeperException.create(KeeperException.Code.CONNECTIONLOSS);
+ newZkSetLatch.countDown();
+ if (!newZkPublished.get()) {
+ newZk.close();
+ throw new InterruptedException("ZooKeeper client is closed");
+ }
+ if (!newZkConnectedLatch.await(sessionTimeoutMs, TimeUnit.MILLISECONDS)) {
+ throw KeeperException.create(KeeperException.Code.CONNECTIONLOSS);
+ }
+ log.info("ZooKeeper session {} is created to {}.",
+ Long.toHexString(newZk.getSessionId()), connectString);
+ return newZk;
}
- // Publish the new instance before releasing the forwarding watcher. waitForConnection() must
- // happen after countDown(), since it depends on the forwarded SyncConnected event.
- zk.set(newZk);
- newZkSetLatch.countDown();
- waitForConnection();
- log.info("ZooKeeper session {} is created to {}.",
- Long.toHexString(newZk.getSessionId()), connectString);
- return newZk;
- }
+ @Override
+ public String toString() {
+ return String.format("ZooKeeper Client Creator (%s)", connectString);
+ }
- @Override
- public String toString() {
- return String.format("ZooKeeper Client Creator (%s)", connectString);
+ }, connectRetryPolicy, rateLimiter, createClientStats);
+ } catch (Exception e) {
+ synchronized (reconnectLock) {
+ if (!closed.get()) {
+ log.error("Gave up reconnecting to ZooKeeper : ", e);
+ Runtime.getRuntime().exit(1);
+ }
}
+ }
- }, connectRetryPolicy, rateLimiter, createClientStats);
- } catch (Exception e) {
- log.error("Gave up reconnecting to ZooKeeper : ", e);
- Runtime.getRuntime().exit(1);
+ synchronized (reconnectLock) {
+ long currentGeneration = currentHandleGeneration.get();
+ if (!closed.get()
+ && pendingReconnectHandleGeneration == currentGeneration
+ && pendingReconnectHandleGeneration != reconnectingHandleGeneration) {
+ reconnectingHandleGeneration = pendingReconnectHandleGeneration;
+ pendingReconnectHandleGeneration = NO_HANDLE_GENERATION;
+ } else {
+ clearReconnectStateLocked();
+ return;
+ }
+ }
}
}
@@ -326,13 +392,30 @@ public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoClo
double rate,
boolean allowReadOnlyMode,
String configPath) throws IOException, QuorumPeerConfig.ConfigException {
- super(connectString, sessionTimeoutMs, watcherManager, allowReadOnlyMode,
+ this(connectString, sessionTimeoutMs, watcherManager,
+ new InitialHandleWatcher(watcherManager, sessionTimeoutMs), connectRetryPolicy,
+ operationRetryPolicy, statsLogger, retryExecThreadCount, rate, allowReadOnlyMode, configPath);
+ }
+
+ private PulsarZooKeeperClient(String connectString,
+ int sessionTimeoutMs,
+ ZooKeeperWatcherBase watcherManager,
+ InitialHandleWatcher initialHandleWatcher,
+ RetryPolicy connectRetryPolicy,
+ RetryPolicy operationRetryPolicy,
+ StatsLogger statsLogger,
+ int retryExecThreadCount,
+ double rate,
+ boolean allowReadOnlyMode,
+ String configPath) throws IOException, QuorumPeerConfig.ConfigException {
+ super(connectString, sessionTimeoutMs, initialHandleWatcher, allowReadOnlyMode,
configPath == null ? null : new ZKClientConfig(configPath));
this.connectString = connectString;
this.configPath = configPath;
this.sessionTimeoutMs = sessionTimeoutMs;
this.allowReadOnlyMode = allowReadOnlyMode;
this.watcherManager = watcherManager;
+ this.initialHandleWatcher = initialHandleWatcher;
this.connectRetryPolicy = connectRetryPolicy;
this.operationRetryPolicy = operationRetryPolicy;
this.rateLimiter = rate > 0 ? RateLimiter.create(rate) : null;
@@ -358,14 +441,23 @@ public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoClo
getACLStats = scopedStatsLogger.getOpStatsLogger("get_acl");
setACLStats = scopedStatsLogger.getOpStatsLogger("set_acl");
syncStats = scopedStatsLogger.getOpStatsLogger("sync");
+ initialHandleWatcher.setClient(this);
}
@Override
public void close() throws InterruptedException {
- closed.set(true);
- connectExecutor.shutdown();
+ synchronized (reconnectLock) {
+ closed.set(true);
+ pendingReconnectHandleGeneration = NO_HANDLE_GENERATION;
+ }
+ connectExecutor.shutdownNow();
retryExecutor.shutdown();
- closeZkHandle();
+ try {
+ closeZkHandle();
+ connectExecutor.awaitTermination(sessionTimeoutMs, TimeUnit.MILLISECONDS);
+ } finally {
+ clearReconnectState();
+ }
}
private void closeZkHandle() throws InterruptedException {
@@ -391,31 +483,196 @@ public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoClo
@Override
public void process(WatchedEvent event) {
- if (event.getType() == EventType.None
- && event.getState() == KeeperState.Expired) {
- onExpired();
+ long eventHandleGeneration;
+ States eventClientState;
+ if (event instanceof HandleStateWatchedEvent handleStateEvent) {
+ eventHandleGeneration = handleStateEvent.getHandleGeneration();
+ eventClientState = handleStateEvent.getClientState();
+ } else {
+ eventHandleGeneration = currentHandleGeneration.get();
+ eventClientState = getState();
+ }
+ if (event.getType() == EventType.None) {
+ if (event.getState() == KeeperState.Expired) {
+ scheduleReconnect("session expired", eventHandleGeneration);
+ } else if (event.getState() == KeeperState.AuthFailed
+ && eventClientState == States.AUTH_FAILED) {
+ scheduleReconnect("authentication failed", eventHandleGeneration);
+ }
}
}
- private void onExpired() {
- if (closed.get()) {
- // we don't schedule any tries if the client is closed.
- return;
+ private void scheduleReconnect(String reason, long eventHandleGeneration) {
+ synchronized (reconnectLock) {
+ if (closed.get()) {
+ // we don't schedule any tries if the client is closed.
+ return;
+ }
+ long currentGeneration = currentHandleGeneration.get();
+ if (eventHandleGeneration != currentGeneration) {
+ if (log.isDebugEnabled()) {
+ log.debug("Ignoring reconnect request from a stale ZooKeeper handle. eventHandleGeneration: {},"
+ + " currentHandleGeneration: {}, reason: {}", eventHandleGeneration, currentGeneration,
+ reason);
+ }
+ return;
+ }
+ if (reconnectScheduled.get()) {
+ if (eventHandleGeneration == reconnectingHandleGeneration
+ || eventHandleGeneration == pendingReconnectHandleGeneration) {
+ if (log.isDebugEnabled()) {
+ log.debug("Ignoring duplicate ZooKeeper reconnect request. handleGeneration: {}, reason: {}",
+ eventHandleGeneration, reason);
+ }
+ return;
+ }
+ pendingReconnectHandleGeneration = eventHandleGeneration;
+ return;
+ }
+ reconnectingHandleGeneration = eventHandleGeneration;
+ reconnectScheduled.set(true);
}
- log.info("ZooKeeper session {} is expired from {}.",
- Long.toHexString(getSessionId()), connectString);
+ log.info("Scheduling ZooKeeper reconnect for session {} to {}. reason: {}",
+ Long.toHexString(getSessionId()), connectString, reason);
try {
connectExecutor.execute(clientCreator);
} catch (RejectedExecutionException ree) {
+ clearReconnectState();
if (!closed.get()) {
log.error("ZooKeeper reconnect task is rejected : ", ree);
}
} catch (Exception t) {
+ clearReconnectState();
log.error("Failed to submit zookeeper reconnect task due to runtime exception : ", t);
}
}
+ private void clearReconnectState() {
+ synchronized (reconnectLock) {
+ clearReconnectStateLocked();
+ }
+ }
+
+ private void clearReconnectStateLocked() {
+ reconnectingHandleGeneration = NO_HANDLE_GENERATION;
+ pendingReconnectHandleGeneration = NO_HANDLE_GENERATION;
+ reconnectScheduled.set(false);
+ }
+
+ @VisibleForTesting
+ ZooKeeper getZooKeeperHandle() {
+ return zk.get();
+ }
+
+ @VisibleForTesting
+ void setZooKeeperHandle(ZooKeeper zkHandle) {
+ zk.set(zkHandle);
+ }
+
+ @VisibleForTesting
+ boolean isReconnectInProgress() {
+ return reconnectScheduled.get();
+ }
+
+ @VisibleForTesting
+ Watcher getInitialHandleWatcher() {
+ return initialHandleWatcher;
+ }
+
+ long getCurrentHandleGeneration() {
+ return currentHandleGeneration.get();
+ }
+
+ CurrentHandleIdentity getCurrentHandleIdentity() {
+ synchronized (reconnectLock) {
+ return getCurrentHandleIdentityLocked();
+ }
+ }
+
+ // The state update must be short and must not invoke listeners or perform blocking work while reconnectLock is
+ // held.
+ boolean runStateUpdateIfCurrentHandle(CurrentHandleIdentity expectedIdentity, Runnable stateUpdate) {
+ synchronized (reconnectLock) {
+ if (!expectedIdentity.equals(getCurrentHandleIdentityLocked())) {
+ return false;
+ }
+ stateUpdate.run();
+ return true;
+ }
+ }
+
+ private CurrentHandleIdentity getCurrentHandleIdentityLocked() {
+ ZooKeeper zkHandle = zk.get();
+ long sessionId = zkHandle == null ? super.getSessionId() : zkHandle.getSessionId();
+ return new CurrentHandleIdentity(currentHandleGeneration.get(), sessionId);
+ }
+
+ record CurrentHandleIdentity(long generation, long sessionId) {
+ }
+
+ static final class HandleStateWatchedEvent extends WatchedEvent {
+ private final long handleGeneration;
+ private final long sessionId;
+ private final States clientState;
+
+ HandleStateWatchedEvent(WatchedEvent event, long handleGeneration, long sessionId, States clientState) {
+ super(event.getType(), event.getState(), event.getPath(), event.getZxid());
+ this.handleGeneration = handleGeneration;
+ this.sessionId = sessionId;
+ this.clientState = clientState;
+ }
+
+ long getHandleGeneration() {
+ return handleGeneration;
+ }
+
+ long getSessionId() {
+ return sessionId;
+ }
+
+ States getClientState() {
+ return clientState;
+ }
+ }
+
+ private static final class InitialHandleWatcher implements Watcher {
+ private final ZooKeeperWatcherBase watcherManager;
+ private final int sessionTimeoutMs;
+ private final CountDownLatch clientSetLatch = new CountDownLatch(1);
+ private PulsarZooKeeperClient client;
+
+ private InitialHandleWatcher(ZooKeeperWatcherBase watcherManager, int sessionTimeoutMs) {
+ this.watcherManager = watcherManager;
+ this.sessionTimeoutMs = sessionTimeoutMs;
+ }
+
+ private void setClient(PulsarZooKeeperClient client) {
+ this.client = client;
+ clientSetLatch.countDown();
+ }
+
+ @Override
+ public void process(WatchedEvent event) {
+ try {
+ if (!clientSetLatch.await(sessionTimeoutMs, TimeUnit.MILLISECONDS)) {
+ log.warn("Timed out waiting for the initial ZooKeeper handle to be initialized, event {}", event);
+ return;
+ }
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ log.warn("Interrupted while waiting for the initial ZooKeeper handle to be initialized, event {}",
+ event, e);
+ return;
+ }
+
+ PulsarZooKeeperClient eventSource = client;
+ watcherManager.process(new HandleStateWatchedEvent(
+ event, INITIAL_HANDLE_GENERATION,
+ eventSource.getInitialHandleSessionId(), eventSource.getInitialHandleState()));
+ }
+ }
+
/**
* A runnable that retries zookeeper operations.
*/
@@ -455,6 +712,10 @@ public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoClo
return zkHandle.getSessionId();
}
+ private long getInitialHandleSessionId() {
+ return super.getSessionId();
+ }
+
@Override
public byte[] getSessionPasswd() {
ZooKeeper zkHandle = zk.get();
@@ -740,6 +1001,10 @@ public class PulsarZooKeeperClient extends ZooKeeper implements Watcher, AutoClo
}
}
+ private States getInitialHandleState() {
+ return super.getState();
+ }
+
@Override
public String toString() {
ZooKeeper zkHandle = zk.get();
diff --git a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java
index 1ba8eb4e3bb5..06516f86ea07 100644
--- a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java
+++ b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKMetadataStore.java
@@ -122,9 +122,22 @@ public class ZKMetadataStore extends AbstractBatchedMetadataStore
}
}
- private void processSessionWatcher(WatchedEvent event) {
+ @VisibleForTesting
+ void processSessionWatcher(WatchedEvent event) {
if (sessionWatcher != null) {
- processEvent(sessionWatcher::process, event);
+ ZKSessionWatcher.SessionEventSnapshot stateSnapshot;
+ if (event instanceof PulsarZooKeeperClient.HandleStateWatchedEvent handleStateEvent) {
+ stateSnapshot = new ZKSessionWatcher.SessionEventSnapshot(event,
+ handleStateEvent.getHandleGeneration(), handleStateEvent.getSessionId(),
+ handleStateEvent.getClientState());
+ } else {
+ long handleGeneration = zkc instanceof PulsarZooKeeperClient
+ ? ((PulsarZooKeeperClient) zkc).getCurrentHandleGeneration()
+ : ZKSessionWatcher.UNKNOWN_HANDLE_GENERATION;
+ stateSnapshot = new ZKSessionWatcher.SessionEventSnapshot(
+ event, handleGeneration, zkc.getSessionId(), zkc.getState());
+ }
+ processEvent(sessionWatcher::process, stateSnapshot);
}
}
@@ -141,15 +154,26 @@ public class ZKMetadataStore extends AbstractBatchedMetadataStore
}
@VisibleForTesting
- @SneakyThrows
public ZKMetadataStore(ZooKeeper zkc, MetadataStoreConfig config, boolean isZkManaged) {
+ this(zkc, config, isZkManaged, true);
+ }
+
+ @VisibleForTesting
+ ZKMetadataStore(ZooKeeper zkc, boolean startSessionWatcherScheduler) {
+ this(zkc, MetadataStoreConfig.builder().build(), false, startSessionWatcherScheduler);
+ }
+
+ @SneakyThrows
+ private ZKMetadataStore(ZooKeeper zkc, MetadataStoreConfig config, boolean isZkManaged,
+ boolean startSessionWatcherScheduler) {
super(config);
this.zkConnectString = null;
this.rootPath = null;
this.metadataStoreConfig = null;
this.isZkManaged = isZkManaged;
this.zkc = zkc;
- this.sessionWatcher = new ZKSessionWatcher(zkc, this::receivedSessionEvent);
+ this.sessionWatcher = new ZKSessionWatcher(
+ zkc, this::receivedSessionEvent, startSessionWatcherScheduler);
zkc.addWatch("/", eventWatcher, AddWatchMode.PERSISTENT_RECURSIVE);
initBatchStrategy();
}
diff --git a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKSessionWatcher.java b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKSessionWatcher.java
index 7b6739799b2d..6e5882419011 100644
--- a/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKSessionWatcher.java
+++ b/pulsar-metadata/src/main/java/org/apache/pulsar/metadata/impl/ZKSessionWatcher.java
@@ -19,7 +19,10 @@
package org.apache.pulsar.metadata.impl;
import static org.apache.pulsar.common.util.Runnables.catchingAndLoggingThrowables;
+import com.google.common.annotations.VisibleForTesting;
import io.netty.util.concurrent.DefaultThreadFactory;
+import java.util.ArrayList;
+import java.util.List;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Executors;
import java.util.concurrent.RejectedExecutionException;
@@ -34,6 +37,7 @@ import org.apache.zookeeper.AsyncCallback.StatCallback;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
import org.apache.zookeeper.Watcher;
+import org.apache.zookeeper.Watcher.Event.KeeperState;
import org.apache.zookeeper.ZooKeeper;
/**
@@ -41,6 +45,7 @@ import org.apache.zookeeper.ZooKeeper;
*/
@Slf4j
public class ZKSessionWatcher implements AutoCloseable, Watcher {
+ static final long UNKNOWN_HANDLE_GENERATION = -1;
private final ZooKeeper zk;
private SessionEvent currentStatus;
@@ -56,6 +61,9 @@ public class ZKSessionWatcher implements AutoCloseable, Watcher {
private final ScheduledFuture<?> task;
private long disconnectedAt = 0;
+ // Highest handle generation whose SyncConnected event has been processed. A terminal event from an older handle
+ // remains relevant until a newer handle has actually reached the session watcher, not merely been published.
+ private long lastConnectedHandleGeneration = UNKNOWN_HANDLE_GENERATION;
// The id of a session that a listener invalidated and that PulsarZooKeeperClient is replacing, or 0. Events of
// that session, such as a connection check that still succeeds before the session is closed, must not
@@ -63,73 +71,145 @@ public class ZKSessionWatcher implements AutoCloseable, Watcher {
private long replacedSessionId = 0;
public ZKSessionWatcher(ZooKeeper zk, Consumer<SessionEvent> sessionListener) {
+ this(zk, sessionListener, true);
+ }
+
+ @VisibleForTesting
+ ZKSessionWatcher(ZooKeeper zk, Consumer<SessionEvent> sessionListener, boolean startScheduler) {
this.zk = zk;
this.monitorTimeoutMillis = zk.getSessionTimeout() * 5 / 6;
this.tickTimeMillis = zk.getSessionTimeout() / 15;
this.sessionListener = sessionListener;
+ this.lastConnectedHandleGeneration = zk instanceof PulsarZooKeeperClient
+ ? ((PulsarZooKeeperClient) zk).getCurrentHandleGeneration()
+ : UNKNOWN_HANDLE_GENERATION;
this.scheduler = Executors
.newSingleThreadScheduledExecutor(new DefaultThreadFactory("metadata-store-zk-session-watcher"));
- this.task =
- scheduler.scheduleWithFixedDelay(
+ this.task = startScheduler
+ ? scheduler.scheduleWithFixedDelay(
catchingAndLoggingThrowables(this::checkConnectionStatus), tickTimeMillis,
tickTimeMillis,
- TimeUnit.MILLISECONDS);
+ TimeUnit.MILLISECONDS)
+ : null;
this.currentStatus = SessionEvent.SessionReestablished;
}
@Override
public void close() throws Exception {
- task.cancel(true);
+ if (task != null) {
+ task.cancel(true);
+ }
scheduler.shutdownNow();
scheduler.awaitTermination(10, TimeUnit.SECONDS);
}
- // task that runs every TICK_TIME to check zk connection
- // NOT ThreadSafe:
- // If zk client can't ensure the order, it may lead to problems.
- // Currently,we only use it in single thread, it will be fine. but we shouldn't leave any potential problems
- // in the future.
- private void checkConnectionStatus() {
+ // Probe the current handle every TICK_TIME. The handle identity is checked atomically with the state update so a
+ // delayed result cannot overwrite the state established by a replacement handle.
+ @VisibleForTesting
+ void checkConnectionStatus() {
try {
- long checkedSessionId = zk.getSessionId();
- CompletableFuture<Watcher.Event.KeeperState> future = new CompletableFuture<>();
+ HandleIdentity checkedHandleIdentity = getCurrentHandleIdentity();
+ CompletableFuture<StateSnapshot> future = new CompletableFuture<>();
zk.exists("/", false, (StatCallback) (rc, path, ctx, stat) -> {
+ ZooKeeper.States clientState = zk.getState();
switch (KeeperException.Code.get(rc)) {
case CONNECTIONLOSS:
- future.complete(Watcher.Event.KeeperState.Disconnected);
+ future.complete(new StateSnapshot(KeeperState.Disconnected, clientState));
break;
case SESSIONEXPIRED:
- future.complete(Watcher.Event.KeeperState.Expired);
+ future.complete(new StateSnapshot(KeeperState.Expired, clientState));
+ break;
+
+ case AUTHFAILED:
+ future.complete(new StateSnapshot(KeeperState.AuthFailed, clientState));
break;
case OK:
default:
- future.complete(Watcher.Event.KeeperState.SyncConnected);
+ future.complete(new StateSnapshot(getKeeperState(clientState), clientState));
}
}, null);
- Watcher.Event.KeeperState zkClientState;
+ StateSnapshot stateSnapshot;
try {
- zkClientState = future.get(tickTimeMillis, TimeUnit.MILLISECONDS);
+ stateSnapshot = future.get(tickTimeMillis, TimeUnit.MILLISECONDS);
} catch (TimeoutException e) {
// Consider zk disconnection if zk operation takes more than TICK_TIME
- zkClientState = Watcher.Event.KeeperState.Disconnected;
+ stateSnapshot = new StateSnapshot(KeeperState.Disconnected, zk.getState());
}
- checkStateIfSameSession(checkedSessionId, zkClientState);
+ checkStateIfSameSession(checkedHandleIdentity, stateSnapshot);
} catch (RejectedExecutionException | InterruptedException e) {
- task.cancel(true);
+ if (task != null) {
+ task.cancel(true);
+ }
} catch (Throwable t) {
log.warn("Error while checking ZK connection status", t);
}
}
@Override
- public synchronized void process(WatchedEvent event) {
- log.info("Got ZK session watch event: {}", event);
- checkState(event.getState());
+ public void process(WatchedEvent event) {
+ long handleGeneration = zk instanceof PulsarZooKeeperClient
+ ? ((PulsarZooKeeperClient) zk).getCurrentHandleGeneration()
+ : UNKNOWN_HANDLE_GENERATION;
+ process(new SessionEventSnapshot(event, handleGeneration, zk.getSessionId(), zk.getState()));
+ }
+
+ synchronized void process(SessionEventSnapshot stateSnapshot) {
+ long currentHandleGeneration = zk instanceof PulsarZooKeeperClient
+ ? ((PulsarZooKeeperClient) zk).getCurrentHandleGeneration()
+ : UNKNOWN_HANDLE_GENERATION;
+ boolean staleHandle = stateSnapshot.handleGeneration() != UNKNOWN_HANDLE_GENERATION
+ && currentHandleGeneration != UNKNOWN_HANDLE_GENERATION
+ && stateSnapshot.handleGeneration() != currentHandleGeneration;
+ if (staleHandle && (!isTerminalInvalidation(stateSnapshot)
+ || stateSnapshot.handleGeneration() < lastConnectedHandleGeneration)) {
+ if (log.isDebugEnabled()) {
+ log.debug("Ignoring ZooKeeper session event {} from a stale handle. eventHandleGeneration: {},"
+ + " currentHandleGeneration: {}, eventSessionId: {}, currentSessionId: {}",
+ stateSnapshot.event(), stateSnapshot.handleGeneration(), currentHandleGeneration,
+ stateSnapshot.sessionId(), zk.getSessionId());
+ }
+ return;
+ }
+ if (!staleHandle && stateSnapshot.event().getState() == KeeperState.SyncConnected
+ && stateSnapshot.handleGeneration() != UNKNOWN_HANDLE_GENERATION) {
+ lastConnectedHandleGeneration = Math.max(lastConnectedHandleGeneration,
+ stateSnapshot.handleGeneration());
+ }
+ log.info("Got ZK session watch event: {}", stateSnapshot.event());
+ checkState(stateSnapshot.event().getState(), stateSnapshot.clientState(), stateSnapshot.sessionId(),
+ sessionListener);
+ }
+
+ private static boolean isTerminalInvalidation(SessionEventSnapshot stateSnapshot) {
+ return stateSnapshot.event().getState() == KeeperState.Expired
+ || (stateSnapshot.event().getState() == KeeperState.AuthFailed
+ && stateSnapshot.clientState() == ZooKeeper.States.AUTH_FAILED);
+ }
+
+ private static KeeperState getKeeperState(ZooKeeper.States state) {
+ if (state == null) {
+ return KeeperState.Disconnected;
+ }
+ switch (state) {
+ case CONNECTED:
+ return KeeperState.SyncConnected;
+ case CONNECTEDREADONLY:
+ return KeeperState.ConnectedReadOnly;
+ case AUTH_FAILED:
+ return KeeperState.AuthFailed;
+ case CLOSED:
+ return KeeperState.Closed;
+ case CONNECTING:
+ case ASSOCIATING:
+ case NOT_CONNECTED:
+ default:
+ return KeeperState.Disconnected;
+ }
}
synchronized void setSessionInvalid() {
@@ -145,55 +225,108 @@ public class ZKSessionWatcher implements AutoCloseable, Watcher {
replacedSessionId = zk.getSessionId();
}
- // PulsarZooKeeperClient publishes the new ZooKeeper instance before forwarding the corresponding session event to
- // watcherManager, so zk.set(newZk) happens-before this watcher observes the new-session event. Keep the session-id
- // check and state transition in the same synchronized section to prevent stale async probes from racing with that
- // event and overwriting the state of the newly established session.
- private synchronized void checkStateIfSameSession(long checkedSessionId,
- Watcher.Event.KeeperState zkClientState) {
- long currentSessionId = zk.getSessionId();
- if (checkedSessionId != currentSessionId) {
- log.warn("Ignoring ZooKeeper session state from a stale session. checkedSessionId: {},"
- + " currentSessionId: {}, zkClientState: {}",
- checkedSessionId, currentSessionId, zkClientState);
+ private HandleIdentity getCurrentHandleIdentity() {
+ if (zk instanceof PulsarZooKeeperClient) {
+ PulsarZooKeeperClient.CurrentHandleIdentity identity =
+ ((PulsarZooKeeperClient) zk).getCurrentHandleIdentity();
+ return new HandleIdentity(identity.generation(), identity.sessionId());
+ }
+ return new HandleIdentity(UNKNOWN_HANDLE_GENERATION, zk.getSessionId());
+ }
+
+ // Capture and compare the handle identity atomically with handle publication. The operation can be retried on a
+ // replacement handle by PulsarZooKeeperClient, but its result is safe to use only while the captured identity is
+ // still current. Keep this check and the state transition in the same synchronized section so a stale probe cannot
+ // overwrite a newly established session.
+ private synchronized void checkStateIfSameSession(HandleIdentity checkedHandleIdentity,
+ StateSnapshot stateSnapshot) {
+ if (zk instanceof PulsarZooKeeperClient) {
+ PulsarZooKeeperClient.CurrentHandleIdentity expectedIdentity =
+ new PulsarZooKeeperClient.CurrentHandleIdentity(checkedHandleIdentity.handleGeneration(),
+ checkedHandleIdentity.sessionId());
+ List<SessionEvent> notifications = new ArrayList<>(2);
+ boolean processed = ((PulsarZooKeeperClient) zk).runStateUpdateIfCurrentHandle(expectedIdentity,
+ () -> processCurrentProbeState(checkedHandleIdentity, stateSnapshot, notifications::add));
+ if (!processed) {
+ logStaleProbe(checkedHandleIdentity, getCurrentHandleIdentity(), stateSnapshot);
+ return;
+ }
+ notifySessionListener(notifications);
return;
}
- checkState(zkClientState);
+
+ HandleIdentity currentHandleIdentity = getCurrentHandleIdentity();
+ if (checkedHandleIdentity.sessionId() != currentHandleIdentity.sessionId()) {
+ logStaleProbe(checkedHandleIdentity, currentHandleIdentity, stateSnapshot);
+ return;
+ }
+ processCurrentProbeState(checkedHandleIdentity, stateSnapshot, sessionListener);
}
- private synchronized void checkState(Watcher.Event.KeeperState zkClientState) {
+ private void processCurrentProbeState(HandleIdentity checkedHandleIdentity, StateSnapshot stateSnapshot,
+ Consumer<SessionEvent> notificationSink) {
+ if (stateSnapshot.keeperState() == KeeperState.SyncConnected
+ && checkedHandleIdentity.handleGeneration() != UNKNOWN_HANDLE_GENERATION) {
+ lastConnectedHandleGeneration = Math.max(lastConnectedHandleGeneration,
+ checkedHandleIdentity.handleGeneration());
+ }
+ checkState(stateSnapshot.keeperState(), stateSnapshot.clientState(), checkedHandleIdentity.sessionId(),
+ notificationSink);
+ }
+
+ private void logStaleProbe(HandleIdentity checkedHandleIdentity, HandleIdentity currentHandleIdentity,
+ StateSnapshot stateSnapshot) {
+ if (log.isDebugEnabled()) {
+ log.debug("Ignoring ZooKeeper session state from a stale ZooKeeper handle. checkedSessionId: {},"
+ + " currentSessionId: {}, checkedHandleGeneration: {}, currentHandleGeneration: {},"
+ + " zkClientState: {}", checkedHandleIdentity.sessionId(),
+ currentHandleIdentity.sessionId(), checkedHandleIdentity.handleGeneration(),
+ currentHandleIdentity.handleGeneration(), stateSnapshot.keeperState());
+ }
+ }
+
+ private synchronized void checkState(KeeperState zkClientState, ZooKeeper.States clientState, long sessionId,
+ Consumer<SessionEvent> notificationSink) {
switch (zkClientState) {
case Expired:
if (currentStatus != SessionEvent.SessionLost) {
log.error("ZooKeeper session expired");
currentStatus = SessionEvent.SessionLost;
- sessionListener.accept(currentStatus);
+ notificationSink.accept(currentStatus);
}
break;
case Disconnected:
- if (disconnectedAt == 0) {
- // this is the first disconnect event, we should monitor the time out from now, so we record the
- // time of disconnect
- disconnectedAt = System.nanoTime();
- }
+ handleDisconnected(notificationSink);
+ break;
- long timeRemainingMillis = monitorTimeoutMillis
- - TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - disconnectedAt);
- if (timeRemainingMillis <= 0 && currentStatus != SessionEvent.SessionLost) {
- log.error("ZooKeeper session reconnection timeout. Notifying session is lost.");
- currentStatus = SessionEvent.SessionLost;
- sessionListener.accept(currentStatus);
- } else if (currentStatus != SessionEvent.SessionLost) {
- log.warn("[{}] ZooKeeper client is disconnected. Waiting to reconnect, time remaining = {} seconds",
- zk.getSessionId(), timeRemainingMillis / 1000.0);
- if (currentStatus == SessionEvent.SessionReestablished) {
- currentStatus = SessionEvent.ConnectionLost;
- sessionListener.accept(currentStatus);
+ case ConnectedReadOnly:
+ handleConnectedReadOnly(notificationSink);
+ break;
+
+ case SyncConnected:
+ handleSyncConnected(sessionId, notificationSink);
+ break;
+
+ case AuthFailed:
+ if (clientState == ZooKeeper.States.AUTH_FAILED) {
+ if (currentStatus != SessionEvent.SessionLost) {
+ log.error("ZooKeeper client authentication failed. Notifying session is lost.");
+ currentStatus = SessionEvent.SessionLost;
+ notificationSink.accept(currentStatus);
+ }
+ } else {
+ if (log.isDebugEnabled()) {
+ log.debug("ZooKeeper client authentication failed; waiting for the client connection state."
+ + " Current status: {}", currentStatus);
}
}
break;
+ case SaslAuthenticated:
+ log.info("ZooKeeper client SASL authentication completed. Current status: {}", currentStatus);
+ break;
+
case Closed:
// A ZooKeeper handle was closed: the previous handle when PulsarZooKeeperClient replaces an expired
// session with a new one, or the client when the metadata store closes. It says nothing about the
@@ -201,32 +334,98 @@ public class ZKSessionWatcher implements AutoCloseable, Watcher {
break;
default:
- if (currentStatus != SessionEvent.SessionReestablished) {
- if (replacedSessionId != 0 && zk.getSessionId() == replacedSessionId) {
- log.info("Ignoring the connection of a ZooKeeper session that is being replaced, sessionId={}",
- replacedSessionId);
- break;
- }
- replacedSessionId = 0;
- // since it reconnected to zoo keeper, we reset the disconnected time
- log.info("ZooKeeper client reconnection with server quorum. Current status: {}", currentStatus);
- disconnectedAt = 0;
-
- SessionEvent previousStatus = currentStatus;
- // Update the status before notifying the listener: a listener that can't use the new session, such
- // as ZKMetadataStore when it fails to recreate its persistent watch, calls setSessionReplaced() or
- // setSessionInvalid() while it's notified, and the status it sets must not be overwritten, or the next
- // session would not send SessionReestablished again.
- currentStatus = SessionEvent.SessionReestablished;
- try {
- sessionListener.accept(SessionEvent.Reconnected);
- } finally {
- if (previousStatus == SessionEvent.SessionLost) {
- sessionListener.accept(SessionEvent.SessionReestablished);
- }
+ log.warn("Ignoring ZooKeeper client state {} that does not indicate a reconnection. Current status: {}",
+ zkClientState, currentStatus);
+ break;
+ }
+ }
+
+ private void handleDisconnected(Consumer<SessionEvent> notificationSink) {
+ if (disconnectedAt == 0) {
+ // this is the first disconnect event, we should monitor the time out from now, so we record the
+ // time of disconnect
+ disconnectedAt = System.nanoTime();
+ }
+
+ long timeRemainingMillis = monitorTimeoutMillis
+ - TimeUnit.NANOSECONDS.toMillis(System.nanoTime() - disconnectedAt);
+ if (timeRemainingMillis <= 0 && currentStatus != SessionEvent.SessionLost) {
+ log.error("ZooKeeper session reconnection timeout. Notifying session is lost.");
+ currentStatus = SessionEvent.SessionLost;
+ notificationSink.accept(currentStatus);
+ } else if (currentStatus != SessionEvent.SessionLost) {
+ log.warn("[{}] ZooKeeper client is disconnected. Waiting to reconnect, time remaining = {} seconds",
+ zk.getSessionId(), timeRemainingMillis / 1000.0);
+ if (currentStatus == SessionEvent.SessionReestablished) {
+ currentStatus = SessionEvent.ConnectionLost;
+ notificationSink.accept(currentStatus);
+ }
+ }
+ }
+
+ private void handleConnectedReadOnly(Consumer<SessionEvent> notificationSink) {
+ disconnectedAt = 0;
+ if (currentStatus == SessionEvent.SessionReestablished) {
+ log.warn("[{}] ZooKeeper client is connected to a read-only server. Waiting for read-write connection",
+ zk.getSessionId());
+ currentStatus = SessionEvent.ConnectionLost;
+ notificationSink.accept(currentStatus);
+ } else {
+ if (log.isDebugEnabled()) {
+ log.debug("[{}] ZooKeeper client remains connected to a read-only server. Current status: {}",
+ zk.getSessionId(), currentStatus);
+ }
+ }
+ }
+
+ private void handleSyncConnected(long sessionId, Consumer<SessionEvent> notificationSink) {
+ if (currentStatus != SessionEvent.SessionReestablished) {
+ if (replacedSessionId != 0 && sessionId == replacedSessionId) {
+ log.info("Ignoring the connection of a ZooKeeper session that is being replaced, sessionId={}",
+ replacedSessionId);
+ return;
+ }
+ replacedSessionId = 0;
+
+ // since it reconnected to zoo keeper, we reset the disconnected time
+ log.info("ZooKeeper client reconnection with server quorum. Current status: {}", currentStatus);
+ disconnectedAt = 0;
+
+ SessionEvent previousStatus = currentStatus;
+ // Update the status before notifying the listener: a listener that can't use the new session, such
+ // as ZKMetadataStore when it fails to recreate its persistent watch, calls setSessionReplaced() or
+ // setSessionInvalid() while it's notified, and the status it sets must not be overwritten, or the next
+ // session would not send SessionReestablished again.
+ currentStatus = SessionEvent.SessionReestablished;
+ try {
+ notificationSink.accept(SessionEvent.Reconnected);
+ } finally {
+ if (previousStatus == SessionEvent.SessionLost) {
+ notificationSink.accept(SessionEvent.SessionReestablished);
}
}
- break;
}
}
+
+ private void notifySessionListener(List<SessionEvent> notifications) {
+ if (notifications.size() == 2) {
+ try {
+ sessionListener.accept(notifications.get(0));
+ } finally {
+ sessionListener.accept(notifications.get(1));
+ }
+ } else if (notifications.size() == 1) {
+ sessionListener.accept(notifications.get(0));
+ }
+ }
+
+ private record StateSnapshot(KeeperState keeperState, ZooKeeper.States clientState) {
+ }
+
+ private record HandleIdentity(long handleGeneration, long sessionId) {
+ }
+
+ record SessionEventSnapshot(WatchedEvent event, long handleGeneration, long sessionId,
+ ZooKeeper.States clientState) {
+ }
}
diff --git a/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/impl/PulsarZooKeeperClientTest.java b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/impl/PulsarZooKeeperClientTest.java
new file mode 100644
index 000000000000..365c550e5f55
--- /dev/null
+++ b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/impl/PulsarZooKeeperClientTest.java
@@ -0,0 +1,278 @@
+/*
+ * 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.pulsar.metadata.impl;
+
+import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.Answers.RETURNS_DEEP_STUBS;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.Mockito.doAnswer;
+import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.when;
+import java.time.Duration;
+import java.util.Collections;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import lombok.Cleanup;
+import org.apache.bookkeeper.stats.OpStatsLogger;
+import org.apache.bookkeeper.stats.StatsLogger;
+import org.apache.pulsar.metadata.BaseMetadataStoreTest;
+import org.apache.zookeeper.WatchedEvent;
+import org.apache.zookeeper.Watcher;
+import org.apache.zookeeper.Watcher.Event.EventType;
+import org.apache.zookeeper.Watcher.Event.KeeperState;
+import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.ZooKeeper.States;
+import org.awaitility.Awaitility;
+import org.testng.annotations.Test;
+
+public class PulsarZooKeeperClientTest extends BaseMetadataStoreTest {
+
+ @Test
+ public void testRunStateUpdateIfCurrentHandleRejectsReplacedHandleIdentity() throws Exception {
+ @Cleanup
+ PulsarZooKeeperClient zkClient = newClient();
+ PulsarZooKeeperClient.CurrentHandleIdentity initialIdentity = zkClient.getCurrentHandleIdentity();
+ AtomicInteger actionRuns = new AtomicInteger();
+
+ assertThat(zkClient.runStateUpdateIfCurrentHandle(initialIdentity, actionRuns::incrementAndGet)).isTrue();
+ assertThat(actionRuns).hasValue(1);
+
+ zkClient.process(sessionEvent(KeeperState.Expired));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> zkClient.getCurrentHandleGeneration() > initialIdentity.generation()
+ && zkClient.getZooKeeperHandle().getState() == States.CONNECTED
+ && !zkClient.isReconnectInProgress());
+
+ assertThat(zkClient.runStateUpdateIfCurrentHandle(initialIdentity, actionRuns::incrementAndGet)).isFalse();
+ assertThat(actionRuns).hasValue(1);
+ }
+
+ @Test
+ public void testTerminalAuthFailedRecreatesZookeeperHandle() throws Exception {
+ @Cleanup
+ PulsarZooKeeperClient zkClient = newClient();
+
+ zkClient.process(sessionEvent(KeeperState.Expired));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> zkClient.getZooKeeperHandle() != null
+ && zkClient.getZooKeeperHandle().getState() == States.CONNECTED
+ && !zkClient.isReconnectInProgress());
+ ZooKeeper previousZooKeeper = zkClient.getZooKeeperHandle();
+ ZooKeeper failedZooKeeper = mock(ZooKeeper.class);
+ when(failedZooKeeper.getState()).thenReturn(States.AUTH_FAILED);
+ when(failedZooKeeper.getSessionId()).thenReturn(previousZooKeeper.getSessionId());
+ zkClient.setZooKeeperHandle(failedZooKeeper);
+ previousZooKeeper.close();
+
+ zkClient.process(sessionEvent(KeeperState.AuthFailed));
+
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> zkClient.getZooKeeperHandle() != failedZooKeeper
+ && zkClient.getZooKeeperHandle().getState() == States.CONNECTED
+ && !zkClient.isReconnectInProgress());
+ }
+
+ @Test
+ public void testInvalidationWhileReconnectCompletesIsNotDropped() throws Exception {
+ CountDownLatch firstReconnectCompleted = new CountDownLatch(1);
+ CountDownLatch releaseFirstReconnect = new CountDownLatch(1);
+ StatsLogger statsLogger = mock(StatsLogger.class, RETURNS_DEEP_STUBS);
+ OpStatsLogger createClientStats = mock(OpStatsLogger.class);
+ when(statsLogger.scope("zk").getOpStatsLogger("create_client")).thenReturn(createClientStats);
+ doAnswer(invocation -> {
+ firstReconnectCompleted.countDown();
+ releaseFirstReconnect.await();
+ return null;
+ }).when(createClientStats).registerSuccessfulEvent(anyLong(), eq(TimeUnit.MICROSECONDS));
+
+ @Cleanup
+ PulsarZooKeeperClient zkClient = PulsarZooKeeperClient.newBuilder()
+ .connectString(zks.getConnectionString())
+ .sessionTimeoutMs(3000)
+ .statsLogger(statsLogger)
+ .build();
+
+ ZooKeeper firstReplacement;
+ try {
+ zkClient.process(sessionEvent(KeeperState.Expired));
+ assertThat(firstReconnectCompleted.await(5, TimeUnit.SECONDS)).isTrue();
+ firstReplacement = zkClient.getZooKeeperHandle();
+ assertThat(firstReplacement.getState()).isEqualTo(States.CONNECTED);
+ zkClient.process(sessionEvent(KeeperState.Expired));
+ } finally {
+ releaseFirstReconnect.countDown();
+ }
+
+ Awaitility.await().atMost(10, TimeUnit.SECONDS)
+ .until(() -> zkClient.getZooKeeperHandle() != firstReplacement
+ && zkClient.getZooKeeperHandle().getState() == States.CONNECTED
+ && !zkClient.isReconnectInProgress());
+ }
+
+ @Test
+ public void testDuplicateInvalidationDoesNotReplaceRecoveredHandleAgain() throws Exception {
+ CountDownLatch firstReconnectCompleted = new CountDownLatch(1);
+ CountDownLatch releaseFirstReconnect = new CountDownLatch(1);
+ StatsLogger statsLogger = mock(StatsLogger.class, RETURNS_DEEP_STUBS);
+ OpStatsLogger createClientStats = mock(OpStatsLogger.class);
+ when(statsLogger.scope("zk").getOpStatsLogger("create_client")).thenReturn(createClientStats);
+ doAnswer(invocation -> {
+ firstReconnectCompleted.countDown();
+ releaseFirstReconnect.await();
+ return null;
+ }).when(createClientStats).registerSuccessfulEvent(anyLong(), eq(TimeUnit.MICROSECONDS));
+
+ @Cleanup
+ PulsarZooKeeperClient zkClient = PulsarZooKeeperClient.newBuilder()
+ .connectString(zks.getConnectionString())
+ .sessionTimeoutMs(3000)
+ .statsLogger(statsLogger)
+ .build();
+ WatchedEvent invalidation = sessionEvent(KeeperState.Expired);
+
+ ZooKeeper firstReplacement;
+ try {
+ zkClient.getInitialHandleWatcher().process(invalidation);
+ assertThat(firstReconnectCompleted.await(5, TimeUnit.SECONDS)).isTrue();
+ firstReplacement = zkClient.getZooKeeperHandle();
+ assertThat(firstReplacement.getState()).isEqualTo(States.CONNECTED);
+
+ // The same terminal event from the original handle is delivered again before reconnect bookkeeping
+ // is cleared. It must not be interpreted as an invalidation of firstReplacement.
+ zkClient.getInitialHandleWatcher().process(invalidation);
+ } finally {
+ releaseFirstReconnect.countDown();
+ }
+
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> !zkClient.isReconnectInProgress());
+ assertThat(zkClient.getZooKeeperHandle()).isSameAs(firstReplacement);
+ assertThat(firstReplacement.getState()).isEqualTo(States.CONNECTED);
+ }
+
+ @Test
+ public void testBlockingChildWatcherDoesNotRejectConnectedReplacement() throws Exception {
+ CountDownLatch childWatcherEntered = new CountDownLatch(1);
+ CountDownLatch releaseChildWatcher = new CountDownLatch(1);
+ AtomicBoolean blockFirstReplacement = new AtomicBoolean(true);
+ Watcher blockingChildWatcher = event -> {
+ if (event instanceof PulsarZooKeeperClient.HandleStateWatchedEvent handleStateEvent
+ && handleStateEvent.getHandleGeneration() > 0
+ && event.getState() == KeeperState.SyncConnected
+ && blockFirstReplacement.compareAndSet(true, false)) {
+ childWatcherEntered.countDown();
+ try {
+ releaseChildWatcher.await(10, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ }
+ };
+
+ @Cleanup
+ PulsarZooKeeperClient zkClient = PulsarZooKeeperClient.newBuilder()
+ .connectString(zks.getConnectionString())
+ .sessionTimeoutMs(1000)
+ .watchers(Collections.singleton(blockingChildWatcher))
+ .build();
+
+ ZooKeeper connectedReplacement;
+ try {
+ zkClient.process(sessionEvent(KeeperState.Expired));
+ assertThat(childWatcherEntered.await(5, TimeUnit.SECONDS)).isTrue();
+ connectedReplacement = zkClient.getZooKeeperHandle();
+ assertThat(connectedReplacement.getState()).isEqualTo(States.CONNECTED);
+
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> !zkClient.isReconnectInProgress());
+ assertThat(zkClient.getZooKeeperHandle()).isSameAs(connectedReplacement);
+ } finally {
+ releaseChildWatcher.countDown();
+ }
+ }
+
+ @Test
+ public void testTerminalAuthFailedWaitsForReplacementConnection() throws Exception {
+ @Cleanup
+ PulsarZooKeeperClient zkClient = newClient();
+ zkClient.process(sessionEvent(KeeperState.Expired));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> zkClient.getZooKeeperHandle() != null
+ && zkClient.getZooKeeperHandle().getState() == States.CONNECTED
+ && !zkClient.isReconnectInProgress());
+ ZooKeeper previousZooKeeper = zkClient.getZooKeeperHandle();
+ ZooKeeper failedZooKeeper = mock(ZooKeeper.class);
+ when(failedZooKeeper.getState()).thenReturn(States.AUTH_FAILED);
+ zkClient.setZooKeeperHandle(failedZooKeeper);
+ previousZooKeeper.close();
+
+ zks.stop();
+ try {
+ zkClient.process(sessionEvent(KeeperState.AuthFailed));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> zkClient.getZooKeeperHandle() != failedZooKeeper);
+ Awaitility.await().during(Duration.ofMillis(500)).atMost(Duration.ofSeconds(2))
+ .until(zkClient::isReconnectInProgress);
+ } finally {
+ zks.start();
+ }
+
+ Awaitility.await().atMost(15, TimeUnit.SECONDS)
+ .until(() -> zkClient.getZooKeeperHandle().getState() == States.CONNECTED
+ && !zkClient.isReconnectInProgress());
+ }
+
+ @Test
+ public void testCloseStopsReconnectWhileReplacementIsConnecting() throws Exception {
+ PulsarZooKeeperClient zkClient = newClient();
+
+ zks.stop();
+ try {
+ zkClient.process(sessionEvent(KeeperState.Expired));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .until(() -> zkClient.getZooKeeperHandle() != null
+ && zkClient.isReconnectInProgress());
+ ZooKeeper replacementZooKeeper = zkClient.getZooKeeperHandle();
+
+ zkClient.close();
+
+ Awaitility.await().during(Duration.ofMillis(500)).atMost(Duration.ofSeconds(2))
+ .until(() -> zkClient.getZooKeeperHandle() == replacementZooKeeper
+ && replacementZooKeeper.getState() == States.CLOSED
+ && !zkClient.isReconnectInProgress());
+ } finally {
+ zkClient.close();
+ zks.start();
+ }
+ }
+
+ private PulsarZooKeeperClient newClient() throws Exception {
+ return PulsarZooKeeperClient.newBuilder()
+ .connectString(zks.getConnectionString())
+ .sessionTimeoutMs(3000)
+ .build();
+ }
+
+ private static WatchedEvent sessionEvent(KeeperState state) {
+ return new WatchedEvent(EventType.None, state, null);
+ }
+}
diff --git a/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/impl/ZKSessionWatcherTest.java b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/impl/ZKSessionWatcherTest.java
index 0fd200cf788a..699eef085f8b 100644
--- a/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/impl/ZKSessionWatcherTest.java
+++ b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/impl/ZKSessionWatcherTest.java
@@ -19,70 +19,582 @@
package org.apache.pulsar.metadata.impl;
import static org.assertj.core.api.Assertions.assertThat;
+import static org.mockito.ArgumentMatchers.any;
+import static org.mockito.ArgumentMatchers.eq;
+import static org.mockito.ArgumentMatchers.isNull;
+import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertTrue;
+import java.util.Arrays;
import java.util.List;
import java.util.concurrent.CopyOnWriteArrayList;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicReference;
-import lombok.Cleanup;
import org.apache.pulsar.metadata.api.extended.SessionEvent;
+import org.apache.zookeeper.AddWatchMode;
+import org.apache.zookeeper.AsyncCallback.StatCallback;
+import org.apache.zookeeper.AsyncCallback.VoidCallback;
+import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.WatchedEvent;
-import org.apache.zookeeper.Watcher;
+import org.apache.zookeeper.Watcher.Event.EventType;
+import org.apache.zookeeper.Watcher.Event.KeeperState;
import org.apache.zookeeper.ZooKeeper;
+import org.apache.zookeeper.ZooKeeper.States;
+import org.apache.zookeeper.client.ZKClientConfig;
+import org.awaitility.Awaitility;
import org.testng.annotations.Test;
+@Test
public class ZKSessionWatcherTest {
- /**
- * When ZKMetadataStore can't use a reestablished session, such as when it fails to recreate its persistent watch,
- * it marks the session replaced while it's notified and makes PulsarZooKeeperClient replace the session. The new
- * session must send SessionReestablished again, so that the listeners, such as the lock managers that revalidate
- * their locks, get it, and the replaced session's own events must not reestablish it first.
- */
- @Test(timeOut = 30_000)
- public void sessionReplacedWhileNotifiedIsReestablishedWithNextSession() throws Exception {
+ @Test
+ public void testClosedEventShouldNotBeTreatedAsReconnectedAfterSessionLost() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ try (ZKSessionWatcher watcher = newSessionWatcher(events)) {
+ watcher.setSessionInvalid();
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Closed, null));
+
+ assertTrue(events.isEmpty(),
+ "Closed is a terminal state for the old ZooKeeper handle and must not be treated as "
+ + "Reconnected or SessionReestablished, but received " + events);
+ }
+ }
+
+ @Test
+ public void testOnlySyncConnectedShouldBeTreatedAsReconnectedAfterSessionLost() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ try (ZKSessionWatcher watcher = newSessionWatcher(events)) {
+ watcher.setSessionInvalid();
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null));
+
+ assertEquals(events, Arrays.asList(SessionEvent.Reconnected, SessionEvent.SessionReestablished));
+ }
+ }
+
+ @Test
+ public void testNonSyncConnectedEventsShouldNotBeTreatedAsReconnectedAfterSessionLost() throws Exception {
+ for (KeeperState keeperState : Arrays.asList(
+ KeeperState.Disconnected,
+ KeeperState.AuthFailed,
+ KeeperState.ConnectedReadOnly,
+ KeeperState.SaslAuthenticated,
+ KeeperState.Closed)) {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ try (ZKSessionWatcher watcher = newSessionWatcher(events)) {
+ watcher.setSessionInvalid();
+ watcher.process(new WatchedEvent(EventType.None, keeperState, null));
+
+ assertTrue(events.stream().noneMatch(event -> event == SessionEvent.Reconnected
+ || event == SessionEvent.SessionReestablished),
+ keeperState + " must not be treated as a ZooKeeper reconnection event, but received "
+ + events);
+ }
+ }
+ }
+
+ @Test
+ public void testAuthFailedProbeShouldMarkSessionLostForTerminalClientState() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.AUTH_FAILED);
+ completeExistsWith(zk, KeeperException.Code.AUTHFAILED);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.checkConnectionStatus();
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.AuthFailed, null));
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(SessionEvent.SessionLost);
+ }
+ }
+
+ @Test
+ public void testAuthFailedStateFromSuccessfulProbeShouldMarkSessionLost() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.AUTH_FAILED);
+ completeExistsWith(zk, KeeperException.Code.OK);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.checkConnectionStatus();
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(SessionEvent.SessionLost);
+ }
+ }
+
+ @Test
+ public void testAuthFailedEventWhileConnectingShouldWaitForConnectionState() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.CONNECTING);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.AuthFailed, null));
+
+ assertThat(events).isEmpty();
+ }
+ }
+
+ @Test
+ public void testTerminalAuthFailedThenSyncConnectedReestablishesSession() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.AUTH_FAILED, States.CONNECTED);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.AuthFailed, null));
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Closed, null));
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null));
+
+ assertThat(events).containsExactly(
+ SessionEvent.SessionLost, SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+ }
+ }
+
+ @Test
+ public void testSessionReplacedWhileListenerIsNotifiedReestablishesWithNextSession() throws Exception {
ZooKeeper zk = mock(ZooKeeper.class);
- // a long session timeout, so that the periodic connection check doesn't run during the test
when(zk.getSessionTimeout()).thenReturn(600_000);
AtomicLong sessionId = new AtomicLong(1);
when(zk.getSessionId()).thenAnswer(invocation -> sessionId.get());
List<SessionEvent> events = new CopyOnWriteArrayList<>();
- AtomicReference<ZKSessionWatcher> watcher = new AtomicReference<>();
+ AtomicReference<ZKSessionWatcher> watcherRef = new AtomicReference<>();
AtomicBoolean replaceOnce = new AtomicBoolean(true);
- @Cleanup
- ZKSessionWatcher sessionWatcher = new ZKSessionWatcher(zk, event -> {
+
+ try (ZKSessionWatcher watcher = new ZKSessionWatcher(zk, event -> {
events.add(event);
if (event == SessionEvent.SessionReestablished && replaceOnce.getAndSet(false)) {
- watcher.get().setSessionReplaced();
+ watcherRef.get().setSessionReplaced();
+ }
+ }, false)) {
+ watcherRef.set(watcher);
+
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Expired, null));
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Closed, null));
+ assertThat(events).containsExactly(SessionEvent.SessionLost);
+
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null));
+ assertThat(events).containsExactly(
+ SessionEvent.SessionLost, SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+
+ // The old session may still report a successful probe or a close event while it is being replaced.
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null));
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Closed, null));
+ assertThat(events).containsExactly(
+ SessionEvent.SessionLost, SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+
+ sessionId.set(2);
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null));
+ assertThat(events).containsExactly(
+ SessionEvent.SessionLost, SessionEvent.Reconnected, SessionEvent.SessionReestablished,
+ SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+ }
+ }
+
+ @Test
+ public void testReplacedSessionIdRejectsOldEventBeforeGenerationPublication() throws Exception {
+ PulsarZooKeeperClient zk = mock(PulsarZooKeeperClient.class);
+ AtomicLong sessionId = new AtomicLong(1);
+ when(zk.getSessionTimeout()).thenReturn(600_000);
+ when(zk.getSessionId()).thenAnswer(invocation -> sessionId.get());
+ when(zk.getCurrentHandleGeneration()).thenReturn(1L);
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+
+ try (ZKSessionWatcher watcher = new ZKSessionWatcher(zk, events::add, false)) {
+ watcher.setSessionReplaced();
+
+ // Publishing a replacement writes the new handle before its generation. During that small window,
+ // the source generation alone still matches, so the source session ID must reject H1's event.
+ sessionId.set(2);
+ watcher.process(new ZKSessionWatcher.SessionEventSnapshot(
+ sessionEvent(KeeperState.SyncConnected), 1, 1, States.CONNECTED));
+
+ assertThat(events).isEmpty();
+ }
+ }
+
+ @Test
+ public void testQueuedTerminalAuthFailedIsProcessedBeforeReplacementConnects() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ PulsarZooKeeperClient zk = mock(PulsarZooKeeperClient.class);
+ AtomicLong sessionId = new AtomicLong(1);
+ AtomicLong handleGeneration = new AtomicLong(1);
+ AtomicReference<States> clientState = new AtomicReference<>(States.AUTH_FAILED);
+ when(zk.getSessionTimeout()).thenReturn(600_000);
+ when(zk.getSessionId()).thenAnswer(invocation -> sessionId.get());
+ when(zk.getCurrentHandleGeneration()).thenAnswer(invocation -> handleGeneration.get());
+ when(zk.getState()).thenAnswer(invocation -> clientState.get());
+ when(zk.getClientConfig()).thenReturn(new ZKClientConfig());
+ doAnswer(invocation -> {
+ VoidCallback callback = invocation.getArgument(3);
+ callback.processResult(KeeperException.Code.OK.intValue(), "/", null);
+ return null;
+ }).when(zk).addWatch(eq("/"), any(), eq(AddWatchMode.PERSISTENT_RECURSIVE), any(), isNull());
+ CountDownLatch eventExecutorBlocked = new CountDownLatch(1);
+ CountDownLatch releaseEventExecutor = new CountDownLatch(1);
+
+ try (ZKMetadataStore store = new ZKMetadataStore(zk, false)) {
+ store.registerSessionListener(events::add);
+ store.getEventExecutor().execute(() -> {
+ eventExecutorBlocked.countDown();
+ try {
+ releaseEventExecutor.await(5, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ });
+ assertThat(eventExecutorBlocked.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Queue H1's terminal failure, then publish H2 before the event executor processes it.
+ store.processSessionWatcher(new PulsarZooKeeperClient.HandleStateWatchedEvent(
+ sessionEvent(KeeperState.AuthFailed), 1, 1, States.AUTH_FAILED));
+ sessionId.set(2);
+ handleGeneration.set(2);
+ clientState.set(States.CONNECTED);
+ releaseEventExecutor.countDown();
+
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .untilAsserted(() -> assertThat(events).containsExactly(SessionEvent.SessionLost));
+
+ // Only H2's connection can recover the logical metadata session.
+ store.processSessionWatcher(new PulsarZooKeeperClient.HandleStateWatchedEvent(
+ sessionEvent(KeeperState.SyncConnected), 2, 2, States.CONNECTED));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .untilAsserted(() -> assertThat(events).containsExactly(
+ SessionEvent.SessionLost,
+ SessionEvent.Reconnected,
+ SessionEvent.SessionReestablished));
+ } finally {
+ releaseEventExecutor.countDown();
+ }
+ }
+
+ @Test
+ public void testTerminalEventFromOlderHandleIsIgnoredAfterReplacementConnects() throws Exception {
+ PulsarZooKeeperClient zk = mock(PulsarZooKeeperClient.class);
+ AtomicLong sessionId = new AtomicLong(1);
+ AtomicLong handleGeneration = new AtomicLong(1);
+ when(zk.getSessionTimeout()).thenReturn(600_000);
+ when(zk.getSessionId()).thenAnswer(invocation -> sessionId.get());
+ when(zk.getCurrentHandleGeneration()).thenAnswer(invocation -> handleGeneration.get());
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+
+ try (ZKSessionWatcher watcher = new ZKSessionWatcher(zk, events::add, false)) {
+ // H2 is published before H1's queued terminal event is processed. H1's session loss still matters
+ // until H2's own SyncConnected event reaches this watcher.
+ sessionId.set(2);
+ handleGeneration.set(2);
+ watcher.process(new ZKSessionWatcher.SessionEventSnapshot(
+ sessionEvent(KeeperState.Expired), 1, 1, States.CLOSED));
+
+ watcher.process(new ZKSessionWatcher.SessionEventSnapshot(
+ sessionEvent(KeeperState.SyncConnected), 2, 2, States.CONNECTED));
+
+ // H1 may deliver a duplicate terminal event after H2 has completed logical recovery.
+ watcher.process(new ZKSessionWatcher.SessionEventSnapshot(
+ sessionEvent(KeeperState.Expired), 1, 1, States.CLOSED));
+
+ assertThat(events).containsExactly(
+ SessionEvent.SessionLost,
+ SessionEvent.Reconnected,
+ SessionEvent.SessionReestablished);
+ }
+ }
+
+ @Test
+ public void testTerminalEventFromOlderHandleIsIgnoredAfterReplacementProbeReconnects() throws Exception {
+ PulsarZooKeeperClient zk = mock(PulsarZooKeeperClient.class);
+ AtomicLong sessionId = new AtomicLong(1);
+ AtomicLong handleGeneration = new AtomicLong(1);
+ when(zk.getSessionTimeout()).thenReturn(600_000);
+ when(zk.getSessionId()).thenAnswer(invocation -> sessionId.get());
+ when(zk.getCurrentHandleGeneration()).thenAnswer(invocation -> handleGeneration.get());
+ when(zk.getCurrentHandleIdentity()).thenAnswer(invocation ->
+ new PulsarZooKeeperClient.CurrentHandleIdentity(handleGeneration.get(), sessionId.get()));
+ when(zk.runStateUpdateIfCurrentHandle(any(), any())).thenAnswer(invocation -> {
+ PulsarZooKeeperClient.CurrentHandleIdentity expectedIdentity = invocation.getArgument(0);
+ if (!expectedIdentity.equals(zk.getCurrentHandleIdentity())) {
+ return false;
}
+ Runnable action = invocation.getArgument(1);
+ action.run();
+ return true;
});
- watcher.set(sessionWatcher);
-
- // the session expired, and PulsarZooKeeperClient replaced it with session 1, closing the previous handle
- sessionWatcher.process(sessionEvent(Watcher.Event.KeeperState.Expired));
- sessionWatcher.process(sessionEvent(Watcher.Event.KeeperState.Closed));
- assertThat(events).containsExactly(SessionEvent.SessionLost);
- sessionWatcher.process(sessionEvent(Watcher.Event.KeeperState.SyncConnected));
- assertThat(events).containsExactly(SessionEvent.SessionLost, SessionEvent.Reconnected,
- SessionEvent.SessionReestablished);
-
- // The listener marked session 1 replaced. A connection check of session 1 that still succeeds, and the
- // Closed event of its handle, don't reestablish it.
- sessionWatcher.process(sessionEvent(Watcher.Event.KeeperState.SyncConnected));
- sessionWatcher.process(sessionEvent(Watcher.Event.KeeperState.Closed));
- assertThat(events).containsExactly(SessionEvent.SessionLost, SessionEvent.Reconnected,
- SessionEvent.SessionReestablished);
-
- // session 2 connects
- sessionId.set(2);
- sessionWatcher.process(sessionEvent(Watcher.Event.KeeperState.SyncConnected));
- assertThat(events).containsExactly(SessionEvent.SessionLost, SessionEvent.Reconnected,
- SessionEvent.SessionReestablished, SessionEvent.Reconnected, SessionEvent.SessionReestablished);
- }
-
- private static WatchedEvent sessionEvent(Watcher.Event.KeeperState state) {
- return new WatchedEvent(Watcher.Event.EventType.None, state, null);
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ completeExistsWith(zk, KeeperException.Code.OK);
+
+ try (ZKSessionWatcher watcher = new ZKSessionWatcher(zk, events::add, false)) {
+ watcher.setSessionInvalid();
+ sessionId.set(2);
+ handleGeneration.set(2);
+ watcher.checkConnectionStatus();
+ assertThat(events).containsExactly(SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+
+ watcher.process(new ZKSessionWatcher.SessionEventSnapshot(
+ sessionEvent(KeeperState.Expired), 1, 1, States.CLOSED));
+ assertThat(events).containsExactly(SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+ }
+ }
+
+ @Test
+ public void testProbeListenerRunsOutsideCurrentHandleUpdate() throws Exception {
+ PulsarZooKeeperClient zk = mock(PulsarZooKeeperClient.class);
+ PulsarZooKeeperClient.CurrentHandleIdentity handleIdentity =
+ new PulsarZooKeeperClient.CurrentHandleIdentity(1, 1);
+ AtomicBoolean updatingCurrentHandle = new AtomicBoolean();
+ AtomicBoolean listenerCalledDuringUpdate = new AtomicBoolean();
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ when(zk.getSessionTimeout()).thenReturn(30_000);
+ when(zk.getSessionId()).thenReturn(1L);
+ when(zk.getCurrentHandleGeneration()).thenReturn(1L);
+ when(zk.getCurrentHandleIdentity()).thenReturn(handleIdentity);
+ when(zk.runStateUpdateIfCurrentHandle(any(), any())).thenAnswer(invocation -> {
+ Runnable stateUpdate = invocation.getArgument(1);
+ updatingCurrentHandle.set(true);
+ try {
+ stateUpdate.run();
+ } finally {
+ updatingCurrentHandle.set(false);
+ }
+ return true;
+ });
+ when(zk.getState()).thenReturn(States.CONNECTING);
+ completeExistsWith(zk, KeeperException.Code.CONNECTIONLOSS);
+
+ try (ZKSessionWatcher watcher = new ZKSessionWatcher(zk, event -> {
+ listenerCalledDuringUpdate.set(updatingCurrentHandle.get());
+ events.add(event);
+ }, false)) {
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(SessionEvent.ConnectionLost);
+ assertThat(listenerCalledDuringUpdate).isFalse();
+ }
+ }
+
+ @Test
+ public void testProbeFromOlderHandleIsIgnoredWhenSessionIdsAreUnassigned() throws Exception {
+ PulsarZooKeeperClient zk = mock(PulsarZooKeeperClient.class);
+ AtomicLong handleGeneration = new AtomicLong(1);
+ when(zk.getSessionTimeout()).thenReturn(600_000);
+ when(zk.getSessionId()).thenReturn(0L);
+ when(zk.getCurrentHandleGeneration()).thenAnswer(invocation -> handleGeneration.get());
+ when(zk.getCurrentHandleIdentity()).thenAnswer(invocation ->
+ new PulsarZooKeeperClient.CurrentHandleIdentity(handleGeneration.get(), 0));
+ when(zk.runStateUpdateIfCurrentHandle(any(), any())).thenAnswer(invocation -> {
+ PulsarZooKeeperClient.CurrentHandleIdentity expectedIdentity = invocation.getArgument(0);
+ if (!expectedIdentity.equals(zk.getCurrentHandleIdentity())) {
+ return false;
+ }
+ Runnable action = invocation.getArgument(1);
+ action.run();
+ return true;
+ });
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ doAnswer(invocation -> {
+ handleGeneration.set(2);
+ StatCallback callback = invocation.getArgument(2);
+ callback.processResult(KeeperException.Code.OK.intValue(), "/", null, null);
+ return null;
+ }).when(zk).exists(eq("/"), eq(false), any(StatCallback.class), isNull());
+
+ try (ZKSessionWatcher watcher = new ZKSessionWatcher(zk, events::add, false)) {
+ watcher.setSessionInvalid();
+ watcher.checkConnectionStatus();
+
+ assertThat(events).isEmpty();
+ }
+ }
+
+ @Test
+ public void testQueuedSyncConnectedFromReplacedSessionIsIgnored() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ PulsarZooKeeperClient zk = mock(PulsarZooKeeperClient.class);
+ AtomicLong sessionId = new AtomicLong(1);
+ AtomicLong handleGeneration = new AtomicLong(1);
+ when(zk.getSessionTimeout()).thenReturn(600_000);
+ when(zk.getSessionId()).thenAnswer(invocation -> sessionId.get());
+ when(zk.getCurrentHandleGeneration()).thenAnswer(invocation -> handleGeneration.get());
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ when(zk.getClientConfig()).thenReturn(new ZKClientConfig());
+ doAnswer(invocation -> {
+ VoidCallback callback = invocation.getArgument(3);
+ callback.processResult(KeeperException.Code.CONNECTIONLOSS.intValue(), "/", null);
+ return null;
+ }).when(zk).addWatch(eq("/"), any(), eq(AddWatchMode.PERSISTENT_RECURSIVE), any(), isNull());
+
+ CountDownLatch eventExecutorBlocked = new CountDownLatch(1);
+ CountDownLatch releaseEventExecutor = new CountDownLatch(1);
+ try (ZKMetadataStore store = new ZKMetadataStore(zk, false)) {
+ store.registerSessionListener(events::add);
+ store.processSessionWatcher(sessionEvent(KeeperState.Expired));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .untilAsserted(() -> assertThat(events).containsExactly(SessionEvent.SessionLost));
+
+ // The persistent watch recreation fails, so session 1 is marked as replaced.
+ store.processSessionWatcher(new PulsarZooKeeperClient.HandleStateWatchedEvent(
+ sessionEvent(KeeperState.SyncConnected), 1, 1, States.CONNECTED));
+ Awaitility.await().atMost(5, TimeUnit.SECONDS)
+ .untilAsserted(() -> assertThat(events).containsExactly(
+ SessionEvent.SessionLost, SessionEvent.Reconnected));
+
+ store.getEventExecutor().execute(() -> {
+ eventExecutorBlocked.countDown();
+ try {
+ releaseEventExecutor.await(5, TimeUnit.SECONDS);
+ } catch (InterruptedException e) {
+ Thread.currentThread().interrupt();
+ }
+ });
+ assertThat(eventExecutorBlocked.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Queue another SyncConnected from session 1, then replace it with session 2 before processing.
+ store.processSessionWatcher(new PulsarZooKeeperClient.HandleStateWatchedEvent(
+ sessionEvent(KeeperState.SyncConnected), 1, 1, States.CONNECTED));
+ CountDownLatch queuedEventProcessed = new CountDownLatch(1);
+ store.getEventExecutor().execute(queuedEventProcessed::countDown);
+ sessionId.set(2);
+ handleGeneration.set(2);
+ releaseEventExecutor.countDown();
+ assertThat(queuedEventProcessed.await(5, TimeUnit.SECONDS)).isTrue();
+
+ // Drain notifications that the queued session event might have produced.
+ CountDownLatch notificationsProcessed = new CountDownLatch(1);
+ store.getEventExecutor().execute(notificationsProcessed::countDown);
+ assertThat(notificationsProcessed.await(5, TimeUnit.SECONDS)).isTrue();
+
+ assertThat(events).containsExactly(SessionEvent.SessionLost, SessionEvent.Reconnected);
+ } finally {
+ releaseEventExecutor.countDown();
+ }
+ }
+
+ @Test
+ public void testConnectedReadOnlyProbeShouldNotBeTreatedAsWritableConnection() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.CONNECTEDREADONLY);
+ completeExistsWith(zk, KeeperException.Code.OK);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.checkConnectionStatus();
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(SessionEvent.ConnectionLost);
+ }
+ }
+
+ @Test
+ public void testConnectedReadOnlyThenWritableProbeDoesNotReestablishSession() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.CONNECTEDREADONLY);
+ completeExistsWith(zk, KeeperException.Code.OK);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.checkConnectionStatus();
+
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(SessionEvent.ConnectionLost, SessionEvent.Reconnected);
+ }
+ }
+
+ @Test
+ public void testReadOnlyProbesWaitForWritableConnectionWithoutSessionLoss() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.CONNECTEDREADONLY);
+ completeExistsWith(zk, KeeperException.Code.OK);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.checkConnectionStatus();
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(SessionEvent.ConnectionLost);
+
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ watcher.checkConnectionStatus();
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null));
+
+ assertThat(events).containsExactly(SessionEvent.ConnectionLost, SessionEvent.Reconnected);
+ }
+ }
+
+ @Test
+ public void testClosedHandleAfterExpirationWaitsForWritableConnectionAndDeduplicatesRecovery() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.CLOSED);
+ completeExistsWith(zk, KeeperException.Code.OK);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Expired, null));
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Closed, null));
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.Closed, null));
+ watcher.checkConnectionStatus();
+ assertThat(events).containsExactly(SessionEvent.SessionLost);
+
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ watcher.checkConnectionStatus();
+ watcher.process(new WatchedEvent(EventType.None, KeeperState.SyncConnected, null));
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(
+ SessionEvent.SessionLost, SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+ }
+ }
+
+ @Test
+ public void testWritableConnectedProbeReestablishesSession() throws Exception {
+ List<SessionEvent> events = new CopyOnWriteArrayList<>();
+ ZooKeeper zk = newSessionZooKeeper();
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ completeExistsWith(zk, KeeperException.Code.OK);
+
+ try (ZKSessionWatcher watcher = newSessionWatcher(zk, events)) {
+ watcher.setSessionInvalid();
+ watcher.checkConnectionStatus();
+
+ assertThat(events).containsExactly(SessionEvent.Reconnected, SessionEvent.SessionReestablished);
+ }
+ }
+
+ private static ZKSessionWatcher newSessionWatcher(List<SessionEvent> events) {
+ return newSessionWatcher(newSessionZooKeeper(), events);
+ }
+
+ private static ZKSessionWatcher newSessionWatcher(ZooKeeper zk, List<SessionEvent> events) {
+ return new ZKSessionWatcher(zk, events::add, false);
+ }
+
+ private static ZooKeeper newSessionZooKeeper() {
+ ZooKeeper zk = mock(ZooKeeper.class);
+ when(zk.getSessionTimeout()).thenReturn(30_000);
+ when(zk.getSessionId()).thenReturn(0x1234L);
+ when(zk.getState()).thenReturn(States.CONNECTED);
+ when(zk.getClientConfig()).thenReturn(new ZKClientConfig());
+ return zk;
+ }
+
+ private static void completeExistsWith(ZooKeeper zk, KeeperException.Code code) {
+ doAnswer(invocation -> {
+ StatCallback callback = invocation.getArgument(2);
+ callback.processResult(code.intValue(), "/", null, null);
+ return null;
+ }).when(zk).exists(eq("/"), eq(false), any(StatCallback.class), isNull());
+ }
+
+ private static WatchedEvent sessionEvent(KeeperState state) {
+ return new WatchedEvent(EventType.None, state, null);
}
}