CustomStopWatch.java

package com.ubercab.network.ramengrpc;

import java.util.concurrent.atomic.AtomicBoolean;

public class CustomStopWatch implements Runnable {

  private final long timeoutMs;
  private final TimeoutCallBack callBack;
  private AtomicBoolean eventTriggered = new AtomicBoolean(false);
  private AtomicBoolean shutdownInvoked = new AtomicBoolean(false);
  private AtomicBoolean paused = new AtomicBoolean(false);

  public CustomStopWatch(TimeoutCallBack callBack, long timeoutMs) {
    this.callBack = callBack;
    this.timeoutMs = timeoutMs;
  }

  @Override
  public void run() {
    try {
      waitUntilTimeout();
    } catch (InterruptedException e) {
      throw new IllegalStateException("IE thrown ", e);
    }
  }

  private synchronized void waitUntilTimeout() throws InterruptedException {
    while (!shutdownInvoked.get()) { // execute until shutdown is invoked
      eventTriggered.set(false); // reset eventTrigger just before going to wait
      this.wait(timeoutMs);
      if (!shutdownInvoked.get() && !paused.get()) { // if shutdown is not called and not paused
        // if shutdown not invoked
        if (!eventTriggered.get()) { // if no event has been triggered within timeout, then invoke
          // timeout callback
          callBack.onTimeout();
          // what to do after calling timeout.
          /*
          // may be put to pause state. Caller is expected to call restart to resume.
          paused.set(true);
          while(!shutdownInvoked.get() && paused.get()) {
          // need a while loop, bcoz, if resetTimer is called while in paused state, should still
          // be in waiting
          this.wait();
          }
           */
          // if above code is uncommented, remove the break;
          break;
        } else { // if event has been triggered, go to next cycle of waiting
          // go back to wait again
        }
      } else if (!shutdownInvoked.get()) { // if paused
        while(!shutdownInvoked.get() && paused.get()) {
          // need a while loop, bcoz, if resetTimer is called while in paused state, should still
          // be in waiting
          this.wait();
        }
      }
    }
  }

  public synchronized void shutdown() {
    shutdownInvoked.set(true);
    this.notifyAll();
  }

  public synchronized void resetTimer() {
    eventTriggered.set(true);
    this.notifyAll();
  }

  public synchronized void paused() {
    paused.set(true);
    notifyAll();
  }

  public synchronized void restart() {
    paused.set(false);
    notifyAll();
  }

  interface TimeoutCallBack {

    void onTimeout();
  }
}

TestCustomStopWatch.java

package com.ubercab.network.ramengrpc;

import com.ubercab.test.UberTestBase;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import org.junit.Test;

public class TestCustomStopWatch extends UberTestBase {

  @Test
  public void simpleTest1() throws InterruptedException {
    CountDownLatch countDownLatch = new CountDownLatch(1);
    TestTimeoutCallBack timeoutCallBack = new TestTimeoutCallBack(System.currentTimeMillis(), countDownLatch);

    CustomStopWatch stopWatch = new CustomStopWatch(timeoutCallBack, 2000);
    new Thread(stopWatch).start();
    Thread.sleep(1000);
    // timeoutCallBack.startTimeMs = System.currentTimeMillis();
    stopWatch.resetTimer();
    countDownLatch.await(10, TimeUnit.SECONDS);

    System.out.println("First timeout completed ");
    // after timeout called, need to again restart the stopwatch
    timeoutCallBack.countDownLatch = new CountDownLatch(1);
    timeoutCallBack.startTimeMs = System.currentTimeMillis();
    stopWatch.restart();
    countDownLatch.await(10, TimeUnit.SECONDS);
  }

  @Test
  public void simpleTest2() throws InterruptedException {
    CountDownLatch countDownLatch = new CountDownLatch(1);
    TestTimeoutCallBack timeoutCallBack = new TestTimeoutCallBack(System.currentTimeMillis(), countDownLatch);

    CustomStopWatch stopWatch = new CustomStopWatch(timeoutCallBack, 2000);
    new Thread(stopWatch).start();
    Thread.sleep(1000);
    timeoutCallBack.startTimeMs = System.currentTimeMillis();
    stopWatch.resetTimer();
    countDownLatch.await(10, TimeUnit.SECONDS);
  }

  @Test
  public void pauseTest() throws InterruptedException {
    CountDownLatch countDownLatch = new CountDownLatch(1);
    TestTimeoutCallBack timeoutCallBack = new TestTimeoutCallBack(System.currentTimeMillis(), countDownLatch);

    CustomStopWatch stopWatch = new CustomStopWatch(timeoutCallBack, 2000);
    new Thread(stopWatch).start();
    Thread.sleep(1000);
    stopWatch.paused();
    timeoutCallBack.startTimeMs = System.currentTimeMillis();
    Thread.sleep(3000);
    System.out.println("stop watch paused");
    stopWatch.restart();
    countDownLatch.await(10, TimeUnit.SECONDS);
  }

  class TestTimeoutCallBack implements CustomStopWatch.TimeoutCallBack {

    long startTimeMs;
    CountDownLatch countDownLatch;

    public TestTimeoutCallBack(long startTimeMs, CountDownLatch countDownLatch) {
      this.startTimeMs = startTimeMs;
      this.countDownLatch = countDownLatch;
    }

    @Override
    public void onTimeout() {
      long curTime = System.currentTimeMillis();
      System.out.println("Timeout called at "+  (curTime - startTimeMs) +", abs time "
          + System.currentTimeMillis());
      this.countDownLatch.countDown();
    }
  }

}