/**
 * 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 org.apache.ambari.server.controller.utilities;

import java.util.Queue;
import java.util.concurrent.ExecutorCompletionService;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;

/**
 * An {@link ExecutorCompletionService} which takes a {@link ThreadPoolExecutor}
 * and uses its thread pool to execute tasks - buffering any tasks that
 * overflow. Such buffered tasks are later re-submitted to the executor when
 * finished tasks are polled or taken.
 * <p/>
 * This class overrides the {@link ThreadPoolExecutor}'s
 * {@link RejectedExecutionHandler} to collect overflowing tasks. When
 * {@link #poll(long, TimeUnit)} is invoked, this class will attempt to
 * re-submit any overflow tasks before waiting the specified amount of time.
 * This will prevent blocking on an empty completion queue since the parent
 * {@link ExecutorCompletionService} doesn't have any idea that there are tasks
 * waiting to be resubmitted.
 * <p/>
 * The {@link ScalingThreadPoolExecutor} can be used in conjunction with this
 * class to provide an efficient buffered, scaling thread pool implementation.
 *
 * @param <V>
 */
public class BufferedThreadPoolExecutorCompletionService<V> extends ExecutorCompletionService<V> {

  private ThreadPoolExecutor executor;
  private Queue<Runnable> overflowQueue;

  public BufferedThreadPoolExecutorCompletionService(ThreadPoolExecutor executor) {
    super(executor);
    this.executor = executor;
    this.overflowQueue = new LinkedBlockingQueue<>();
    this.executor.setRejectedExecutionHandler(new RejectedExecutionHandler() {
      /**
       * Once the ThreadPoolExecutor is at full capacity, it starts to reject
       * submissions which are queued for later submission.
       */
      @Override
      public void rejectedExecution(Runnable r, ThreadPoolExecutor executor) {
        overflowQueue.add(r);
      }
    });
  }

  @Override
  public Future<V> take() throws InterruptedException {
    Future<V> take = super.take();
    if (!executor.isTerminating() && !overflowQueue.isEmpty() && executor.getActiveCount() < executor.getMaximumPoolSize()) {
      Runnable overflow = overflowQueue.poll();
      if (overflow != null) {
        executor.execute(overflow);
      }
    }
    return take;
  }

  @Override
  public Future<V> poll() {
    Future<V> poll = super.poll();
    if (!executor.isTerminating() && !overflowQueue.isEmpty() && executor.getActiveCount() < executor.getMaximumPoolSize()) {
      Runnable overflow = overflowQueue.poll();
      if (overflow != null) {
        executor.execute(overflow);
      }
    }
    return poll;
  }

  /**
   * {@inheritDoc}
   * <p/>
   * The goal of this method is to prevent blocking if there are tasks waiting
   * in the overflow queue. In the event that the overflow queue was populated,
   * we should not blindly wait on the parent
   * {@link ExecutorCompletionService#poll(long, TimeUnit)} method. Instead, we
   * should ensure that we have submitted at least one of our own tasks for
   * completion.
   */
  @Override
  public Future<V> poll(long timeout, TimeUnit unit) throws InterruptedException {
    // first poll for anything that's already completed and do a short-circuit return
    Future<V> poll = super.poll();
    if( null != poll ) {
      // there's something to return; that's great, but let's also see if we can
      // submit anything in the overflow queue back to the completion service
      if (!executor.isTerminating() && !overflowQueue.isEmpty()
          && executor.getActiveCount() < executor.getMaximumPoolSize()) {
        Runnable overflow = overflowQueue.poll();
        if (overflow != null) {
          executor.execute(overflow);
        }
      }

      // return the future
      return poll;
    }

    // nothing completed yet, so check active thread count - if there is nothing
    // working either, then that means that the parent completion service thinks
    // it's done - we should submit our own tasks
    if (executor.getActiveCount() == 0) {
      if (!executor.isTerminating() && !overflowQueue.isEmpty()) {
        Runnable overflow = overflowQueue.poll();
        if (overflow != null) {
          executor.execute(overflow);
        }
      }
    }

    // now that we've confirmed that either the parent completion service is
    // still working or we submitted our own task, we can poll with a timeout
    poll = super.poll(timeout, unit);
    return poll;
  }
}