@Bean(destroyMethod = "close")
@ConditionalOnMissingBean
@SneakyThrows
public CuratorFramework curatorFramework(
RetryPolicy retryPolicy, ZookeeperDiscoveryProperties properties) {
CuratorFrameworkFactory.Builder builder = CuratorFrameworkFactory.builder();
if (ensembleProvider != null) {
builder.ensembleProvider(ensembleProvider);
}
CuratorFramework curator =
builder
.retryPolicy(retryPolicy)
.connectString(zookeeperProperties.getConnectString())
.build();
curator.start();
log.trace(
"blocking until connected to zookeeper for "
+ properties.getBlockUntilConnectedWait()
+ properties.getBlockUntilConnectedUnit());
curator.blockUntilConnected(
properties.getBlockUntilConnectedWait(), properties.getBlockUntilConnectedUnit());
log.trace("connected to zookeeper");
return curator;
}
void doEvaluate(Statement base) throws Throwable {
try {
cluster = new TestingCluster(3);
cluster.start();
client = newClient(cluster.getConnectString(), new RetryOneTime(200 /* ms */));
client.start();
checkState(
client.blockUntilConnected(5, TimeUnit.SECONDS),
"failed to connect to zookeeper in 5 seconds");
base.evaluate();
} catch (InterruptedException e) {
Thread.currentThread().interrupt();
throw new IllegalStateException("Interrupted while connecting to ZooKeeper", e);
} finally {
client.close();
cluster.close();
}
}
@BeforeMethod
public void initStuff() throws Exception {
LOG.info("Creating ZK server instance listening in port {}...", ZK_PORT);
while (zkServer == null) {
try {
zkServer = new TestingServer(ZK_PORT);
} catch (BindException e) {
System.err.println("Getting bind exception - retrying to allocate server");
zkServer = null;
}
}
LOG.info("ZK Server Started @ {}", zkServer.getConnectString());
RetryPolicy retryPolicy = new ExponentialBackoffRetry(1000, 3);
LOG.info("Creating Zookeeper Client connected to {}", ZK_CLUSTER);
zkClient =
CuratorFrameworkFactory.builder()
.namespace("omid")
.connectString(ZK_CLUSTER)
.retryPolicy(retryPolicy)
.build();
zkClient.start();
zkClient.blockUntilConnected();
LOG.info("Creating Internal Zookeeper Client connected to {}", ZK_CLUSTER);
storageInternalZKClient =
Mockito.spy(
CuratorFrameworkFactory.builder()
.namespace("omid")
.connectString(ZK_CLUSTER)
.retryPolicy(retryPolicy)
.build());
storage = new ZKTimestampStorage(storageInternalZKClient);
}
@Before
public void setUp() throws Exception {
req = EasyMock.createStrictMock(HttpServletRequest.class);
supervisorManager = EasyMock.createMock(SupervisorManager.class);
taskLockbox = EasyMock.createStrictMock(TaskLockbox.class);
taskLockbox.syncFromStorage();
EasyMock.expectLastCall().atLeastOnce();
taskLockbox.add(EasyMock.<Task>anyObject());
EasyMock.expectLastCall().atLeastOnce();
taskLockbox.remove(EasyMock.<Task>anyObject());
EasyMock.expectLastCall().atLeastOnce();
// for second Noop Task directly added to deep storage.
taskLockbox.add(EasyMock.<Task>anyObject());
EasyMock.expectLastCall().atLeastOnce();
taskLockbox.remove(EasyMock.<Task>anyObject());
EasyMock.expectLastCall().atLeastOnce();
taskActionClientFactory = EasyMock.createStrictMock(TaskActionClientFactory.class);
EasyMock.expect(taskActionClientFactory.create(EasyMock.<Task>anyObject()))
.andReturn(null)
.anyTimes();
EasyMock.replay(taskLockbox, taskActionClientFactory);
taskStorage = new HeapMemoryTaskStorage(new TaskStorageConfig(null));
runTaskCountDownLatches = new CountDownLatch[2];
runTaskCountDownLatches[0] = new CountDownLatch(1);
runTaskCountDownLatches[1] = new CountDownLatch(1);
taskCompletionCountDownLatches = new CountDownLatch[2];
taskCompletionCountDownLatches[0] = new CountDownLatch(1);
taskCompletionCountDownLatches[1] = new CountDownLatch(1);
announcementLatch = new CountDownLatch(1);
IndexerZkConfig indexerZkConfig =
new IndexerZkConfig(new ZkPathsConfig(), null, null, null, null, null);
setupServerAndCurator();
curator.start();
curator.blockUntilConnected();
curator.create().creatingParentsIfNeeded().forPath(indexerZkConfig.getLeaderLatchPath());
druidNode = new DruidNode("hey", "what", 1234);
ServiceEmitter serviceEmitter = new NoopServiceEmitter();
taskMaster =
new TaskMaster(
new TaskQueueConfig(null, new Period(1), null, new Period(10)),
taskLockbox,
taskStorage,
taskActionClientFactory,
druidNode,
indexerZkConfig,
new TaskRunnerFactory<MockTaskRunner>() {
@Override
public MockTaskRunner build() {
return new MockTaskRunner(runTaskCountDownLatches, taskCompletionCountDownLatches);
}
},
curator,
new NoopServiceAnnouncer() {
@Override
public void announce(DruidNode node) {
announcementLatch.countDown();
}
},
serviceEmitter,
supervisorManager);
EmittingLogger.registerEmitter(serviceEmitter);
}