org.apache.curator.framework.recipes.shared.SharedCount Java Examples

@Test
public void givenRunningZookeeper_whenUsingSharedCounter_thenCounterIsIncrement() throws Exception {
    try (CuratorFramework client = newClient()) {
        client.start();

        try (SharedCount counter = new SharedCount(client, "/counters/A", 0)) {
            counter.start();

            counter.setCount(0);
            counter.setCount(counter.getCount() + 1);

            assertThat(counter.getCount()).isEqualTo(1);
        }

    }
}
 

/**
 * Increments the shared distributed counter named {@code counterName} by one.
 */
public void incrementCounter(String counterName) throws Exception {
    ArgumentChecker.notNull(counterName);
    
    SharedCount count = sharedCounters.get(counterName);
    Preconditions.checkArgument(count != null, "Invalid counter name: " + counterName + ". Shared counter may be down.");
    
    VersionedValue<Integer> currentCount = count.getVersionedValue();
    int newCount = currentCount.getValue() + 1;
    int tries = 0;
    while (!count.trySetCount(currentCount, newCount)) {
        currentCount = count.getVersionedValue();
        newCount = currentCount.getValue() + 1;
        if (++tries >= maxRetries) {
            // We've exceeded our max tries to update the counter. Try to re-register it and also throw an exception to
            // indicate that we didn't necessarily update the shared count.
            sharedCounters.remove(counterName);
            count.removeListener(sharedCountListeners.get(counterName));
            count.close();
            reregisterCounter(counterName, sharedCountListeners.get(counterName), newCount);
            throw new IllegalStateException("Unable to increment shared counter " + counterName + " after " + maxRetries
                            + " attempts. Zookeeper connection may be down.");
        }
    }
    localCounters.put(counterName, newCount);
}
 

private void reregisterCounter(String counterName, SharedCountListener listener, int seedValue) throws Exception {
    ArgumentChecker.notNull(counterName, listener);
    
    SharedCount count = new SharedCount(curatorClient, ZKPaths.makePath("/counters", counterName), seedValue);
    count.start();
    sharedCounters.put(counterName, count);
    sharedCountListeners.put(counterName, listener);
    localCounters.put(counterName, count.getCount());
    
    count.addListener(new SharedCountListener() {
        @Override
        public void countHasChanged(SharedCountReader sharedCountReader, int i) throws Exception {}
        
        @Override
        public void stateChanged(CuratorFramework curatorFramework, ConnectionState connectionState) {
            if (connectionState == ConnectionState.RECONNECTED) {
                try {
                    reregisterCounter(counterName, sharedCountListeners.get(counterName), localCounters.get(counterName));
                } catch (Exception e) {
                    System.err.println("Unable to re-register counter " + counterName + ": " + e.getMessage());
                }
            }
        }
    });
    count.addListener(listener);
}
 

/**
 * Registers a distributed shared counter named {@code counterName}. This counter can be watched on many servers, and can be used to coordinate local
 * in-memory global operations.
 * 
 * @param counterName
 *            the name of the counter
 * @param listener
 *            a listener that is called when the counter value changes
 */
public void registerCounter(String counterName, SharedCountListener listener) throws Exception {
    ArgumentChecker.notNull(counterName, listener);
    
    SharedCount count = new SharedCount(curatorClient, ZKPaths.makePath("/counters", counterName), 1);
    count.start();
    sharedCounters.put(counterName, count);
    localCounters.put(counterName, count.getCount());
    
    count.addListener(listener);
}
 

private void incrSharedCount(SharedCount sharedCount) throws Exception {
  while (true) {
    // Loop until we successfully increment the counter
    VersionedValue<Integer> versionedValue = sharedCount.getVersionedValue();
    if (sharedCount.trySetCount(versionedValue, versionedValue.getValue() + 1)) {
      break;
    }
  }
}
 

private void incrSharedCount(SharedCount sharedCount) throws Exception {
  while (true) {
    // Loop until we successfully increment the counter
    VersionedValue<Integer> versionedValue = sharedCount.getVersionedValue();
    if (sharedCount.trySetCount(versionedValue, versionedValue.getValue() + 1)) {
      break;
    }
  }
}
 

/**
 * Increments the shared distributed counter named {@code counterName} by one.
 */
