Search in sources :

Example 1 with PropertyEventShutdownCallback

use of com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback in project rest.li by linkedin.

the class ZooKeeperEphemeralStoreTest method testPutGetRemovePartial.

@Test(groups = { "small", "back-end" })
public void testPutGetRemovePartial() throws InterruptedException, IOException, PropertyStoreException, ExecutionException {
    ZooKeeperEphemeralStore<String> store = getStore();
    store.put("service-1", "1");
    store.put("service-1", "2");
    store.put("service-2", "3");
    assertTrue(store.get("service-1").equals("1,2") || store.get("service-1").equals("2,1"));
    assertEquals(store.get("service-2"), "3");
    assertNull(store.get("service-3"));
    store.removePartial("service-1", "2");
    assertEquals(store.get("service-1"), "1");
    store.remove("service-2");
    assertNull(store.get("service-2"));
    final CountDownLatch latch = new CountDownLatch(1);
    store.shutdown(new PropertyEventShutdownCallback() {

        @Override
        public void done() {
            latch.countDown();
        }
    });
    if (!latch.await(5, TimeUnit.SECONDS)) {
        fail("unable to shut down");
    }
}
Also used : PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) CountDownLatch(java.util.concurrent.CountDownLatch) Test(org.testng.annotations.Test)

Example 2 with PropertyEventShutdownCallback

use of com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback in project rest.li by linkedin.

the class TogglingStoreTest method testShutdownDisabled.

@Test(groups = { "small", "back-end" })
public void testShutdownDisabled() throws InterruptedException {
    TogglingStore<String> store = getStore();
    store.setEnabled(false);
    final CountDownLatch latch = new CountDownLatch(1);
    store.shutdown(new PropertyEventShutdownCallback() {

        @Override
        public void done() {
            latch.countDown();
        }
    });
    if (!latch.await(5, TimeUnit.SECONDS)) {
        fail("unable to shut down store");
    }
}
Also used : PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) CountDownLatch(java.util.concurrent.CountDownLatch) PropertyStoreTest(com.linkedin.d2.discovery.stores.PropertyStoreTest) Test(org.testng.annotations.Test)

Example 3 with PropertyEventShutdownCallback

use of com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback in project rest.li by linkedin.

the class SimpleLoadBalancerState method shutdown.

@Override
public void shutdown(final PropertyEventShutdownCallback shutdown) {
    trace(_log, "shutdown");
    // shutdown all three registries, all tracker clients, and the event thread
    _executor.execute(new PropertyEvent("shutdown load balancer state") {

        @Override
        public void innerRun() {
            // Need to shutdown loadBalancerStrategies before the transportClients are shutdown
            for (Map<String, LoadBalancerStrategy> strategyEntry : _serviceStrategies.values()) {
                strategyEntry.values().forEach(LoadBalancerStrategy::shutdown);
            }
            // put all tracker clients into a single set for convenience
            Set<TransportClient> transportClients = new HashSet<>();
            for (Map<String, TransportClient> clientsByScheme : _serviceClients.values()) {
                transportClients.addAll(clientsByScheme.values());
            }
            Callback<None> trackerCallback = Callbacks.countDown(Callbacks.<None>adaptSimple(new SimpleCallback() {

                @Override
                public void onDone() {
                    shutdown.done();
                }
            }), transportClients.size());
            info(_log, "shutting down cluster clients");
            for (TransportClient transportClient : transportClients) {
                transportClient.shutdown(trackerCallback);
            }
            // so it is needed to notify all the listeners
            for (SimpleLoadBalancerStateListener listener : _listeners) {
                // Notify the strategy removal
                for (Map.Entry<String, Map<String, LoadBalancerStrategy>> serviceStrategy : _serviceStrategies.entrySet()) {
                    for (Map.Entry<String, LoadBalancerStrategy> strategyEntry : serviceStrategy.getValue().entrySet()) {
                        listener.onStrategyRemoved(serviceStrategy.getKey(), strategyEntry.getKey(), strategyEntry.getValue());
                    }
                    // Also notify the client removal
                    Map<URI, TrackerClient> trackerClients = _trackerClients.get(serviceStrategy.getKey());
                    if (trackerClients != null) {
                        for (TrackerClient client : trackerClients.values()) {
                            listener.onClientRemoved(serviceStrategy.getKey(), client);
                        }
                    }
                }
            }
            // When SimpleLoadBalancerStateis shutdown, all the cluster listener also need to be notified.
            for (LoadBalancerClusterListener clusterListener : _clusterListeners) {
                for (String clusterName : _clusterInfo.keySet()) {
                    clusterListener.onClusterRemoved(clusterName);
                }
            }
        }
    });
}
Also used : Set(java.util.Set) HashSet(java.util.HashSet) TransportClient(com.linkedin.r2.transport.common.bridge.client.TransportClient) LoadBalancerClusterListener(com.linkedin.d2.balancer.LoadBalancerClusterListener) PropertyEvent(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEvent) RelativeLoadBalancerStrategy(com.linkedin.d2.balancer.strategies.relative.RelativeLoadBalancerStrategy) LoadBalancerStrategy(com.linkedin.d2.balancer.strategies.LoadBalancerStrategy) PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) Callback(com.linkedin.common.callback.Callback) SimpleCallback(com.linkedin.common.callback.SimpleCallback) TrackerClient(com.linkedin.d2.balancer.clients.TrackerClient) Map(java.util.Map) ConcurrentHashMap(java.util.concurrent.ConcurrentHashMap) HashMap(java.util.HashMap) None(com.linkedin.common.util.None) SimpleCallback(com.linkedin.common.callback.SimpleCallback)

