How to use

ExtendedLatch

in

com.dremio.common.concurrent

/**
 * Waits until it is safe to exit. Blocks until all currently running fragments have completed.
 *
 * <p>This is intended to be used by {@link com.dremio.exec.server.SabotNode#close()}.</p>
 */
public void waitToExit() {
 synchronized(this) {
  if (externalIdToForeman.isEmpty()) {
   return;
  }
  exitLatch = new ExtendedLatch();
 }
 // Wait for at most 5 seconds or until the latch is released.
 exitLatch.awaitUninterruptibly(5000);
}

 /**
  * Await without interruption. In the case of interruption, log a warning and continue to wait.
  */
 public void awaitUninterruptibly() {
  while (true) {
   try {
    await();
    return;
   } catch (final InterruptedException e) {
    // if we're still not ready, the while loop will cause us to wait again
    logger.warn("Interrupted while waiting for event latch.", e);
   }
  }
 }
}

@Override
public void countDown() {
 Preconditions.checkNotNull(latch, "Latch not initialized in %s at %s.", siteClass.getSimpleName(), desc);
 Preconditions.checkArgument(latch.getCount() > 0, "Counting down on latch more than intended.");
 latch.countDown();
}

final ExtendedLatch endpointLatch = new ExtendedLatch(numIntFragments);
final FragmentSubmitFailures fragmentSubmitFailures = new FragmentSubmitFailures();
if(numIntFragments > 0 && !endpointLatch.awaitUninterruptibly(timeout)){
 long numberRemaining = endpointLatch.getCount();
 throw UserException.connectionError()
   .message(

@Override
public void initialize(final int count) {
 Preconditions.checkArgument(latch == null, "Latch can be initialized only once at %s in %s.", desc,
  siteClass.getSimpleName());
 Preconditions.checkArgument(count > 0, "Count has to be a positive integer at %s in %s.", desc,
  siteClass.getSimpleName());
 latch = new ExtendedLatch(count);
}

@Override
public void awaitUninterruptibly() {
 Preconditions.checkNotNull(latch, "Latch not initialized in %s at %s.", siteClass.getSimpleName(), desc);
 latch.awaitUninterruptibly();
}

 public void unpause() {
  latch.countDown();
 }
}

@Test
public void pauseInjected() {
 final long expectedDuration = 1000L;
 final ExtendedLatch trigger = new ExtendedLatch(1);
 final Pointer<Exception> ex = new Pointer<>();
 final String controls = Controls.newBuilder()
  .addPause(DummyClass.class, DummyClass.PAUSES)
  .build();
 ControlsInjectionUtil.setControls(session, controls);
 final QueryContext queryContext = new QueryContext(session, nodes[0].getContext(), QueryId.getDefaultInstance());
 (new ResumingThread(queryContext, trigger, ex, expectedDuration)).start();
 // test that the pause happens
 final DummyClass dummyClass = new DummyClass(queryContext, trigger);
 final long actualDuration = dummyClass.pauses();
 assertTrue(String.format("Test should stop for at least %d milliseconds.", expectedDuration),
  expectedDuration <= actualDuration);
 assertTrue("No exception should be thrown.", ex.value == null);
 try {
  queryContext.close();
 } catch (final Exception e) {
  fail("Failed to close query context: " + e);
 }
}

public void pause() {
 if (!injectNow()) {
  return;
 }
 latch.awaitUninterruptibly();
}

/**
 * If it is safe to exit, and the exitLatch is in use, signals it so that waitToExit() will
 * unblock.
 */
private void indicateIfSafeToExit() {
 synchronized(this) {
  if (exitLatch != null) {
   if (externalIdToForeman.isEmpty()) {
    exitLatch.countDown();
   }
  }
 }
}

/**
 * Waits until it is safe to exit. Blocks until all currently running fragments have completed.
 *
 * <p>This is intended to be used by {@link com.dremio.exec.server.SabotNode#close()}.</p>
 */
public void waitToExit() {
 synchronized(this) {
  if (fragmentExecutors == null || fragmentExecutors.size() == 0) {
   return;
  }
  exitLatch = new ExtendedLatch();
 }
 // Wait for at most 5 seconds or until the latch is released.
 exitLatch.awaitUninterruptibly(5000);
}

@Override
public void await() throws InterruptedException {
 Preconditions.checkNotNull(latch, "Latch not initialized in %s at %s.", siteClass.getSimpleName(), desc);
 try {
  latch.await();
 } catch (final InterruptedException e) {
  logger.warn("Interrupted while awaiting in {} at {}.", siteClass.getSimpleName(), desc);
  throw e;
 }
}

final ExtendedLatch trigger = new ExtendedLatch(1);
final Pointer<Exception> ex = new Pointer<>();
final QueryContext queryContext = new QueryContext(session, nodeContext1, QueryId.getDefaultInstance());
final ExtendedLatch trigger = new ExtendedLatch(1);
final QueryContext queryContext = new QueryContext(session, nodeContext2, QueryId.getDefaultInstance());

 @Override
 public void run() {
  latch.awaitUninterruptibly();
  final long startTime = System.currentTimeMillis();
  for (int i = 0; i < count; i++) {
   (new Thread() {
    @Override
    public void run() {
     dummyClass.countDown();
    }
   }).start();
  }
  final long endTime = System.currentTimeMillis();
  countingDownTime.value = (endTime - startTime);
 }
}

/**
 * If it is safe to exit, and the exitLatch is in use, signals it so that waitToExit() will
 * unblock.
 */
private void indicateIfSafeToExit() {
 synchronized(this) {
  if (exitLatch != null) {
   if (fragmentExecutors.size() == 0) {
    exitLatch.countDown();
   }
  }
 }
}

}else{
 parallelism = Math.min(parallelism,  runnables.size());
 final ExtendedLatch latch = new ExtendedLatch(runnables.size());
 final ExecutorService threadPool = Executors.newFixedThreadPool(parallelism);
 try{
  if (!latch.awaitUninterruptibly(timeout)) {

/**
 * Await without interruption for a given time.
 * @param waitMillis
 *          Time in milliseconds to wait
 * @return Whether the countdown reached zero or not.
 */
public boolean awaitUninterruptibly(long waitMillis) {
 final long targetMillis = System.currentTimeMillis() + waitMillis;
 while (System.currentTimeMillis() < targetMillis) {
  final long wait = targetMillis - System.currentTimeMillis();
  if (wait < 1) {
   return false;
  }
  try {
   return await(wait, TimeUnit.MILLISECONDS);
  } catch (final InterruptedException e) {
   // if we weren't ready, the while loop will continue to wait
  }
 }
 return false;
}

final ExtendedLatch trigger = new ExtendedLatch(1);
final Pointer<Long> countingDownTime = new Pointer<>();

 @Override
 public void run() {
  latch.awaitUninterruptibly();
  try {
   Thread.sleep(millis);
  } catch (final InterruptedException ex) {
   this.ex.value = ex;
  }
  context.getExecutionControls().unpauseAll();
 }
}

public void interruptiblePause() throws InterruptedException {
 if (!injectNow()) {
  return;
 }
 latch.await();
}