public void decrementCounter(String counterName) throws Exception {
    ArgumentChecker.notNull(counterName);
    
    SharedCount count = sharedCounters.get(counterName);
    Preconditions.checkArgument(count != null, "Invalid counter name: " + counterName);
    
    int newCount = count.getCount() - 1;
    localCounters.put(counterName, newCount);
    while (!count.trySetCount(newCount)) {
        newCount = count.getCount() - 1;
    }
}
 

/**
 * Increments the shared distributed counter named {@code counterName} by one.
 */
public void incrementCounter(String counterName) throws Exception {
    ArgumentChecker.notNull(counterName);
    
    SharedCount count = sharedCounters.get(counterName);
    Preconditions.checkArgument(count != null, "Invalid counter name: " + counterName);
    
    int newCount = count.getCount() + 1;
    localCounters.put(counterName, newCount);
    while (!count.trySetCount(newCount)) {
        newCount = count.getCount() + 1;
    }
}
 

private void shutdownCounters() {
    for (String counterName : localCounters.keySet()) {
        SharedCount count = sharedCounters.remove(counterName);
        localCounters.put(counterName, count.getCount());
        try {
            count.removeListener(sharedCountListeners.get(counterName));
            count.close();
        } catch (IOException e) {
            // ignore -- we're going to abandon this counter.
            log.warn("Error closing counter " + counterName + " after connection lost.", e);
        }
    }
}
 

/**
 * Creates a {@link ZooKeeperSharedCount} to store a shared count between multiple instances.
 *
 * @param path to the shared count in ZooKeeper
 * @param seedCount for the shared count
 * @return a shared count
 */
public ZooKeeperSharedCount createSharedCount(String path, int seedCount) {
	return new ZooKeeperSharedCount(
		new SharedCount(
			facade,
			path,
			seedCount));
}
 

/**
 * Creates a {@link ZooKeeperSharedCount} to store a shared count between multiple instances.
 *
 * @param path to the shared count in ZooKeeper
 * @param seedCount for the shared count
 * @return a shared count
 */
public ZooKeeperSharedCount createSharedCount(String path, int seedCount) {
	return new ZooKeeperSharedCount(
		new SharedCount(
			facade,
			path,
			seedCount));
}
 

public ZooKeeperSharedCount(SharedCount sharedCount) {
	this.sharedCount = Preconditions.checkNotNull(sharedCount);
}
 

@Test
public void testThreadedLeaseIncrease() throws Exception
{
    final Timing timing = new Timing();
    CuratorFramework client = CuratorFrameworkFactory.newClient(server.getConnectString(), timing.session(), timing.connection(), new RetryOneTime(1));
    try
    {
        client.start();

        final SharedCount count = new SharedCount(client, "/foo/count", 1);
        count.start();

        final InterProcessSemaphoreV2 semaphore = new InterProcessSemaphoreV2(client, "/test", count);

        ExecutorService service = Executors.newCachedThreadPool();

        final CountDownLatch latch1 = new CountDownLatch(1);
        final CountDownLatch latch2 = new CountDownLatch(1);
        Future<Object> future1 = service.submit
            (
                new Callable<Object>()
                {
                    @Override
                    public Object call() throws Exception
                    {
                        Lease lease = semaphore.acquire(timing.seconds(), TimeUnit.SECONDS);
                        Assert.assertNotNull(lease);
                        latch1.countDown();
                        lease = semaphore.acquire(timing.forWaiting().seconds(), TimeUnit.SECONDS);
                        Assert.assertNotNull(lease);
                        latch2.countDown();
                        return null;
                    }
                }
            );
        Future<Object> future2 = service.submit
            (
                new Callable<Object>()
                {
                    @Override
                    public Object call() throws Exception
                    {
                        Assert.assertTrue(latch1.await(timing.forWaiting().seconds(), TimeUnit.SECONDS));
                        timing.sleepABit(); // make sure second acquire is waiting
                        Assert.assertTrue(count.trySetCount(2));
                        //Make sure second acquire takes less than full waiting time:
                        timing.sleepABit();
                        Assert.assertTrue(latch2.await(0, TimeUnit.SECONDS));
                        return null;
                    }
                }
            );

        future1.get();
        future2.get();
    }
    finally
    {
        CloseableUtils.closeQuietly(client);
    }
}
 

public ZooKeeperSharedCount(SharedCount sharedCount) {
	this.sharedCount = Preconditions.checkNotNull(sharedCount);
}
 

public SharedCount createSharedCounter(String path, int seedValue) {
  return new SharedCount(client, path, seedValue);
}
 

