package com.bazaarvoice.ostrich.pool; import com.bazaarvoice.ostrich.HealthCheckResults; import com.bazaarvoice.ostrich.HostDiscovery; import com.bazaarvoice.ostrich.LoadBalanceAlgorithm; import com.bazaarvoice.ostrich.PartitionContext; import com.bazaarvoice.ostrich.RetryPolicy; import com.bazaarvoice.ostrich.ServiceCallback; import com.bazaarvoice.ostrich.ServiceEndPoint; import com.bazaarvoice.ostrich.ServiceFactory; import com.bazaarvoice.ostrich.ServicePoolStatistics; import com.bazaarvoice.ostrich.exceptions.MaxRetriesException; import com.bazaarvoice.ostrich.exceptions.NoAvailableHostsException; import com.bazaarvoice.ostrich.exceptions.NoSuitableHostsException; import com.bazaarvoice.ostrich.exceptions.OnlyBadHostsException; import com.bazaarvoice.ostrich.exceptions.ServiceException; import com.bazaarvoice.ostrich.healthcheck.FixedHealthCheckRetryDelay; import com.bazaarvoice.ostrich.partition.PartitionFilter; import com.codahale.metrics.MetricRegistry; import com.google.common.base.Ticker; import com.google.common.collect.ImmutableList; import com.google.common.collect.Iterables; import com.google.common.collect.Sets; import org.junit.After; import org.junit.Before; import org.junit.Test; import org.mockito.ArgumentCaptor; import org.mockito.Matchers; import org.mockito.invocation.InvocationOnMock; import org.mockito.stubbing.Answer; import java.io.IOException; import java.util.Collections; import java.util.Set; import java.util.concurrent.CountDownLatch; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; import java.util.concurrent.atomic.AtomicBoolean; import static org.junit.Assert.assertEquals; import static org.junit.Assert.assertFalse; import static org.junit.Assert.assertSame; import static org.junit.Assert.assertTrue; import static org.junit.Assert.fail; import static org.mockito.Matchers.any; import static org.mockito.Matchers.anyInt; import static org.mockito.Matchers.anyLong; import static org.mockito.Matchers.eq; import static org.mockito.Matchers.same; import static org.mockito.Mockito.mock; import static org.mockito.Mockito.never; import static org.mockito.Mockito.reset; import static org.mockito.Mockito.times; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.verifyZeroInteractions; import static org.mockito.Mockito.when; public abstract class AbstractServicePoolTestingHarness { protected static final ServiceEndPoint FOO_ENDPOINT = mock(ServiceEndPoint.class); protected static final ServiceEndPoint BAR_ENDPOINT = mock(ServiceEndPoint.class); protected static final ServiceEndPoint BAZ_ENDPOINT = mock(ServiceEndPoint.class); protected static final Service FOO_SERVICE = mock(Service.class); protected static final Service BAR_SERVICE = mock(Service.class); protected static final Service BAZ_SERVICE = mock(Service.class); protected static final RetryPolicy NEVER_RETRY = mock(RetryPolicy.class); protected final ServiceCachingPolicy UNLIMITED_CACHING = getServiceCachingPolicy(); protected Ticker _ticker; protected HostDiscovery _hostDiscovery; protected PartitionFilter _partitionFilter; protected LoadBalanceAlgorithm _loadBalanceAlgorithm; protected ServiceFactory<Service> _serviceFactory; protected ScheduledExecutorService _healthCheckExecutor; protected MetricRegistry _registry; protected ServicePool<Service> _pool; protected abstract ServiceCachingPolicy getServiceCachingPolicy(); protected abstract ServiceFactory<Service> getServiceFactoryMock(); @SuppressWarnings("unchecked") @Before public void setup() { // // This setup method takes the approach of building a reasonably useful ServicePool using mocks that can then be // customized by individual test methods to add whatever functionality they need to (or ignored completely). // _ticker = mock(Ticker.class); _hostDiscovery = mock(HostDiscovery.class); when(_hostDiscovery.getHosts()).thenReturn(ImmutableList.of(FOO_ENDPOINT, BAR_ENDPOINT, BAZ_ENDPOINT)); _partitionFilter = mock(PartitionFilter.class); when(_partitionFilter.filter(any(Iterable.class), any(PartitionContext.class))) .thenAnswer(new Answer<Object>() { @Override public Iterable<ServiceEndPoint> answer(InvocationOnMock invocation) throws Throwable { return (Iterable<ServiceEndPoint>) invocation.getArguments()[0]; } }); _loadBalanceAlgorithm = mock(LoadBalanceAlgorithm.class); when(_loadBalanceAlgorithm.choose(any(Iterable.class), any(ServicePoolStatistics.class))) .thenAnswer(new Answer<ServiceEndPoint>() { @Override public ServiceEndPoint answer(InvocationOnMock invocation) throws Throwable { // Always choose the first end point. This is probably fine since most tests will have just a single // end point available anyways. Iterable<ServiceEndPoint> endPoints = (Iterable<ServiceEndPoint>) invocation.getArguments()[0]; return endPoints.iterator().next(); } }); _serviceFactory = getServiceFactoryMock(); when(_serviceFactory.getServiceName()).thenReturn(Service.class.getSimpleName()); when(_serviceFactory.create(FOO_ENDPOINT)).thenReturn(FOO_SERVICE); when(_serviceFactory.create(BAR_ENDPOINT)).thenReturn(BAR_SERVICE); when(_serviceFactory.create(BAZ_ENDPOINT)).thenReturn(BAZ_SERVICE); when(_serviceFactory.isRetriableException(any(Exception.class))).thenReturn(true); _healthCheckExecutor = mock(ScheduledExecutorService.class); _registry = new MetricRegistry(); _pool = new ServicePool<>(_ticker, _hostDiscovery, false, _serviceFactory, UNLIMITED_CACHING, _partitionFilter, _loadBalanceAlgorithm, _healthCheckExecutor, true, FixedHealthCheckRetryDelay.ZERO, _registry); } @After public void teardown() throws IOException { _pool.close(); _hostDiscovery.close(); } @Test public void testCallInvokedWithCorrectService() { Service expectedService = mock(Service.class); // Wire our expected service into the system when(_serviceFactory.create(FOO_ENDPOINT)).thenReturn(expectedService); // Don't leak service end points in real code!!! This is just a test case. Service actualService = _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Service>() { @Override public Service call(Service s) { return s; } }); assertSame(expectedService, actualService); } @Test(expected = NoAvailableHostsException.class) public void testThrowsNoAvailableHostsExceptionWhenNoEndPointsAvailable() { // Host discovery sees no end points... when(_hostDiscovery.getHosts()).thenReturn(ImmutableList.<ServiceEndPoint>of()); _pool.execute(NEVER_RETRY, null); } @Test(expected = OnlyBadHostsException.class) public void testThrowsOnlyBadHostsExceptionWhenOnlyBadEndPointsAvailable() { // Exhaust all of the end points... int numEndPointsAvailable = Iterables.size(_hostDiscovery.getHosts()); for (int i = 0; i < numEndPointsAvailable; i++) { try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); // should have propagated service exception } catch (MaxRetriesException e) { // Expected } } // This should trigger a service exception because there are no more available end points. _pool.execute(NEVER_RETRY, null); } @Test(expected = NoSuitableHostsException.class) public void testNullPartitionFilter() { reset(_partitionFilter); when(_partitionFilter.filter(Matchers.<Iterable<ServiceEndPoint>>any(), any(PartitionContext.class))) .thenReturn(null); boolean called = _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Boolean>() { @Override public Boolean call(Service service) throws ServiceException { return true; } }); assertFalse(called); } @Test(expected = NoSuitableHostsException.class) public void testEmptyPartition() { reset(_partitionFilter); when(_partitionFilter.filter(Matchers.<Iterable<ServiceEndPoint>>any(), any(PartitionContext.class))) .thenReturn(ImmutableList.<ServiceEndPoint>of()); boolean called = _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Boolean>() { @Override public Boolean call(Service service) throws ServiceException { return true; } }); assertFalse(called); } @SuppressWarnings ({"unchecked", "rawType"}) @Test public void testPartitionFilter() { reset(_partitionFilter); when(_partitionFilter.filter(Matchers.<Iterable<ServiceEndPoint>>any(), any(PartitionContext.class))) .thenReturn(ImmutableList.of(BAR_ENDPOINT)); PartitionContext context = mock(PartitionContext.class); _pool.execute(context, NEVER_RETRY, new ServiceCallback<Service, Boolean>() { @Override public Boolean call(Service service) throws ServiceException { return true; } }); // Verify the PartitionFilter was called with the correct arguments ArgumentCaptor<Iterable> filterEndPoints = ArgumentCaptor.forClass(Iterable.class); verify(_partitionFilter).filter(filterEndPoints.capture(), eq(context)); assertEquals(ImmutableList.of(FOO_ENDPOINT, BAR_ENDPOINT, BAZ_ENDPOINT), ImmutableList.copyOf(filterEndPoints.getValue())); // Verify that the result of the PartitionFilter was passed on as input to the LoadBalanceAlgorithm ArgumentCaptor<Iterable> balanceEndPoints = ArgumentCaptor.forClass(Iterable.class); verify(_loadBalanceAlgorithm).choose(balanceEndPoints.capture(), any(ServicePoolStatistics.class)); assertEquals(ImmutableList.of(BAR_ENDPOINT), ImmutableList.copyOf(balanceEndPoints.getValue())); } @Test(expected = NoSuitableHostsException.class) public void testThrowsNoSuitableHostsExceptionWhenLoadBalancerReturnsNull() { // Reset the load balance algorithm's setup and make it always return null. reset(_loadBalanceAlgorithm); when(_loadBalanceAlgorithm.choose(Matchers.<Iterable<ServiceEndPoint>>any(), any(ServicePoolStatistics.class))) .thenReturn(null); boolean called = _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Boolean>() { @Override public Boolean call(Service service) throws ServiceException { return true; } }); assertFalse(called); } @Test public void testDoesNotRetryOnCallbackSuccess() { RetryPolicy retry = mock(RetryPolicy.class); _pool.execute(retry, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) { return null; } }); // Should have never called the retry strategy for anything. This might change in the future if we implement // circuit breakers. verifyZeroInteractions(retry); } @Test public void testAttemptsToRetryOnRetriableException() { when(_serviceFactory.isRetriableException(any(Exception.class))).thenReturn(true); RetryPolicy retry = mock(RetryPolicy.class); when(retry.allowRetry(anyInt(), anyLong())).thenReturn(false); try { _pool.execute(retry, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) { throw new ServiceException(); } }); fail(); } catch (MaxRetriesException expected) { // We expect a service exception to happen since we're not going to be allowed to retry at all. // Make sure we asked the retry strategy if it was okay to retry one time (it said no). verify(retry).allowRetry(eq(1), anyLong()); } } @Test public void testDoesNotAttemptToRetryOnNonRetriableException() { when(_serviceFactory.isRetriableException(any(Exception.class))).thenReturn(false); RetryPolicy retry = mock(RetryPolicy.class); try { _pool.execute(retry, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new NullPointerException(); } }); fail(); } catch (NullPointerException expected) { verifyZeroInteractions(retry); } } @Test public void testKeepsRetryingUntilRetryPolicyReturnsFalse() { RetryPolicy retry = mock(RetryPolicy.class); when(retry.allowRetry(anyInt(), anyLong())).thenReturn(true, true, false); try { _pool.execute(retry, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); } catch (MaxRetriesException expected) { // Make sure we tried 3 times. verify(retry).allowRetry(eq(3), anyLong()); } } @Test public void testRetriesWithDifferentServiceEndPoints() { RetryPolicy retry = mock(RetryPolicy.class); when(retry.allowRetry(anyInt(), anyLong())).thenReturn(true, true, false); // Each end point has a specific service that it's supposed to return. Remember each service we've seen so // that we can make sure we saw the correct ones. final Set<Service> seenServices = Sets.newHashSet(); try { _pool.execute(retry, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { seenServices.add(service); throw new ServiceException(); } }); fail(); } catch (MaxRetriesException expected) { assertEquals(Sets.newHashSet(FOO_SERVICE, BAR_SERVICE, BAZ_SERVICE), seenServices); } } @Test public void testMaxRetriesExceptionIncludesUnderlyingCause() { final RuntimeException e = new RuntimeException(); try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw e; } }); fail(); } catch (MaxRetriesException expected) { assertSame(e, expected.getCause()); } } @Test public void testOnlyBadHostsExceptionIncludesUnderlyingCauseIfItMadeARequest() { // Exhaust all but one of the available end points... int numEndPointsAvailable = Iterables.size(_hostDiscovery.getHosts()); for (int i = 0; i < numEndPointsAvailable - 1; i++) { try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); // should have propagated service exception } catch (MaxRetriesException e) { // Expected } } // Executing a request only has one choice of endpoint now... RetryPolicy retry = mock(RetryPolicy.class); when(retry.allowRetry(anyInt(), anyLong())).thenReturn(true, true, false); final RuntimeException e = new RuntimeException(); try { _pool.execute(retry, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw e; } }); } catch (OnlyBadHostsException expected) { assertSame(e, expected.getCause()); } } @Test public void testHealthCheckIsRescheduled() { ArgumentCaptor<ServicePool.HealthCheckVerifier> healthCheckVerifierCaptor = ArgumentCaptor.forClass(ServicePool.HealthCheckVerifier.class); verify(_healthCheckExecutor).scheduleAtFixedRate(healthCheckVerifierCaptor.capture(), anyLong(), anyLong(), any(TimeUnit.class)); ServicePool.HealthCheckVerifier healthCheckVerifier = healthCheckVerifierCaptor.getValue(); when(_healthCheckExecutor.submit(any(ServicePool.HealthCheck.class))).thenThrow(NullPointerException.class); try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); } catch (MaxRetriesException expected) { // Expected exception } verify(_healthCheckExecutor, times(1)).submit(any(ServicePool.HealthCheck.class)); // Health check was not scheduled, so run the verify process to reschedule. // Scheduling fails, so next verify run will schedule again. healthCheckVerifier.run(); verify(_healthCheckExecutor, times(2)).submit(any(ServicePool.HealthCheck.class)); // Ensure next scheduling attempt is successful reset(_healthCheckExecutor); when(_healthCheckExecutor.submit(any(ServicePool.HealthCheck.class))).thenReturn(mock(Future.class)); // Verification successfully schedules the health check healthCheckVerifier.run(); verify(_healthCheckExecutor, times(1)).submit(any(ServicePool.HealthCheck.class)); // Running verification again does not schedule the health check again healthCheckVerifier.run(); verify(_healthCheckExecutor, times(1)).submit(any(ServicePool.HealthCheck.class)); } @Test public void testSubmitsHealthCheckOnRetriableException() { try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); } catch (MaxRetriesException expected) { // Expected exception } // Make sure we added a health check. verify(_healthCheckExecutor).submit(any(com.bazaarvoice.ostrich.pool.ServicePool.HealthCheck.class)); } @Test public void testDoesNotSubmitHealthCheckOnNonRetriableException() { when(_serviceFactory.isRetriableException(any(Exception.class))).thenReturn(false); try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new NullPointerException(); } }); fail(); } catch (NullPointerException expected) { // Expected exception } // Make sure we didn't add a health check. verify(_healthCheckExecutor, never()).submit(any(com.bazaarvoice.ostrich.pool.ServicePool.HealthCheck.class)); } @Test public void testStatsPassedToLoadBalancer() { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { return null; } }); verify(_loadBalanceAlgorithm).choose(Matchers.<Iterable<ServiceEndPoint>>any(), same(_pool.getServicePoolStatistics())); } @Test public void testStatsNumActiveInstancesIncrementsDuringExecute() { // Make sure we only get FOO_ENDPOINT. reset(_loadBalanceAlgorithm); when(_loadBalanceAlgorithm.choose(Matchers.<Iterable<ServiceEndPoint>>any(), any(ServicePoolStatistics.class))) .thenReturn(FOO_ENDPOINT); final ServicePoolStatistics servicePoolStatistics = _pool.getServicePoolStatistics(); int numActiveInitially = servicePoolStatistics.getNumActiveInstances(FOO_ENDPOINT); int numActiveDuringExecute = _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Integer>() { @Override public Integer call(Service service) throws ServiceException { return servicePoolStatistics.getNumActiveInstances(FOO_ENDPOINT); } }); assertEquals(numActiveInitially + 1, numActiveDuringExecute); } @Test public void testCheckForHealthyEndPointWhenEmpty() { when(_hostDiscovery.getHosts()).thenReturn(Collections.<ServiceEndPoint>emptySet()); assertTrue(Iterables.isEmpty(_pool.checkForHealthyEndPoint().getAllResults())); } @Test public void testCheckForHealthyEndPointWhenHealthy() { when(_serviceFactory.isHealthy(any(ServiceEndPoint.class))).thenReturn(true); HealthCheckResults results = _pool.checkForHealthyEndPoint(); assertTrue(results.hasHealthyResult()); assertTrue(Iterables.isEmpty(results.getUnhealthyResults())); } @Test public void testCheckForHealthyEndPointWhenUnhealthy() { when(_serviceFactory.isHealthy(any(ServiceEndPoint.class))).thenReturn(false); HealthCheckResults results = _pool.checkForHealthyEndPoint(); assertFalse(results.hasHealthyResult()); assertFalse(Iterables.isEmpty(results.getUnhealthyResults())); } @Test public void testCheckForHealthyEndPointRetriesWhenUnhealthy() { when(_serviceFactory.isHealthy(any(ServiceEndPoint.class))).thenReturn(false, true); HealthCheckResults results = _pool.checkForHealthyEndPoint(); assertTrue(results.hasHealthyResult()); assertFalse(Iterables.isEmpty(results.getUnhealthyResults())); } @Test public void testCheckForHealthyEndPointMarksEndPointBad() { when(_serviceFactory.isHealthy(any(ServiceEndPoint.class))).thenReturn(false); _pool.checkForHealthyEndPoint(); assertTrue(_pool.getBadEndPoints().containsAll(Sets.newHashSet(_pool.getAllEndPoints()))); } @Test public void testCheckForHealthyEndPointNotBeholdenToLoadBalancer() { reset(_loadBalanceAlgorithm); when(_loadBalanceAlgorithm.choose(Matchers.<Iterable<ServiceEndPoint>>any(), any(ServicePoolStatistics.class))) .thenReturn(null); assertFalse(Iterables.isEmpty(_pool.checkForHealthyEndPoint().getAllResults())); } @Test public void testCheckForHealthyEndPointRetriableException() { when(_serviceFactory.isRetriableException(any(Exception.class))).thenReturn(true); when(_serviceFactory.isHealthy(any(ServiceEndPoint.class))).thenThrow(new RuntimeException()).thenReturn(true); assertTrue(_pool.checkForHealthyEndPoint().hasHealthyResult()); } @Test public void testCheckForHealthyEndPointNonRetriableException() { when(_serviceFactory.isRetriableException(any(Exception.class))).thenReturn(false); when(_serviceFactory.isHealthy(any(ServiceEndPoint.class))).thenThrow(new RuntimeException()).thenReturn(true); assertFalse(_pool.checkForHealthyEndPoint().hasHealthyResult()); } @Test public void testCallsHealthCheckAfterRetriableException() throws InterruptedException { final AtomicBoolean healthCheckCalled = new AtomicBoolean(false); when(_serviceFactory.isHealthy(any(ServiceEndPoint.class))).then(new Answer<Boolean>() { @Override public Boolean answer(InvocationOnMock invocation) throws Throwable { healthCheckCalled.set(true); return false; } }); try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); } catch (MaxRetriesException expected) { // Expected exception } _pool.forceHealthChecks(); assertTrue(healthCheckCalled.get()); } @Test public void testDoesNotCallHealthCheckAfterNonRetriableException() throws InterruptedException { when(_serviceFactory.isRetriableException(any(Exception.class))).thenReturn(false); try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new NullPointerException(); } }); fail(); } catch (NullPointerException expected) { // Expected exception } // Make sure we never tried to call any health checks verify(_serviceFactory, never()).isHealthy(any(ServiceEndPoint.class)); } @Test public void testAllowsEndPointToBeUsedAgainAfterSuccessfulHealthCheck() { // Only allow BAZ to have a valid health check -- we know based on the load balance strategy that this // will be the last failed end point when(_serviceFactory.isHealthy(eq(BAZ_ENDPOINT))).thenReturn(true); // Exhaust all of the end points... for (ServiceEndPoint ignored : _hostDiscovery.getHosts()) { try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); // should have propagated service exception } catch (MaxRetriesException e) { // Expected } } _pool.forceHealthChecks(); // BAZ should still be healthy, so this shouldn't throw an exception. Service usedService = _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Service>() { @Override public Service call(Service service) throws ServiceException { return service; } }); assertSame(BAZ_SERVICE, usedService); } @Test public void testBatchHealthCheckAllowsEndPointToBeUsedAgainAfterSuccessfulHealthCheck() { // Exhaust all of the end points... int numEndPointsAvailable = Iterables.size(_hostDiscovery.getHosts()); for (int i = 0; i < numEndPointsAvailable; i++) { try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); // should have propagated service exception } catch (MaxRetriesException e) { // Expected } } // Set it up so that when we health check FOO, that it becomes healthy. when(_serviceFactory.isHealthy(FOO_ENDPOINT)).thenReturn(true); _pool.forceHealthChecks(); assertSame(FOO_SERVICE, _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Service>() { @Override public Service call(Service service) throws ServiceException { return service; } })); } @SuppressWarnings("unchecked") @Test public void testBadEndPointIsNoLongerHealthCheckedAfterHostDiscoveryRemovesIt() { // Redefine the end points that HostDiscovery knows about to be only FOO when(_hostDiscovery.getHosts()).thenReturn(ImmutableList.of(FOO_ENDPOINT)); // Make it so that FOO is considered bad... try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); // should have propagated service exception } catch (MaxRetriesException e) { // Expected } // At this point the health check for FOO would have been executed since it just failed. We're going to // reset the serviceFactory mock at this point so that we forget that that has happened. Once we reset it, // it's not going to be good for much, but the rest of this test fortunately doesn't interact with it. reset(_serviceFactory); // Capture the end point listener that was registered with HostDiscovery ArgumentCaptor<HostDiscovery.EndPointListener> listener = ArgumentCaptor.forClass( HostDiscovery.EndPointListener.class); verify(_hostDiscovery).addListener(listener.capture()); // Now, have HostDiscovery fire an event saying that FOO has been removed listener.getValue().onEndPointRemoved(FOO_ENDPOINT); _pool.forceHealthChecks(); verify(_serviceFactory, never()).isHealthy(eq(FOO_ENDPOINT)); } @Test public void testBadEndPointDisappearingFromHostDiscoveryDuringCallback() { // Redefine the end points that HostDiscovery knows about to be only FOO when(_hostDiscovery.getHosts()).thenReturn(ImmutableList.of(FOO_ENDPOINT)); // Capture the end point listener that was registered with HostDiscovery final ArgumentCaptor<HostDiscovery.EndPointListener> listener = ArgumentCaptor.forClass( HostDiscovery.EndPointListener.class); verify(_hostDiscovery).addListener(listener.capture()); try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { // Have HostDiscovery tell the ServicePool that FOO is gone. listener.getValue().onEndPointRemoved(FOO_ENDPOINT); // Now fail this request, this shouldn't result with any bad end points throw new ServiceException(); } }); fail(); // should have propagated service exception } catch (MaxRetriesException e) { // Expected } // At this point the bad end points list should be empty assertTrue(_pool.getBadEndPoints().isEmpty()); } @Test public void testIsHealthyHandlesExceptions() { when(_serviceFactory.isHealthy(FOO_ENDPOINT)).thenThrow(new RuntimeException()); // Even though an exception was thrown we shouldn't see it, instead false should be returned from checkHealth assertFalse(_pool.checkHealth(FOO_ENDPOINT).isHealthy()); } @Test public void testCloseMultipleTimes() { _pool.close(); _pool.close(); } @Test public void testDoesNotShutdownHealthCheckExecutorOnClose() { ServicePool<Service> pool = new ServicePool<>(_ticker, _hostDiscovery, false, _serviceFactory, ServiceCachingPolicyBuilder.NO_CACHING, _partitionFilter, _loadBalanceAlgorithm, _healthCheckExecutor, false, FixedHealthCheckRetryDelay.ZERO, _registry); pool.close(); verify(_healthCheckExecutor, never()).shutdown(); verify(_healthCheckExecutor, never()).shutdownNow(); } @Test public void testDoesShutdownHealthCheckExecutorOnClose() { ServicePool<Service> pool = new ServicePool<>(_ticker, _hostDiscovery, false, _serviceFactory, ServiceCachingPolicyBuilder.NO_CACHING, _partitionFilter, _loadBalanceAlgorithm, _healthCheckExecutor, true, FixedHealthCheckRetryDelay.ZERO, _registry); pool.close(); verify(_healthCheckExecutor, never()).shutdown(); verify(_healthCheckExecutor).shutdownNow(); } @Test public void testInterruptsHealthCheckOnClose() throws InterruptedException { // Make it so that when we health check FOO that we block until an interrupted exception occurs final CountDownLatch inHealthCheckLatch = new CountDownLatch(1); final CountDownLatch interruptedLatch = new CountDownLatch(1); when(_serviceFactory.isHealthy(FOO_ENDPOINT)).thenAnswer(new Answer<Boolean>() { @Override public Boolean answer(InvocationOnMock invocation) { inHealthCheckLatch.countDown(); synchronized (this) { try { this.wait(); } catch (InterruptedException e) { interruptedLatch.countDown(); } } return false; } }); // Redefine the end points that HostDiscovery knows about to be only FOO when(_hostDiscovery.getHosts()).thenReturn(ImmutableList.of(FOO_ENDPOINT)); ServicePool<Service> pool = new ServicePool<>(_ticker, _hostDiscovery, false, _serviceFactory, ServiceCachingPolicyBuilder.NO_CACHING, _partitionFilter, _loadBalanceAlgorithm, Executors.newScheduledThreadPool(1), true, new FixedHealthCheckRetryDelay(100, TimeUnit.MILLISECONDS), _registry); // Make it so that FOO needs to be health checked... try { pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); } catch (MaxRetriesException e) { // Expected } // The health check should be running now... assertTrue(inHealthCheckLatch.await(10, TimeUnit.SECONDS)); // And it should get interrupted on close... pool.close(); assertTrue(interruptedLatch.await(10, TimeUnit.SECONDS)); } @Test public void testValidEndPointCount() { assertEquals(3, _pool.getNumValidEndPoints()); assertEquals(0, _pool.getNumBadEndPoints()); } @Test public void testBadEndPointCount() { // Only allow BAZ to have a valid health check -- we know based on the load balance strategy that this // will be the last failed end point when(_serviceFactory.isHealthy(eq(BAZ_ENDPOINT))).thenReturn(true); // Exhaust all of the end points... int numEndPointsAvailable = Iterables.size(_hostDiscovery.getHosts()); for (int i = 0; i < numEndPointsAvailable; i++) { try { _pool.execute(NEVER_RETRY, new ServiceCallback<Service, Void>() { @Override public Void call(Service service) throws ServiceException { throw new ServiceException(); } }); fail(); // should have propagated service exception } catch (MaxRetriesException e) { // Expected } } _pool.forceHealthChecks(); // Only BAZ should be healthy assertEquals(1, _pool.getNumValidEndPoints()); assertEquals(2, _pool.getNumBadEndPoints()); } // A dummy interface for testing... protected static interface Service { } }