public AutoCloseable trackHeartbeats(int numHeartbeats) {
    activeHeartbeats.getAndAdd(numHeartbeats);
    return () -> activeHeartbeats.getAndAdd(-numHeartbeats);
  }
@Test
  public void testTrackHeartbeats() throws InterruptedException {
    CountDownLatch processCall = new CountDownLatch(1);
    CountDownLatch callProcessing = new CountDownLatch(1);
    CountDownLatch processingDone = new CountDownLatch(1);
    int numHeartbeats = 5;
    getDataProcessor.submit(
        () -> {
          try (AutoCloseable ignored = getDataMetricTracker.trackHeartbeats(numHeartbeats)) {
            callProcessing.countDown();
            processCall.await();
          } catch (Exception e) {
            // Do nothing.
          }
          processingDone.countDown();
        });

    callProcessing.await();
    ThrottlingGetDataMetricTracker.ReadOnlySnapshot metricsWhileProcessing =
        getDataMetricTracker.getMetricsSnapshot();

    assertThat(metricsWhileProcessing.activeHeartbeats()).isEqualTo(5);

    // Free the thread inside the AutoCloseable, wait for processingDone and check that metrics gets
    // decremented
    processCall.countDown();
    processingDone.await();
    ThrottlingGetDataMetricTracker.ReadOnlySnapshot metricsAfterProcessing =
        getDataMetricTracker.getMetricsSnapshot();
    assertThat(metricsAfterProcessing.activeHeartbeats()).isEqualTo(0);
  }