Java Code Examples for org.apache.twill.zookeeper.ZKClientService#stopAndWait()

@Test(timeout = 20000)
public void testReentrant() {
  // Test the lock is reentrant from the same thread
  ZKClientService zkClient = createZKClient();
  try {
    ReentrantDistributedLock lock = new ReentrantDistributedLock(zkClient, "reentrant");
    lock.lock();
    try {
      try {
        lock.lock();
      } finally {
        lock.unlock();
      }
    } finally {
      lock.unlock();
    }
  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test (timeout = 20000)
public void testMultiThreads() throws InterruptedException {
  // Test the lock mechanism between multiple threads
  ZKClientService zkClient = createZKClient();
  try {
    // Create two threads and compete for the lock
    final ReentrantDistributedLock lock = new ReentrantDistributedLock(zkClient, "multiThreads");
    final CountDownLatch acquired = new CountDownLatch(1);
    Thread t = new Thread() {
      @Override
      public void run() {
        lock.lock();
        try {
          acquired.countDown();
        } finally {
          lock.unlock();
        }
      }
    };

    lock.lock();
    try {
      t.start();
      // Wait for the thread to get the lock, should fail.
      Assert.assertFalse(acquired.await(1, TimeUnit.SECONDS));
    } finally {
      lock.unlock();
    }

    Assert.assertTrue(acquired.await(5, TimeUnit.SECONDS));
    t.join();

  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test (timeout = 20000)
public void testLockInterrupt() throws InterruptedException {
  // Test lock interruption on multiple threads.
  ZKClientService zkClient = createZKClient();
  try {
    final ReentrantDistributedLock lock = new ReentrantDistributedLock(zkClient, "/interrupt");

    // Create a new thread to acquire the same lock interruptibly.
    lock.lock();
    try {
      final CountDownLatch lockAcquired = new CountDownLatch(1);
      final CountDownLatch lockInterrupted = new CountDownLatch(1);
      Thread t = new Thread() {
        @Override
        public void run() {
          try {
            lock.lockInterruptibly();
            try {
              lockAcquired.countDown();
            } finally {
              lock.unlock();
            }
          } catch (InterruptedException e) {
            lockInterrupted.countDown();
          }
        }
      };

      t.start();
      t.interrupt();
      Assert.assertFalse(lockAcquired.await(2, TimeUnit.SECONDS));
      Assert.assertTrue(lockInterrupted.await(2, TimeUnit.SECONDS));
    } finally {
      lock.unlock();
    }
  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test
public void testControllerListener() throws InterruptedException {
  InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
  zkServer.startAndWait();

  LOG.info("ZKServer: " + zkServer.getConnectionStr());
  try {
    RunId runId = RunIds.generate();
    ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClientService.startAndWait();

    Service service = createService(zkClientService, runId);
    service.startAndWait();

    final CountDownLatch runLatch = new CountDownLatch(1);
    TwillController controller = getController(zkClientService, "testControllerListener", runId);
    controller.onRunning(new Runnable() {
      @Override
      public void run() {
        runLatch.countDown();
      }
    }, Threads.SAME_THREAD_EXECUTOR);

    Assert.assertTrue(runLatch.await(2, TimeUnit.SECONDS));

    service.stopAndWait();

    zkClientService.stopAndWait();
  } finally {
    zkServer.stopAndWait();
  }
}
 

@Test
public void testController() throws ExecutionException, InterruptedException, TimeoutException {
  InMemoryZKServer zkServer = InMemoryZKServer.builder().build();
  zkServer.startAndWait();

  LOG.info("ZKServer: " + zkServer.getConnectionStr());

  try {
    RunId runId = RunIds.generate();
    ZKClientService zkClientService = ZKClientService.Builder.of(zkServer.getConnectionStr()).build();
    zkClientService.startAndWait();

    Service service = createService(zkClientService, runId);
    service.startAndWait();

    TwillController controller = getController(zkClientService, "testController", runId);
    controller.sendCommand(Command.Builder.of("test").build()).get(2, TimeUnit.SECONDS);
    controller.terminate().get(2, TimeUnit.SECONDS);

    final CountDownLatch terminateLatch = new CountDownLatch(1);
    service.addListener(new ServiceListenerAdapter() {
      @Override
      public void terminated(Service.State from) {
        terminateLatch.countDown();
      }
    }, Threads.SAME_THREAD_EXECUTOR);

    Assert.assertTrue(service.state() == Service.State.TERMINATED || terminateLatch.await(2, TimeUnit.SECONDS));

    zkClientService.stopAndWait();

  } finally {
    zkServer.stopAndWait();
  }
}
 

@Test (timeout = 20000)
public void testTryLockMultiClients() throws InterruptedException {
  // Test tryLock on multiple clients
  ZKClientService zkClient1 = createZKClient();
  ZKClientService zkClient2 = createZKClient();
  try {
    final ReentrantDistributedLock lock1 = new ReentrantDistributedLock(zkClient1, "/multiTrylock");
    final ReentrantDistributedLock lock2 = new ReentrantDistributedLock(zkClient2, "/multiTrylock");

    // Create a new thread to acquire the same lock by using tryLock
    final CountDownLatch lockAcquired = new CountDownLatch(1);
    final CountDownLatch lockFailed = new CountDownLatch(2);
    Thread t = new Thread() {
      @Override
      public void run() {
        try {
          if (!lock2.tryLock()) {
            lockFailed.countDown();
          }
          if (!lock2.tryLock(1, TimeUnit.SECONDS)) {
            lockFailed.countDown();
          }
          if (lock2.tryLock(5000, TimeUnit.SECONDS)) {
            // Test the reentrant park as well
            if (lock2.tryLock(1, TimeUnit.SECONDS)) {
              if (lock2.tryLock()) {
                lockAcquired.countDown();
                lock2.unlock();
              }
              lock2.unlock();
            }
            lock2.unlock();
          }
        } catch (InterruptedException e) {
          // Shouldn't happen
          LOG.error(e.getMessage(), e);
        }
      }
    };

    lock1.lock();
    try {
      t.start();
      Assert.assertTrue(lockFailed.await(5000, TimeUnit.SECONDS));
    } finally {
      lock1.unlock();
    }
    Assert.assertTrue(lockAcquired.await(10000, TimeUnit.SECONDS));
    t.join();

    // Try to lock again and it should success immediately.
    Assert.assertTrue(lock1.tryLock());
    lock1.unlock();
  } finally {
    zkClient1.stopAndWait();
    zkClient2.stopAndWait();
  }
}
 

@Test
public void testBrokerChange() throws Exception {
  // Create a new namespace in ZK for Kafka server for this test case
  String connectionStr = zkServer.getConnectionStr() + "/broker_change";
  ZKClientService zkClient = ZKClientService.Builder.of(connectionStr).build();
  zkClient.startAndWait();
  zkClient.create("/", null, CreateMode.PERSISTENT).get();

  // Start a new kafka server
  File logDir = TMP_FOLDER.newFolder();
  EmbeddedKafkaServer server = new EmbeddedKafkaServer(generateKafkaConfig(connectionStr, logDir));
  server.startAndWait();

  // Start a Kafka client
  KafkaClientService kafkaClient = new ZKKafkaClientService(zkClient);
  kafkaClient.startAndWait();

  // Attach a consumer
  final BlockingQueue<String> consumedMessages = Queues.newLinkedBlockingQueue();
  kafkaClient.getConsumer()
    .prepare().addFromBeginning("test", 0).consume(new KafkaConsumer.MessageCallback() {
    @Override
    public long onReceived(Iterator<FetchedMessage> messages) {
      long nextOffset = -1L;
      while (messages.hasNext()) {
        FetchedMessage message = messages.next();
        nextOffset = message.getNextOffset();
        consumedMessages.add(Charsets.UTF_8.decode(message.getPayload()).toString());
      }
      return nextOffset;
    }

    @Override
    public void finished() {
      // No-op
    }
  });

  // Get a publisher and publish a message
  KafkaPublisher publisher = kafkaClient.getPublisher(KafkaPublisher.Ack.FIRE_AND_FORGET, Compression.NONE);
  publisher.prepare("test").add(Charsets.UTF_8.encode("Message 0"), 0).send().get();

  // Should receive one message
  Assert.assertEquals("Message 0", consumedMessages.poll(5, TimeUnit.SECONDS));

  // Now shutdown and restart the server on different port
  server.stopAndWait();
  server = new EmbeddedKafkaServer(generateKafkaConfig(connectionStr, logDir));
  server.startAndWait();

  // Wait a little while to make sure changes is reflected in broker service
  TimeUnit.SECONDS.sleep(3);

  // Now publish again with the same publisher. It should succeed and the consumer should receive the message.
  publisher.prepare("test").add(Charsets.UTF_8.encode("Message 1"), 0).send().get();
  Assert.assertEquals("Message 1", consumedMessages.poll(5, TimeUnit.SECONDS));

  kafkaClient.stopAndWait();
  zkClient.stopAndWait();
  server.stopAndWait();
}
 

@Test
public void testKafkaClientReconnect() throws Exception {
  String topic = "backoff";
  Properties kafkaServerConfig = generateKafkaConfig(zkServer.getConnectionStr() + "/backoff");
  EmbeddedKafkaServer server = new EmbeddedKafkaServer(kafkaServerConfig);

  ZKClientService zkClient = ZKClientService.Builder.of(zkServer.getConnectionStr() + "/backoff").build();
  zkClient.startAndWait();
  try {
    zkClient.create("/", null, CreateMode.PERSISTENT).get();

    ZKKafkaClientService kafkaClient = new ZKKafkaClientService(zkClient);
    kafkaClient.startAndWait();

    try {
      server.startAndWait();
      try {
        // Publish a messages
        createPublishThread(kafkaClient, topic, Compression.NONE, "First message", 1).start();

        // Create a consumer
        final BlockingQueue<String> queue = new LinkedBlockingQueue<>();
        Cancellable cancel = kafkaClient.getConsumer().prepare().add(topic, 0, 0)
          .consume(new KafkaConsumer.MessageCallback() {
            @Override
            public long onReceived(Iterator<FetchedMessage> messages) {
              long nextOffset = -1L;
              while (messages.hasNext()) {
                FetchedMessage message = messages.next();
                nextOffset = message.getNextOffset();
                queue.offer(Charsets.UTF_8.decode(message.getPayload()).toString());
              }
              return nextOffset;
            }

            @Override
            public void finished() {
            }
          });

        // Wait for the first message
        Assert.assertEquals("0 First message", queue.poll(60, TimeUnit.SECONDS));

        // Shutdown the server
        server.stopAndWait();

        // Start the server again.
        // Needs to create a new instance with the same config since guava service cannot be restarted
        server = new EmbeddedKafkaServer(kafkaServerConfig);
        server.startAndWait();

        // Wait a little while to make sure changes is reflected in broker service
        TimeUnit.SECONDS.sleep(3);

        // Publish another message
        createPublishThread(kafkaClient, topic, Compression.NONE, "Second message", 1).start();

        // Should be able to get the second message
        Assert.assertEquals("0 Second message", queue.poll(60, TimeUnit.SECONDS));

        cancel.cancel();
      } finally {
        kafkaClient.stopAndWait();
      }
    } finally {
      server.stopAndWait();
    }
  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test
public void testZKCleanup() throws Exception {
  final ZKClientService zkClient = ZKClientService.Builder.of(getZKConnectionString() + "/twill").build();
  zkClient.startAndWait();

  try {
    TwillRunner runner = getTwillRunner();

    // Start an application and stop it.
    TwillController controller = runner.prepare(new EchoServer())
      .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
      .withApplicationArguments("echo")
      .withArguments("EchoServer", "echo2")
      .start();

    Iterable<Discoverable> echoServices = controller.discoverService("echo");
    Assert.assertTrue(waitForSize(echoServices, 1, 120));

    controller.terminate().get();

    // Verify the ZK node gets cleanup
    waitFor(null, new Callable<Stat>() {
      @Override
      public Stat call() throws Exception {
        return zkClient.exists("/EchoServer").get();
      }
    }, 10000, 100, TimeUnit.MILLISECONDS);

    // Start two instances of the application and stop one of it
    List<TwillController> controllers = new ArrayList<>();
    for (int i = 0; i < 2; i++) {
      controllers.add(runner.prepare(new EchoServer())
                        .addLogHandler(new PrinterLogHandler(new PrintWriter(System.out, true)))
                        .withApplicationArguments("echo")
                        .withArguments("EchoServer", "echo2")
                        .start());
    }

    // There should be two instances up and running.
    echoServices = controllers.get(1).discoverService("echo");
    Assert.assertTrue(waitForSize(echoServices, 2, 120));

    // Stop one instance of the app
    controllers.get(0).terminate().get();

    // Verify the ZK node should still be there
    Assert.assertNotNull(zkClient.exists("/EchoServer").get());

    // We should still be able to do discovery, which depends on the ZK node.
    echoServices = controllers.get(1).discoverService("echo");
    Assert.assertTrue(waitForSize(echoServices, 1, 120));

    // Stop second instance of the app
    controllers.get(1).terminate().get();

    // Verify the ZK node gets cleanup
    waitFor(null, new Callable<Stat>() {
      @Override
      public Stat call() throws Exception {
        return zkClient.exists("/EchoServer").get();
      }
    }, 10000, 100, TimeUnit.MILLISECONDS);

  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test(timeout = 150000)
public void testDisconnect() throws IOException, InterruptedException {
  File zkDataDir = tmpFolder.newFolder();
  InMemoryZKServer ownZKServer = InMemoryZKServer.builder().setDataDir(zkDataDir).build();
  ownZKServer.startAndWait();
  try {
    ZKClientService zkClient = ZKClientService.Builder.of(ownZKServer.getConnectionStr()).build();
    zkClient.startAndWait();

    try {
      final Semaphore leaderSem = new Semaphore(0);
      final Semaphore followerSem = new Semaphore(0);

      LeaderElection leaderElection = new LeaderElection(zkClient, "/testDisconnect", new ElectionHandler() {
        @Override
        public void leader() {
          leaderSem.release();
        }

        @Override
        public void follower() {
          followerSem.release();
        }
      });
      leaderElection.start();

      leaderSem.tryAcquire(20, TimeUnit.SECONDS);

      int zkPort = ownZKServer.getLocalAddress().getPort();

      // Disconnect by shutting the server and restart it on the same port
      ownZKServer.stopAndWait();

      // Right after disconnect, it should become follower
      followerSem.tryAcquire(20, TimeUnit.SECONDS);

      ownZKServer = InMemoryZKServer.builder().setDataDir(zkDataDir).setPort(zkPort).build();
      ownZKServer.startAndWait();

      // Right after reconnect, it should be leader again.
      leaderSem.tryAcquire(20, TimeUnit.SECONDS);

      // Now disconnect it again, but then cancel it before reconnect, it shouldn't become leader
      ownZKServer.stopAndWait();

      // Right after disconnect, it should become follower
      followerSem.tryAcquire(20, TimeUnit.SECONDS);

      ListenableFuture<?> cancelFuture = leaderElection.stop();

      ownZKServer = InMemoryZKServer.builder().setDataDir(zkDataDir).setPort(zkPort).build();
      ownZKServer.startAndWait();

      Futures.getUnchecked(cancelFuture);

      // After reconnect, it should not be leader
      Assert.assertFalse(leaderSem.tryAcquire(10, TimeUnit.SECONDS));
    } finally {
      zkClient.stopAndWait();
    }
  } finally {
    ownZKServer.stopAndWait();
  }
}
 

@Test (timeout = 60000)
public void testLockSessionExpire() throws Exception {
  ZKClientService zkClient1 = createZKClient();
  final ZKClientService zkClient2 = createZKClient();

  try {
    Lock lock1 = new ReentrantDistributedLock(zkClient1, "/lockexpire");
    lock1.lock();
    try {
      final CountDownLatch started = new CountDownLatch(1);
      final CountDownLatch failure = new CountDownLatch(1);
      Thread t = new Thread() {
        @Override
        public void run() {
          try {
            started.countDown();
            Lock lock2 = new ReentrantDistributedLock(zkClient2, "/lockexpire");
            lock2.lock();
            try {
              // No-op
            } finally {
              lock2.unlock();
            }
          } catch (Exception ex) {
            failure.countDown();
          }
        }
      };
      t.start();

      Assert.assertTrue(started.await(5, TimeUnit.SECONDS));

      // Wait until the lock2 is waiting.
      while (zkClient1.getChildren("/lockexpire").get().getChildren().size() != 2) {
        TimeUnit.SECONDS.sleep(1);
      }

      // Session timeout for zkclient2, the lock acquisition should failed
      KillZKSession.kill(zkClient2.getZooKeeperSupplier().get(), zkClient2.getConnectString(), 10000);
      Assert.assertTrue(failure.await(10, TimeUnit.SECONDS));

    } finally {
      lock1.unlock();
    }

  } finally {
    zkClient1.stopAndWait();
    zkClient2.stopAndWait();
  }
}
 

@Test (timeout = 60000)
public void testLockRace() throws Exception {
  // Test for multiple clients race on the lock
  // This is for the case when a lock owner release the lock (node deleted)
  // while the other client tries to watch on the node
  ZKClientService zkClient1 = createZKClient();
  ZKClientService zkClient2 = createZKClient();
  try {
    final Lock[] locks = new Lock[] {
      new ReentrantDistributedLock(zkClient1, "/lockrace"),
      new ReentrantDistributedLock(zkClient2, "/lockrace")
    };

    // Have two clients fight for the lock
    final CyclicBarrier barrier = new CyclicBarrier(2);
    final CountDownLatch lockLatch = new CountDownLatch(2);
    for (int i = 0; i < 2; i++) {
      final int threadId = i;
      Thread t = new Thread() {
        @Override
        public void run() {
          try {
            barrier.await();
            for (int i = 0; i < 100; i++) {
              Lock lock = locks[threadId];
              lock.lock();
              try {
                // A short sleep to make the race possible to happen
                Uninterruptibles.sleepUninterruptibly(1, TimeUnit.MILLISECONDS);
              } finally {
                lock.unlock();
              }
            }
            lockLatch.countDown();
          } catch (Exception e) {
            LOG.error("Exception", e);
          }
        }
      };
      t.start();
    }

    Assert.assertTrue(lockLatch.await(30, TimeUnit.SECONDS));

  } finally {
    zkClient1.stopAndWait();
    zkClient2.stopAndWait();
  }
}
 

@VisibleForTesting
public static void doMain(boolean verbose, Configuration conf) throws Exception {
  LOG.info("Starting tx server client test.");
  Injector injector = Guice.createInjector(
    new ConfigModule(conf),
    new ZKModule(),
    new DiscoveryModules().getDistributedModules(),
    new TransactionModules().getDistributedModules(),
    new TransactionClientModule()
  );

  ZKClientService zkClient = injector.getInstance(ZKClientService.class);
  zkClient.startAndWait();

  try {
    TransactionServiceClient client = injector.getInstance(TransactionServiceClient.class);
    LOG.info("Starting tx...");
    Transaction tx = client.startShort();
    if (verbose) {
      LOG.info("Started tx details: " + tx.toString());
    } else {
      LOG.info("Started tx: " + tx.getTransactionId() +
                 ", readPointer: " + tx.getReadPointer() +
                 ", invalids: " + tx.getInvalids().length +
                 ", inProgress: " + tx.getInProgress().length);
    }
    try {
      LOG.info("Checking if canCommit tx...");
      client.canCommitOrThrow(tx, Collections.<byte[]>emptyList());
      LOG.info("canCommit: success");
      LOG.info("Committing tx...");
      client.commitOrThrow(tx);
      LOG.info("Committed tx: success");
    } catch (TransactionConflictException e) {
      LOG.info("Aborting tx...");
      client.abort(tx);
      LOG.info("Aborted tx...");
    }
  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test (timeout = 20000)
public void testTryLock() throws InterruptedException {
  // Test tryLock on multiple threads
  ZKClientService zkClient = createZKClient();
  try {
    final ReentrantDistributedLock lock = new ReentrantDistributedLock(zkClient, "/trylock");

    // Create a new thread to acquire the same lock by using tryLock
    final CountDownLatch lockAcquired = new CountDownLatch(1);
    final CountDownLatch lockFailed = new CountDownLatch(2);
    Thread t = new Thread() {
      @Override
      public void run() {
        try {
          if (!lock.tryLock()) {
            lockFailed.countDown();
          }
          if (!lock.tryLock(1, TimeUnit.SECONDS)) {
            lockFailed.countDown();
          }
          if (lock.tryLock(5, TimeUnit.SECONDS)) {
            // Test the reentrant park as well
            if (lock.tryLock(1, TimeUnit.SECONDS)) {
              if (lock.tryLock()) {
                lockAcquired.countDown();
                lock.unlock();
              }
              lock.unlock();
            }
            lock.unlock();
          }
        } catch (InterruptedException e) {
          // Shouldn't happen
          LOG.error(e.getMessage(), e);
        }
      }
    };

    lock.lock();
    try {
      t.start();
      Assert.assertTrue(lockFailed.await(5, TimeUnit.SECONDS));
    } finally {
      lock.unlock();
    }
    Assert.assertTrue(lockAcquired.await(2, TimeUnit.SECONDS));
    t.join();

    // Try to lock again and it should success immediately.
    Assert.assertTrue(lock.tryLock());
    lock.unlock();
  } finally {
    zkClient.stopAndWait();
  }
}
 

@Test (timeout = 200000)
public void testLockInterruptMultiClients() throws InterruptedException {
  // Test lock interruption from multiple clients
  ZKClientService zkClient1 = createZKClient();
  ZKClientService zkClient2 = createZKClient();
  try {
    final ReentrantDistributedLock lock1 = new ReentrantDistributedLock(zkClient1, "/multiInterrupt");
    final ReentrantDistributedLock lock2 = new ReentrantDistributedLock(zkClient2, "/multiInterrupt");

    lock1.lock();
    // Create a new thread to acquire a lock interruptibly
    try {
      final CountDownLatch lockAcquired = new CountDownLatch(1);
      final CountDownLatch lockInterrupted = new CountDownLatch(1);
      Thread t = new Thread() {
        @Override
        public void run() {
          try {
            lock2.lockInterruptibly();
            try {
              lockAcquired.countDown();
            } finally {
              lock2.unlock();
            }
          } catch (InterruptedException e) {
            lockInterrupted.countDown();
          }
        }
      };

      t.start();

      // Sleep for 2 seconds to make sure the lock2 in the thread already entered the distributed lock phase.
      TimeUnit.SECONDS.sleep(2);
      t.interrupt();
      Assert.assertFalse(lockAcquired.await(2, TimeUnit.SECONDS));
      Assert.assertTrue(lockInterrupted.await(2, TimeUnit.SECONDS));

    } finally {
      lock1.unlock();
    }

    // After lock1 unlocked, try to lock with lock2 again and it should succeed.
    // This is to verify when a lock is interrupted, it cleanup the states in ZK
    lock2.lock();
    lock2.unlock();

  } finally {
    zkClient1.stopAndWait();
    zkClient2.stopAndWait();
  }
}
 

@Test (timeout = 20000)
public void testReentrantMultiClients() throws InterruptedException {
  // Test the reentrant lock/unlock mechanism works correctly across multiple instances of lock
  // This is to simulate the lock between multiple processes
  ZKClientService zkClient1 = createZKClient();
  ZKClientService zkClient2 = createZKClient();
  try {
    // Create two distributed locks
    ReentrantDistributedLock lock1 = new ReentrantDistributedLock(zkClient1, "multiClients");
    final ReentrantDistributedLock lock2 = new ReentrantDistributedLock(zkClient2, "multiClients");

    // Create a thread for locking with lock2
    final CountDownLatch lockAcquired = new CountDownLatch(1);
    Thread t = new Thread() {
      @Override
      public void run() {
        lock2.lock();
        try {
          lockAcquired.countDown();
        } finally {
          lock2.unlock();
        }
      }
    };

    // Lock with lock1
    lock1.lock();
    try {
      // Start the lock2 thread. It should be blocked at acquiring the lock
      t.start();
      Assert.assertFalse(lockAcquired.await(1, TimeUnit.SECONDS));

      // Lock with lock1 again and unlock. The inner unlock shouldn't release the distributed lock
      lock1.lock();
      lock1.unlock();

      Assert.assertFalse(lockAcquired.await(1, TimeUnit.SECONDS));
    } finally {
      lock1.unlock();
    }

    // Now lock1 is unlocked, the lock2 thread should be able to proceed
    Assert.assertTrue(lockAcquired.await(5, TimeUnit.SECONDS));
    t.join();
  } finally {
    zkClient1.stopAndWait();
    zkClient2.stopAndWait();
  }
}
 

@Test (timeout = 30000)
public void testDoubleRegister() throws Exception {
  Map.Entry<DiscoveryService, DiscoveryServiceClient> entry = create();
  try {
    DiscoveryService discoveryService = entry.getKey();
    DiscoveryServiceClient discoveryServiceClient = entry.getValue();

    // Register on the same host port, it shouldn't fail.
    Cancellable cancellable = register(discoveryService, "test_double_reg", "localhost", 54321);
    Cancellable cancellable2 = register(discoveryService, "test_double_reg", "localhost", 54321);

    ServiceDiscovered discoverables = discoveryServiceClient.discover("test_double_reg");

    Assert.assertTrue(waitTillExpected(1, discoverables));

    cancellable.cancel();
    cancellable2.cancel();

    // Register again with two different clients, but killing session of the first one.
    final ZKClientService zkClient2 = ZKClientServices.delegate(
      ZKClients.retryOnFailure(
        ZKClients.reWatchOnExpire(
          ZKClientService.Builder.of(zkServer.getConnectionStr()).build()),
        RetryStrategies.fixDelay(1, TimeUnit.SECONDS)));
    zkClient2.startAndWait();

    try (ZKDiscoveryService discoveryService2 = new ZKDiscoveryService(zkClient2)) {
      cancellable2 = register(discoveryService2, "test_multi_client", "localhost", 54321);

      // Schedule a thread to shutdown zkClient2.
      new Thread() {
        @Override
        public void run() {
          try {
            TimeUnit.SECONDS.sleep(2);
            zkClient2.stopAndWait();
          } catch (InterruptedException e) {
            LOG.error(e.getMessage(), e);
          }
        }
      }.start();

      // This call would block until zkClient2 is shutdown.
      cancellable = register(discoveryService, "test_multi_client", "localhost", 54321);
      cancellable.cancel();
    } finally {
      zkClient2.stopAndWait();
    }
  } finally {
    closeServices(entry);
  }
}