From a9d798a532fcfff8f703db1455fb4d94c5b16e33 Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Thu, 10 Apr 2025 17:47:41 +0300 Subject: [PATCH 1/2] [improve][test] Use configured session timeout for MockZooKeeper and TestZKServer in PulsarTestContext --- .../broker/testcontext/PulsarTestContext.java | 142 ++++++++++++------ .../MockZooKeeperMetadataStoreProvider.java | 1 + .../org/apache/zookeeper/MockZooKeeper.java | 8 +- .../zookeeper/MockZooKeeperSession.java | 12 +- 4 files changed, 116 insertions(+), 47 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java index f8205a2e939a3..26acae1f820ef 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java @@ -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. * @@ -255,6 +267,7 @@ public static class Builder { protected boolean configOverrideCalled = false; protected Function brokerServiceCustomizer = Function.identity(); protected PulsarTestContext otherContextToClose; + protected WithMockZooKeeperOrTestZKServer withMockZooKeeperOrTestZKServer; /** * Initialize the ServiceConfiguration with default values. @@ -475,31 +488,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 @@ -518,27 +514,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. @@ -626,6 +609,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); } @@ -650,6 +634,57 @@ public final PulsarTestContext build() { return super.build(); } + void createWithMockZooKeeperOrTestZKServerInstances() { + if (withMockZooKeeperOrTestZKServer == null) { + return; + } + int sessionTimeout = (int) super.config.getMetadataStoreSessionTimeoutMillis(); + try { + if (withMockZooKeeperOrTestZKServer.isMockZooKeeper()) { + mockZooKeeper(createMockZooKeeper(sessionTimeout)); + if (withMockZooKeeperOrTestZKServer + == WithMockZooKeeperOrTestZKServer.MOCKZOOKEEPER_SEPARATE_GLOBAL) { + mockZooKeeperGlobal(createMockZooKeeper(sessionTimeout)); + } + } else if (withMockZooKeeperOrTestZKServer.isTestZKServer()) { + testZKServer(createTestZookeeper(sessionTimeout)); + if (withMockZooKeeperOrTestZKServer + == WithMockZooKeeperOrTestZKServer.TEST_ZK_SERVER_SEPARATE_GLOBAL) { + testZKServerGlobal(createTestZookeeper(sessionTimeout)); + } + } + } 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 -> { @@ -714,28 +749,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); } @@ -765,16 +802,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); @@ -786,9 +837,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); diff --git a/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/MockZooKeeperMetadataStoreProvider.java b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/MockZooKeeperMetadataStoreProvider.java index 994a97c2b1053..e507ae4544557 100644 --- a/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/MockZooKeeperMetadataStoreProvider.java +++ b/pulsar-metadata/src/test/java/org/apache/pulsar/metadata/MockZooKeeperMetadataStoreProvider.java @@ -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; } diff --git a/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeper.java b/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeper.java index e124699ee1383..2682f038df29a 100644 --- a/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeper.java +++ b/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeper.java @@ -137,6 +137,7 @@ public List getChildren() { private ThreadLocal inExecutorThreadLocal; private int referenceCount; private List 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 @@ -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; } @@ -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 { diff --git a/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeperSession.java b/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeperSession.java index c812423b7280d..766f70979aa0e 100644 --- a/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeperSession.java +++ b/testmocks/src/main/java/org/apache/zookeeper/MockZooKeeperSession.java @@ -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); } @@ -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 From bf0ba417fcdc0fd0815edae08f7ac7e1f4d08e1e Mon Sep 17 00:00:00 2001 From: Lari Hotari Date: Fri, 11 Apr 2025 10:04:54 +0300 Subject: [PATCH 2/2] Fix reusing MockZooKeeper / TestZKServer after broker restart in tests --- .../broker/testcontext/PulsarTestContext.java | 26 ++++++++++++++++--- 1 file changed, 22 insertions(+), 4 deletions(-) diff --git a/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java b/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java index 26acae1f820ef..be204ff3cdd30 100644 --- a/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java +++ b/pulsar-broker/src/test/java/org/apache/pulsar/broker/testcontext/PulsarTestContext.java @@ -424,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()); @@ -641,16 +643,32 @@ void createWithMockZooKeeperOrTestZKServerInstances() { int sessionTimeout = (int) super.config.getMetadataStoreSessionTimeoutMillis(); try { if (withMockZooKeeperOrTestZKServer.isMockZooKeeper()) { - mockZooKeeper(createMockZooKeeper(sessionTimeout)); + if (super.mockZooKeeper == null) { + mockZooKeeper(createMockZooKeeper(sessionTimeout)); + } else { + log.warn("Skipping creating mockZooKeeper, already set"); + } if (withMockZooKeeperOrTestZKServer == WithMockZooKeeperOrTestZKServer.MOCKZOOKEEPER_SEPARATE_GLOBAL) { - mockZooKeeperGlobal(createMockZooKeeper(sessionTimeout)); + if (super.mockZooKeeperGlobal == null) { + mockZooKeeperGlobal(createMockZooKeeper(sessionTimeout)); + } else { + log.warn("Skipping creating mockZooKeeperGlobal, already set"); + } } } else if (withMockZooKeeperOrTestZKServer.isTestZKServer()) { - testZKServer(createTestZookeeper(sessionTimeout)); + if (super.testZKServer == null) { + testZKServer(createTestZookeeper(sessionTimeout)); + } else { + log.warn("Skipping creating testZKServer, already set"); + } if (withMockZooKeeperOrTestZKServer == WithMockZooKeeperOrTestZKServer.TEST_ZK_SERVER_SEPARATE_GLOBAL) { - testZKServerGlobal(createTestZookeeper(sessionTimeout)); + if (super.testZKServerGlobal == null) { + testZKServerGlobal(createTestZookeeper(sessionTimeout)); + } else { + log.warn("Skipping creating testZKServerGlobal, already set"); + } } } } catch (Exception e) {