Skip to content

[improve][test] Use configured session timeout for MockZooKeeper and TestZKServer in PulsarTestContext #24171

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
lhotari merged 2 commits into from
Apr 11, 2025
Merged

[improve][test] Use configured session timeout for MockZooKeeper and TestZKServer in PulsarTestContext #24171

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
a9d798a
[improve][test] Use configured session timeout for MockZooKeeper and …
lhotari Apr 10, 2025
bf0ba41
Fix reusing MockZooKeeper / TestZKServer after broker restart in tests
lhotari Apr 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
160 changes: 115 additions & 45 deletions pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -239,6 +239,18 @@ public ServerCnx createServerCnxSpy() {
getPulsarService());
}

private enum WithMockZooKeeperOrTestZKServer {
MOCKZOOKEEPER, MOCKZOOKEEPER_SEPARATE_GLOBAL, TEST_ZK_SERVER, TEST_ZK_SERVER_SEPARATE_GLOBAL;

boolean isMockZooKeeper() {
return this == MOCKZOOKEEPER || this == MOCKZOOKEEPER_SEPARATE_GLOBAL;
}

boolean isTestZKServer() {
return this == TEST_ZK_SERVER || this == TEST_ZK_SERVER_SEPARATE_GLOBAL;
}
}

/**
* A builder for a PulsarTestContext.
*
Expand All @@ -255,6 +267,7 @@ public static class Builder {
protected boolean configOverrideCalled = false;
protected Function<BrokerService, BrokerService> brokerServiceCustomizer = Function.identity();
protected PulsarTestContext otherContextToClose;
protected WithMockZooKeeperOrTestZKServer withMockZooKeeperOrTestZKServer;

/**
* Initialize the ServiceConfiguration with default values.
Expand Down Expand Up @@ -411,11 +424,13 @@ public Builder pulsarServiceCustomizer(
public Builder reuseMockBookkeeperAndMetadataStores(PulsarTestContext otherContext) {
bookKeeperClient(otherContext.getBookKeeperClient());
if (otherContext.getMockZooKeeper() != null) {
withMockZooKeeperOrTestZKServer = null;
mockZooKeeper(otherContext.getMockZooKeeper());
if (otherContext.getMockZooKeeperGlobal() != null) {
mockZooKeeperGlobal(otherContext.getMockZooKeeperGlobal());
}
} else if (otherContext.getTestZKServer() != null) {
withMockZooKeeperOrTestZKServer = null;
testZKServer(otherContext.getTestZKServer());
if (otherContext.getTestZKServerGlobal() != null) {
testZKServerGlobal(otherContext.getTestZKServerGlobal());
Expand Down Expand Up @@ -475,31 +490,14 @@ public Builder withMockZookeeper() {
* @return the builder
*/
public Builder withMockZookeeper(boolean useSeparateGlobalZk) {
try {
mockZooKeeper(createMockZooKeeper());
if (useSeparateGlobalZk) {
mockZooKeeperGlobal(createMockZooKeeper());
}
} catch (Exception e) {
throw new RuntimeException(e);
if (useSeparateGlobalZk) {
withMockZooKeeperOrTestZKServer = WithMockZooKeeperOrTestZKServer.MOCKZOOKEEPER_SEPARATE_GLOBAL;
} else {
withMockZooKeeperOrTestZKServer = WithMockZooKeeperOrTestZKServer.MOCKZOOKEEPER;
}
return this;
}

private MockZooKeeper createMockZooKeeper() throws Exception {
MockZooKeeper zk = MockZooKeeper.newInstance();
initializeZookeeper(zk);
registerCloseable(zk::shutdown);
return zk;
}

private static void initializeZookeeper(ZooKeeper zk) throws KeeperException, InterruptedException {
ZkUtils.createFullPathOptimistic(zk, "/ledgers/available/192.168.1.1:" + 5000,
"".getBytes(StandardCharsets.UTF_8), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);

zk.create("/ledgers/LAYOUT", "1\nflat:1".getBytes(StandardCharsets.UTF_8), ZooDefs.Ids.OPEN_ACL_UNSAFE,
CreateMode.PERSISTENT);
}

/**
* Configure this PulsarTestContext to use a test ZooKeeper instance which is
Expand All @@ -518,27 +516,14 @@ public Builder withTestZookeeper() {
* @return the builder
*/
public Builder withTestZookeeper(boolean useSeparateGlobalZk) {
try {
testZKServer(createTestZookeeper());
if (useSeparateGlobalZk) {
testZKServerGlobal(createTestZookeeper());
}
} catch (Exception e) {
throw new RuntimeException(e);
if (useSeparateGlobalZk) {
withMockZooKeeperOrTestZKServer = WithMockZooKeeperOrTestZKServer.TEST_ZK_SERVER_SEPARATE_GLOBAL;
} else {
withMockZooKeeperOrTestZKServer = WithMockZooKeeperOrTestZKServer.TEST_ZK_SERVER;
}
return this;
}

private TestZKServer createTestZookeeper() throws Exception {
TestZKServer testZKServer = new TestZKServer();
try (ZooKeeper zkc = new ZooKeeper(testZKServer.getConnectionString(), 5000, event -> {
})) {
initializeZookeeper(zkc);
}
registerCloseable(testZKServer);
return testZKServer;
}

/**
* Applicable only when PulsarTestContext is not startable. This will configure mocks
* for PulsarTestResources and related classes.
Expand Down Expand Up @@ -626,6 +611,7 @@ public final PulsarTestContext build() {
if (configOverrideCustomizer != null) {
configOverrideCustomizer.accept(super.config);
}
createWithMockZooKeeperOrTestZKServerInstances();
if (super.managedLedgerStorage != null && !MockUtil.isMock(super.managedLedgerStorage)) {
super.managedLedgerStorage = spyConfig.getManagedLedgerStorage().spy(super.managedLedgerStorage);
}
Expand All @@ -650,6 +636,73 @@ public final PulsarTestContext build() {
return super.build();
}

void createWithMockZooKeeperOrTestZKServerInstances() {
if (withMockZooKeeperOrTestZKServer == null) {
return;
}
int sessionTimeout = (int) super.config.getMetadataStoreSessionTimeoutMillis();
try {
if (withMockZooKeeperOrTestZKServer.isMockZooKeeper()) {
if (super.mockZooKeeper == null) {
mockZooKeeper(createMockZooKeeper(sessionTimeout));
} else {
log.warn("Skipping creating mockZooKeeper, already set");
}
if (withMockZooKeeperOrTestZKServer
== WithMockZooKeeperOrTestZKServer.MOCKZOOKEEPER_SEPARATE_GLOBAL) {
if (super.mockZooKeeperGlobal == null) {
mockZooKeeperGlobal(createMockZooKeeper(sessionTimeout));
} else {
log.warn("Skipping creating mockZooKeeperGlobal, already set");
}
}
} else if (withMockZooKeeperOrTestZKServer.isTestZKServer()) {
if (super.testZKServer == null) {
testZKServer(createTestZookeeper(sessionTimeout));
} else {
log.warn("Skipping creating testZKServer, already set");
}
if (withMockZooKeeperOrTestZKServer
== WithMockZooKeeperOrTestZKServer.TEST_ZK_SERVER_SEPARATE_GLOBAL) {
if (super.testZKServerGlobal == null) {
testZKServerGlobal(createTestZookeeper(sessionTimeout));
} else {
log.warn("Skipping creating testZKServerGlobal, already set");
}
}
}
} catch (Exception e) {
throw new RuntimeException(e);
}
}

private MockZooKeeper createMockZooKeeper(int sessionTimeout) throws Exception {
MockZooKeeper zk = MockZooKeeper.newInstance();
zk.setSessionTimeout(sessionTimeout);
initializeZookeeper(zk);
registerCloseable(zk::shutdown);
return zk;
}

// this might not be required at all, but it's kept here as an example
private static void initializeZookeeper(ZooKeeper zk) throws KeeperException, InterruptedException {
ZkUtils.createFullPathOptimistic(zk, "/ledgers/available/192.168.1.1:" + 5000,
"".getBytes(StandardCharsets.UTF_8), ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT);

zk.create("/ledgers/LAYOUT", "1\nflat:1".getBytes(StandardCharsets.UTF_8), ZooDefs.Ids.OPEN_ACL_UNSAFE,
CreateMode.PERSISTENT);
}

private TestZKServer createTestZookeeper(int sessionTimeout) throws Exception {
TestZKServer testZKServer = new TestZKServer();
try (ZooKeeper zkc = new ZooKeeper(testZKServer.getConnectionString(), sessionTimeout, event -> {
})) {
initializeZookeeper(zkc);
}
registerCloseable(testZKServer);
return testZKServer;
}

protected void handlePreallocatePorts(ServiceConfiguration config) {
if (super.preallocatePorts) {
config.getBrokerServicePort().ifPresent(portNumber -> {
Expand Down Expand Up @@ -714,28 +767,30 @@ private void initializeCommonPulsarServices(SpyConfig spyConfig) {
if (super.localMetadataStore == null || super.configurationMetadataStore == null) {
if (super.mockZooKeeper != null) {
MetadataStoreExtended mockZookeeperMetadataStore =
createMockZookeeperMetadataStore(super.mockZooKeeper, MetadataStoreConfig.METADATA_STORE);
createMockZookeeperMetadataStore(super.mockZooKeeper, super.config,
MetadataStoreConfig.METADATA_STORE);
if (super.localMetadataStore == null) {
localMetadataStore(mockZookeeperMetadataStore);
}
if (super.configurationMetadataStore == null) {
if (super.mockZooKeeperGlobal != null) {
configurationMetadataStore(createMockZookeeperMetadataStore(super.mockZooKeeperGlobal,
MetadataStoreConfig.CONFIGURATION_METADATA_STORE));
super.config, MetadataStoreConfig.CONFIGURATION_METADATA_STORE));
} else {
configurationMetadataStore(mockZookeeperMetadataStore);
}
}
} else if (super.testZKServer != null) {
MetadataStoreExtended testZookeeperMetadataStore =
createTestZookeeperMetadataStore(super.testZKServer, MetadataStoreConfig.METADATA_STORE);
createTestZookeeperMetadataStore(super.testZKServer, super.config,
MetadataStoreConfig.METADATA_STORE);
if (super.localMetadataStore == null) {
localMetadataStore(testZookeeperMetadataStore);
}
if (super.configurationMetadataStore == null) {
if (super.testZKServerGlobal != null) {
configurationMetadataStore(createTestZookeeperMetadataStore(super.testZKServerGlobal,
MetadataStoreConfig.CONFIGURATION_METADATA_STORE));
super.config, MetadataStoreConfig.CONFIGURATION_METADATA_STORE));
} else {
configurationMetadataStore(testZookeeperMetadataStore);
}
Expand Down Expand Up @@ -765,16 +820,30 @@ private void initializeCommonPulsarServices(SpyConfig spyConfig) {
}
}

private MetadataStoreConfig createMetadataStoreConfig(ServiceConfiguration config, String metadataStoreName) {
return MetadataStoreConfig.builder()
.sessionTimeoutMillis((int) config.getMetadataStoreSessionTimeoutMillis())
.allowReadOnlyOperations(config.isMetadataStoreAllowReadOnlyOperations())
.batchingEnabled(config.isMetadataStoreBatchingEnabled())
.batchingMaxDelayMillis(config.getMetadataStoreBatchingMaxDelayMillis())
.batchingMaxOperations(config.getMetadataStoreBatchingMaxOperations())
.batchingMaxSizeKb(config.getMetadataStoreBatchingMaxSizeKb())
.metadataStoreName(metadataStoreName)
.build();
}

private MetadataStoreExtended createMockZookeeperMetadataStore(MockZooKeeper mockZooKeeper,
ServiceConfiguration config,
String metadataStoreName) {
// provide a unique session id for each instance
MockZooKeeperSession mockZooKeeperSession = MockZooKeeperSession.newInstance(mockZooKeeper, false);
mockZooKeeperSession.setSessionTimeout((int) config.getMetadataStoreSessionTimeoutMillis());
registerCloseable(() -> {
mockZooKeeperSession.close();
resetSpyOrMock(mockZooKeeperSession);
});
ZKMetadataStore zkMetadataStore = new ZKMetadataStore(mockZooKeeperSession,
MetadataStoreConfig.builder().metadataStoreName(metadataStoreName).build());
ZKMetadataStore zkMetadataStore =
new ZKMetadataStore(mockZooKeeperSession, createMetadataStoreConfig(config, metadataStoreName));
registerCloseable(() -> {
zkMetadataStore.close();
resetSpyOrMock(zkMetadataStore);
Expand All @@ -786,9 +855,10 @@ private MetadataStoreExtended createMockZookeeperMetadataStore(MockZooKeeper moc

@SneakyThrows
private MetadataStoreExtended createTestZookeeperMetadataStore(TestZKServer zkServer,
ServiceConfiguration config,
String metadataStoreName) {
MetadataStoreExtended store = MetadataStoreExtended.create("zk:" + zkServer.getConnectionString(),
MetadataStoreConfig.builder().metadataStoreName(metadataStoreName).build());
createMetadataStoreConfig(config, metadataStoreName));
registerCloseable(store);
MetadataStoreExtended nonClosingProxy =
NonClosingProxyHandler.createNonClosingProxy(store, MetadataStoreExtended.class);
Expand Down
1 change: 1 addition & 0 deletions ...metadata/src/test/java/org/apache/pulsar/metadata/MockZooKeeperMetadataStoreProvider.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,7 @@ public MetadataStore create(String metadataURL, MetadataStoreConfig metadataStor
MockZooKeeper mockZooKeeper = mockZooKeepers.computeIfAbsent(metadataURL,
k -> MockZooKeeper.newInstance().registerCloseable(() -> mockZooKeepers.remove(k)));
MockZooKeeperSession mockZooKeeperSession = MockZooKeeperSession.newInstance(mockZooKeeper, true);
mockZooKeeperSession.setSessionTimeout(metadataStoreConfig.getSessionTimeoutMillis());
ZKMetadataStore zkMetadataStore = new ZKMetadataStore(mockZooKeeperSession, metadataStoreConfig, true);
return zkMetadataStore;
}
Expand Down
8 changes: 7 additions & 1 deletion testmocks/src/main/java/org/apache/zookeeper/MockZooKeeper.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,6 +137,7 @@ public List<String> getChildren() {
private ThreadLocal<Boolean> inExecutorThreadLocal;
private int referenceCount;
private List<AutoCloseable> closeables;
private int sessionTimeout;

//see details of Objenesis caching - http://objenesis.org/details.html
//see supported jvms - https://github.com/easymock/objenesis/blob/master/SupportedJVMs.md
Expand Down Expand Up @@ -188,6 +189,7 @@ private static MockZooKeeper createMockZooKeeperInstance(int readOpDelayMs) {
zk.readOpDelayMs = readOpDelayMs;
zk.sequentialIdGenerator = new AtomicLong();
zk.closeables = new ArrayList<>();
zk.sessionTimeout = 30_000;
return zk;
}

Expand All @@ -204,7 +206,11 @@ private void init() {

@Override
public int getSessionTimeout() {
return 30_000;
return sessionTimeout;
}

public void setSessionTimeout(int sessionTimeout) {
this.sessionTimeout = sessionTimeout;
}

private MockZooKeeper(String quorum) throws Exception {
Expand Down
12 changes: 11 additions & 1 deletion testmocks/src/main/java/org/apache/zookeeper/MockZooKeeperSession.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -48,6 +48,8 @@ public class MockZooKeeperSession extends ZooKeeper {

private boolean closeMockZooKeeperOnClose;

private int sessionTimeout = -1;

public static MockZooKeeperSession newInstance(MockZooKeeper mockZooKeeper) {
return newInstance(mockZooKeeper, true);
}
Expand All @@ -74,7 +76,15 @@ private MockZooKeeperSession(String quorum) throws Exception {

@Override
public int getSessionTimeout() {
return mockZooKeeper.getSessionTimeout();
if (sessionTimeout > 0) {
return sessionTimeout;
} else {
return mockZooKeeper.getSessionTimeout();
}
}

public void setSessionTimeout(int sessionTimeout) {
this.sessionTimeout = sessionTimeout;
}

@Override
Expand Down