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); } } } }
@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); }