private void serverAddedSegment(final DruidServerMetadata server, final DataSegment segment) {
String segmentId = segment.getIdentifier();
synchronized (lock) {
log.debug("Adding segment[%s] for server[%s]", segment, server);
ServerSelector selector = selectors.get(segmentId);
if (selector == null) {
selector = new ServerSelector(segment, tierSelectorStrategy);
VersionedIntervalTimeline<String, ServerSelector> timeline =
timelines.get(segment.getDataSource());
if (timeline == null) {
timeline = new VersionedIntervalTimeline<>(Ordering.natural());
timelines.put(segment.getDataSource(), timeline);
}
timeline.add(
segment.getInterval(),
segment.getVersion(),
segment.getShardSpec().createChunk(selector));
selectors.put(segmentId, selector);
}
QueryableDruidServer queryableDruidServer = clients.get(server.getName());
if (queryableDruidServer == null) {
queryableDruidServer = addServer(baseView.getInventoryValue(server.getName()));
}
selector.addServerAndUpdateSegment(queryableDruidServer, segment);
}
}
public void clear() {
synchronized (lock) {
final Iterator<String> clientsIter = clients.keySet().iterator();
while (clientsIter.hasNext()) {
clientsIter.remove();
}
timelines.clear();
final Iterator<ServerSelector> selectorsIter = selectors.values().iterator();
while (selectorsIter.hasNext()) {
final ServerSelector selector = selectorsIter.next();
selectorsIter.remove();
while (!selector.isEmpty()) {
final QueryableDruidServer pick = selector.pick();
selector.removeServer(pick);
}
}
}
}
private void serverRemovedSegment(DruidServerMetadata server, DataSegment segment) {
String segmentId = segment.getIdentifier();
final ServerSelector selector;
synchronized (lock) {
log.debug("Removing segment[%s] from server[%s].", segmentId, server);
selector = selectors.get(segmentId);
if (selector == null) {
log.warn("Told to remove non-existant segment[%s]", segmentId);
return;
}
QueryableDruidServer queryableDruidServer = clients.get(server.getName());
if (!selector.removeServer(queryableDruidServer)) {
log.warn(
"Asked to disassociate non-existant association between server[%s] and segment[%s]",
server, segmentId);
}
if (selector.isEmpty()) {
VersionedIntervalTimeline<String, ServerSelector> timeline =
timelines.get(segment.getDataSource());
selectors.remove(segmentId);
final PartitionChunk<ServerSelector> removedPartition =
timeline.remove(
segment.getInterval(),
segment.getVersion(),
segment.getShardSpec().createChunk(selector));
if (removedPartition == null) {
log.warn(
"Asked to remove timeline entry[interval: %s, version: %s] that doesn't exist",
segment.getInterval(), segment.getVersion());
}
}
}
}
@Test
public void testRun() throws Exception {
HttpClient httpClient = EasyMock.createMock(HttpClient.class);
final URL url = new URL("http://foo/druid/v2/");
SettableFuture<InputStream> futureResult = SettableFuture.create();
Capture<Request> capturedRequest = EasyMock.newCapture();
EasyMock.expect(
httpClient.go(
EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject()))
.andReturn(futureResult)
.times(1);
SettableFuture futureException = SettableFuture.create();
EasyMock.expect(
httpClient.go(
EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject()))
.andReturn(futureException)
.times(1);
EasyMock.expect(
httpClient.go(
EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject()))
.andReturn(SettableFuture.create())
.atLeastOnce();
EasyMock.replay(httpClient);
final ServerSelector serverSelector =
new ServerSelector(
new DataSegment(
"test",
new Interval("2013-01-01/2013-01-02"),
new DateTime("2013-01-01").toString(),
Maps.<String, Object>newHashMap(),
Lists.<String>newArrayList(),
Lists.<String>newArrayList(),
NoneShardSpec.instance(),
0,
0L),
new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy()));
DirectDruidClient client1 =
new DirectDruidClient(
new ReflectionQueryToolChestWarehouse(),
QueryRunnerTestHelper.NOOP_QUERYWATCHER,
new DefaultObjectMapper(),
httpClient,
"foo",
new NoopServiceEmitter());
DirectDruidClient client2 =
new DirectDruidClient(
new ReflectionQueryToolChestWarehouse(),
QueryRunnerTestHelper.NOOP_QUERYWATCHER,
new DefaultObjectMapper(),
httpClient,
"foo2",
new NoopServiceEmitter());
QueryableDruidServer queryableDruidServer1 =
new QueryableDruidServer(
new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0),
client1);
serverSelector.addServerAndUpdateSegment(queryableDruidServer1, serverSelector.getSegment());
QueryableDruidServer queryableDruidServer2 =
new QueryableDruidServer(
new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0),
client2);
serverSelector.addServerAndUpdateSegment(queryableDruidServer2, serverSelector.getSegment());
TimeBoundaryQuery query = Druids.newTimeBoundaryQueryBuilder().dataSource("test").build();
HashMap<String, List> context = Maps.newHashMap();
Sequence s1 = client1.run(query, context);
Assert.assertTrue(capturedRequest.hasCaptured());
Assert.assertEquals(url, capturedRequest.getValue().getUrl());
Assert.assertEquals(HttpMethod.POST, capturedRequest.getValue().getMethod());
Assert.assertEquals(1, client1.getNumOpenConnections());
// simulate read timeout
Sequence s2 = client1.run(query, context);
Assert.assertEquals(2, client1.getNumOpenConnections());
futureException.setException(new ReadTimeoutException());
Assert.assertEquals(1, client1.getNumOpenConnections());
// subsequent connections should work
Sequence s3 = client1.run(query, context);
Sequence s4 = client1.run(query, context);
Sequence s5 = client1.run(query, context);
Assert.assertTrue(client1.getNumOpenConnections() == 4);
// produce result for first connection
futureResult.set(
new ByteArrayInputStream(
"[{\"timestamp\":\"2014-01-01T01:02:03Z\", \"result\": 42.0}]".getBytes()));
List<Result> results = Sequences.toList(s1, Lists.<Result>newArrayList());
Assert.assertEquals(1, results.size());
Assert.assertEquals(new DateTime("2014-01-01T01:02:03Z"), results.get(0).getTimestamp());
Assert.assertEquals(3, client1.getNumOpenConnections());
client2.run(query, context);
client2.run(query, context);
Assert.assertTrue(client2.getNumOpenConnections() == 2);
Assert.assertTrue(serverSelector.pick() == queryableDruidServer2);
EasyMock.verify(httpClient);
}
@Test
public void testQueryInterruptionExceptionLogMessage() throws JsonProcessingException {
HttpClient httpClient = EasyMock.createMock(HttpClient.class);
SettableFuture<Object> interruptionFuture = SettableFuture.create();
Capture<Request> capturedRequest = EasyMock.newCapture();
String hostName = "localhost:8080";
EasyMock.expect(
httpClient.go(
EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject()))
.andReturn(interruptionFuture)
.anyTimes();
EasyMock.replay(httpClient);
DataSegment dataSegment =
new DataSegment(
"test",
new Interval("2013-01-01/2013-01-02"),
new DateTime("2013-01-01").toString(),
Maps.<String, Object>newHashMap(),
Lists.<String>newArrayList(),
Lists.<String>newArrayList(),
NoneShardSpec.instance(),
0,
0L);
final ServerSelector serverSelector =
new ServerSelector(
dataSegment,
new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy()));
DirectDruidClient client1 =
new DirectDruidClient(
new ReflectionQueryToolChestWarehouse(),
QueryRunnerTestHelper.NOOP_QUERYWATCHER,
new DefaultObjectMapper(),
httpClient,
hostName,
new NoopServiceEmitter());
QueryableDruidServer queryableDruidServer =
new QueryableDruidServer(
new DruidServer("test1", hostName, 0, "historical", DruidServer.DEFAULT_TIER, 0),
client1);
serverSelector.addServerAndUpdateSegment(queryableDruidServer, dataSegment);
TimeBoundaryQuery query = Druids.newTimeBoundaryQueryBuilder().dataSource("test").build();
HashMap<String, List> context = Maps.newHashMap();
interruptionFuture.set(new ByteArrayInputStream("{\"error\":\"testing\"}".getBytes()));
Sequence results = client1.run(query, context);
QueryInterruptedException actualException = null;
try {
Sequences.toList(results, Lists.newArrayList());
} catch (QueryInterruptedException e) {
actualException = e;
}
Assert.assertNotNull(actualException);
Assert.assertEquals(actualException.getMessage(), QueryInterruptedException.UNKNOWN_EXCEPTION);
Assert.assertEquals(actualException.getCauseMessage(), "testing");
Assert.assertEquals(actualException.getHost(), hostName);
EasyMock.verify(httpClient);
}
@Test
public void testCancel() throws Exception {
HttpClient httpClient = EasyMock.createStrictMock(HttpClient.class);
Capture<Request> capturedRequest = EasyMock.newCapture();
ListenableFuture<Object> cancelledFuture = Futures.immediateCancelledFuture();
SettableFuture<Object> cancellationFuture = SettableFuture.create();
EasyMock.expect(
httpClient.go(
EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject()))
.andReturn(cancelledFuture)
.once();
EasyMock.expect(
httpClient.go(
EasyMock.capture(capturedRequest), EasyMock.<HttpResponseHandler>anyObject()))
.andReturn(cancellationFuture)
.once();
EasyMock.replay(httpClient);
final ServerSelector serverSelector =
new ServerSelector(
new DataSegment(
"test",
new Interval("2013-01-01/2013-01-02"),
new DateTime("2013-01-01").toString(),
Maps.<String, Object>newHashMap(),
Lists.<String>newArrayList(),
Lists.<String>newArrayList(),
NoneShardSpec.instance(),
0,
0L),
new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy()));
DirectDruidClient client1 =
new DirectDruidClient(
new ReflectionQueryToolChestWarehouse(),
QueryRunnerTestHelper.NOOP_QUERYWATCHER,
new DefaultObjectMapper(),
httpClient,
"foo",
new NoopServiceEmitter());
QueryableDruidServer queryableDruidServer1 =
new QueryableDruidServer(
new DruidServer("test1", "localhost", 0, "historical", DruidServer.DEFAULT_TIER, 0),
client1);
serverSelector.addServerAndUpdateSegment(queryableDruidServer1, serverSelector.getSegment());
TimeBoundaryQuery query = Druids.newTimeBoundaryQueryBuilder().dataSource("test").build();
HashMap<String, List> context = Maps.newHashMap();
cancellationFuture.set(
new StatusResponseHolder(HttpResponseStatus.OK, new StringBuilder("cancelled")));
Sequence results = client1.run(query, context);
Assert.assertEquals(HttpMethod.DELETE, capturedRequest.getValue().getMethod());
Assert.assertEquals(0, client1.getNumOpenConnections());
QueryInterruptedException exception = null;
try {
Sequences.toList(results, Lists.newArrayList());
} catch (QueryInterruptedException e) {
exception = e;
}
Assert.assertNotNull(exception);
EasyMock.verify(httpClient);
}