IteratorExecutorTest.java

/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *    http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package gobblin.util.executors;

import java.util.Iterator;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Condition;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.Test;

import com.google.common.base.Predicate;
import com.google.common.util.concurrent.ThreadFactoryBuilder;

import gobblin.testing.AssertWithBackoff;


public class IteratorExecutorTest {

  private final AtomicInteger nextCallCount = new AtomicInteger(0);
  private final AtomicInteger completedCount = new AtomicInteger(0);
  private final Logger log = LoggerFactory.getLogger(IteratorExecutorTest.class);

  @Test
  public void test() throws Exception {

    TestIterator iterator = new TestIterator(5);

    final IteratorExecutor<Void> executor = new IteratorExecutor<>(iterator, 2, new ThreadFactoryBuilder().build());

    log.info("Starting executor");
    ExecutorService singleThread = Executors.newSingleThreadExecutor();
    final Future<?> future = singleThread.submit(new Runnable() {
      @Override
      public void run() {
        try {
          executor.execute();
        } catch (InterruptedException ie) {
          throw new RuntimeException(ie);
        }
      }
    });

    log.info("Only two threads, so exactly two tasks retrieved");
    verify(2, 0);
    Assert.assertFalse(future.isDone());
    Assert.assertTrue(iterator.hasNext());

    log.info("end one of the tasks");
    iterator.endOneCallable();

    log.info("three tasks retrieved, one completed");
    verify(3, 1);
    Assert.assertFalse(future.isDone());
    Assert.assertTrue(iterator.hasNext());

    log.info("end two tasks");
    iterator.endOneCallable();
    iterator.endOneCallable();

    log.info("5 (all) tasks retrieved, 3 completed");
    verify(5, 3);
    Assert.assertFalse(future.isDone());
    Assert.assertFalse(iterator.hasNext());

    log.info("end two tasks");
    iterator.endOneCallable();
    iterator.endOneCallable();

    log.info("all tasks completed, check future is done");
    verify(5, 5);

    AssertWithBackoff.assertTrue(new Predicate<Void>() {
      @Override public boolean apply(Void input) {
        return future.isDone();
      }
    }, 10000, "future done", log, 2, 1000);

    Assert.assertTrue(future.isDone());

    log.info("done.");
  }

  /**
   * Verify exactly retrieved tasks have been retrieved, and exactly completed tasks have completed.
   */
  private void verify(final int retrieved, final int completed) throws Exception {
    AssertWithBackoff.assertTrue(new Predicate<Void>() {
      @Override public boolean apply(Void input) {
        log.info("Waiting for {} retrieved -> {} and {} completed -> {}", retrieved,
                 nextCallCount.get(), completed, completedCount.get());
        return (nextCallCount.get() == retrieved && completedCount.get() == completed);
      }
    }, 30000, "Waiting for callcount retrieved=" + retrieved + "/completed=" + completed, log,
        1.5, 1000);
  }

  /**
   * Runnable iterator. Runnables block until a call to {@link #endOneCallable()}.
   */
  private class TestIterator implements Iterator<Callable<Void>> {

    private final int maxCallables;
    private Lock lock = new ReentrantLock();
    private Condition endCallTransitionCondition = lock.newCondition();
    private boolean endCallEnqueued = false;

    public TestIterator(int maxCallables) {
      this.maxCallables = maxCallables;
    }

    @Override
    public boolean hasNext() {
      return nextCallCount.get() < maxCallables;
    }

    @Override
    public Callable<Void> next() {
      nextCallCount.incrementAndGet();

      return new Callable<Void>() {
        @Override
        public Void call()
            throws Exception {

          log.debug("Blocking at {}", nextCallCount.get());
          lock.lock();
          try {
            // wait for endCall request
            while (!endCallEnqueued) {
              if (!endCallTransitionCondition.await(30, TimeUnit.SECONDS)) {
                throw new TimeoutException("Waiting for endCall to be enqueued: retrieved=" +
                      IteratorExecutorTest.this.nextCallCount.get()
                      + " completed=" + IteratorExecutorTest.this.completedCount.get()
                      + " endCallEnqueued=" + endCallEnqueued);
              }
            }
            completedCount.incrementAndGet();
            endCallEnqueued = false;
            endCallTransitionCondition.signalAll();
          } finally {
            lock.unlock();
          }

          log.debug("Completed {}", completedCount.get());
          return null;
        }
      };
    }

    public void endOneCallable() throws InterruptedException, TimeoutException {
      log.debug("End one {}", nextCallCount.get());
      lock.lock();
      try {
        // wait till the last endCall was processed
        while (endCallEnqueued) {
          if (!endCallTransitionCondition.await(30, TimeUnit.SECONDS)) {
            throw new TimeoutException("Waiting for endCall to be processed: retrieved=" +
                IteratorExecutorTest.this.nextCallCount.get()
                + " completed=" + IteratorExecutorTest.this.completedCount.get()
                + " endCallEnqueued=" + endCallEnqueued);
          }
        }
        endCallEnqueued = true;
        endCallTransitionCondition.signalAll();
      }
      finally {
        lock.unlock();
      }
    }

    @Override
    public void remove() {
      throw new UnsupportedOperationException();
    }
  }

}