Example 4 with PropertyEventShutdownCallback

use of com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback in project rest.li by linkedin.

the class LoadBalancerSimulator method shutdown.

public void shutdown() throws Exception {
    _clockedExecutor.shutdown();
    final CountDownLatch latch = new CountDownLatch(1);
    PropertyEventShutdownCallback callback = () -> latch.countDown();
    _loadBalancer.shutdown(callback);
    if (!latch.await(60, TimeUnit.SECONDS)) {
        Assert.fail("unable to shutdown state");
    }
    _log.info("LoadBalancer Shutdown @ {}", _clockedExecutor.currentTimeMillis());
}
Also used : PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) CountDownLatch(java.util.concurrent.CountDownLatch)

Example 5 with PropertyEventShutdownCallback

use of com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback in project rest.li by linkedin.

the class SimpleLoadBalancerTest method testClusterInfoProviderRegisterClusterListener.

/**
 * The Register cluster Listener code is already tested in SimpleLoadBalancerStateTest, this is here for testing the
 * SimpleLoadBalancer API exposing this.
 */
@Test
public void testClusterInfoProviderRegisterClusterListener() throws InterruptedException, ExecutionException, ServiceUnavailableException {
    MockStore<ServiceProperties> serviceRegistry = new MockStore<>();
    MockStore<ClusterProperties> clusterRegistry = new MockStore<>();
    MockStore<UriProperties> uriRegistry = new MockStore<>();
    SimpleLoadBalancer loadBalancer = setupLoadBalancer(null, serviceRegistry, clusterRegistry, uriRegistry);
    FutureCallback<None> balancerCallback = new FutureCallback<>();
    loadBalancer.start(balancerCallback);
    balancerCallback.get();
    MockClusterListener testClusterListener = new MockClusterListener();
    loadBalancer.registerClusterListener(testClusterListener);
    loadBalancer.listenToCluster(CLUSTER1_NAME, false, new LoadBalancerState.NullStateListenerCallback());
    clusterRegistry.put(CLUSTER1_NAME, new ClusterProperties(CLUSTER1_NAME, Collections.emptyList(), Collections.emptyMap(), Collections.emptySet(), NullPartitionProperties.getInstance(), Collections.emptyList(), new HashMap<>(), false));
    Assert.assertEquals(testClusterListener.getClusterAddedCount(CLUSTER1_NAME), 1, "expected add count of 1");
    Assert.assertEquals(testClusterListener.getClusterRemovedCount(CLUSTER1_NAME), 0, "expected remove count of 0");
    final CountDownLatch latch = new CountDownLatch(1);
    PropertyEventShutdownCallback callback = latch::countDown;
    loadBalancer.shutdown(callback);
    if (!latch.await(60, TimeUnit.SECONDS)) {
        fail("unable to shutdown state");
    }
    Assert.assertEquals(testClusterListener.getClusterAddedCount(CLUSTER1_NAME), 1, "expected add count of 1 after shutdown");
    Assert.assertEquals(testClusterListener.getClusterRemovedCount(CLUSTER1_NAME), 1, "expected remove count of 1 after shutdown");
}
Also used : PropertyEventShutdownCallback(com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback) HashMap(java.util.HashMap) MockStore(com.linkedin.d2.discovery.stores.mock.MockStore) CountDownLatch(java.util.concurrent.CountDownLatch) LoadBalancerState(com.linkedin.d2.balancer.LoadBalancerState) ServiceProperties(com.linkedin.d2.balancer.properties.ServiceProperties) UriProperties(com.linkedin.d2.balancer.properties.UriProperties) ClusterProperties(com.linkedin.d2.balancer.properties.ClusterProperties) None(com.linkedin.common.util.None) FutureCallback(com.linkedin.common.callback.FutureCallback) Test(org.testng.annotations.Test)

Aggregations

PropertyEventShutdownCallback (com.linkedin.d2.discovery.event.PropertyEventThread.PropertyEventShutdownCallback)17 CountDownLatch (java.util.concurrent.CountDownLatch)15 Test (org.testng.annotations.Test)8 HashMap (java.util.HashMap)6 None (com.linkedin.common.util.None)5 ClusterProperties (com.linkedin.d2.balancer.properties.ClusterProperties)5 ServiceProperties (com.linkedin.d2.balancer.properties.ServiceProperties)5 UriProperties (com.linkedin.d2.balancer.properties.UriProperties)5 HashSet (java.util.HashSet)5 Map (java.util.Map)5 FutureCallback (com.linkedin.common.callback.FutureCallback)4 LoadBalancerStrategy (com.linkedin.d2.balancer.strategies.LoadBalancerStrategy)4 URIRequest (com.linkedin.d2.balancer.util.URIRequest)4 MockStore (com.linkedin.d2.discovery.stores.mock.MockStore)4 RequestContext (com.linkedin.r2.message.RequestContext)4 TransportClientFactory (com.linkedin.r2.transport.common.TransportClientFactory)4 URI (java.net.URI)4 DarkClusterConfigMap (com.linkedin.d2.DarkClusterConfigMap)3 ServiceUnavailableException (com.linkedin.d2.balancer.ServiceUnavailableException)3 RewriteLoadBalancerClient (com.linkedin.d2.balancer.clients.RewriteLoadBalancerClient)3