@Test
public void testThreadedLeaseIncrease() throws Exception
{
    final Timing timing = new Timing();
    CuratorFramework client = CuratorFrameworkFactory.newClient(server.getConnectString(), timing.session(), timing.connection(), new RetryOneTime(1));
    try
    {
        client.start();

        final SharedCount count = new SharedCount(client, "/foo/count", 1);
        count.start();

        final InterProcessSemaphoreV2 semaphore = new InterProcessSemaphoreV2(client, "/test", count);

        ExecutorService service = Executors.newCachedThreadPool();

        final CountDownLatch latch1 = new CountDownLatch(1);
        final CountDownLatch latch2 = new CountDownLatch(1);
        Future<Object> future1 = service.submit
            (
                new Callable<Object>()
                {
                    @Override
                    public Object call() throws Exception
                    {
                        Lease lease = semaphore.acquire(timing.seconds(), TimeUnit.SECONDS);
                        Assert.assertNotNull(lease);
                        latch1.countDown();
                        lease = semaphore.acquire(timing.forWaiting().seconds(), TimeUnit.SECONDS);
                        Assert.assertNotNull(lease);
                        latch2.countDown();
                        return null;
                    }
                }
            );
        Future<Object> future2 = service.submit
            (
                new Callable<Object>()
                {
                    @Override
                    public Object call() throws Exception
                    {
                        Assert.assertTrue(latch1.await(timing.forWaiting().seconds(), TimeUnit.SECONDS));
                        timing.sleepABit(); // make sure second acquire is waiting
                        Assert.assertTrue(count.trySetCount(2));
                        //Make sure second acquire takes less than full waiting time:
                        timing.sleepABit();
                        Assert.assertTrue(latch2.await(0, TimeUnit.SECONDS));
                        return null;
                    }
                }
            );

        future1.get();
        future2.get();

        count.close();
    }
    finally
    {
        TestCleanState.closeAndTestClean(client);
    }
}
 

public static void main(String[] args) throws IOException, Exception {
	final Random rand = new Random();
	SharedCounterExample example = new SharedCounterExample();
	try (TestingServer server = new TestingServer()) {
		CuratorFramework client = CuratorFrameworkFactory.newClient(server.getConnectString(), new ExponentialBackoffRetry(1000, 3));
		client.start();
		
		SharedCount baseCount = new SharedCount(client, PATH, 0);
		baseCount.addListener(example);
		baseCount.start();
		
		List<SharedCount> examples = Lists.newArrayList();
		ExecutorService service = Executors.newFixedThreadPool(QTY);
		for (int i = 0; i < QTY; ++i) {
			final SharedCount count = new SharedCount(client, PATH, 0);
			examples.add(count);
			Callable<Void> task = new Callable<Void>() {
				@Override
				public Void call() throws Exception {
					count.start();
					Thread.sleep(rand.nextInt(10000));
					System.out.println("Increment:" + count.trySetCount(count.getVersionedValue(), count.getCount() + rand.nextInt(10)));
					return null;
				}
			};
			service.submit(task);
		}
		
		
		
		service.shutdown();
		service.awaitTermination(10, TimeUnit.MINUTES);
		
		for (int i = 0; i < QTY; ++i) {
			examples.get(i).close();
		}
		baseCount.close();
	}


}
 

/**
 * Creates a {@link ZooKeeperCheckpointIDCounter} instance.
 *
 * @param client      Curator ZooKeeper client
 * @param counterPath ZooKeeper path for the counter. It's sufficient to have a path per-job.
 */
public ZooKeeperCheckpointIDCounter(CuratorFramework client, String counterPath) {
	this.client = checkNotNull(client, "Curator client");
	this.counterPath = checkNotNull(counterPath, "Counter path");
	this.sharedCount = new SharedCount(client, counterPath, 1);
}
 

/**
 * Creates a {@link ZooKeeperCheckpointIDCounter} instance.
 *
 * @param client      Curator ZooKeeper client
 * @param counterPath ZooKeeper path for the counter. It's sufficient to have a path per-job.
 */
public ZooKeeperCheckpointIDCounter(CuratorFramework client, String counterPath) {
	this.client = checkNotNull(client, "Curator client");
	this.counterPath = checkNotNull(counterPath, "Counter path");
	this.sharedCount = new SharedCount(client, counterPath, 1);
}