/*
 * 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.
 */

#ifndef THRIFT_ASYNC_TEVENTJOBQUEUE_H_
#define THRIFT_ASYNC_TEVENTJOBQUEUE_H_ 1

#include <folly/io/async/NotificationQueue.h>
#include <thrift/lib/cpp/Thrift.h>
#include <folly/io/async/EventBase.h>
#include <thrift/lib/cpp/concurrency/Thread.h>
#include <thrift/lib/cpp/concurrency/PosixThreadFactory.h>

namespace apache { namespace thrift { namespace async {

/**
 * Similar to concurrency::Runnable, but run has access to a EventBase to do
 * asynchronous work.
 *
 * There is one queue per worker thread, so work is distributed not
 * when a worker is free, but when a task is enqueued.  This is slightly
 * less fair than concurrency::Runnable, but increases the performance
 * for short jobs by about 3x due to the lack of locking.
 */
class TEventRunnable {
 public:
  TEventRunnable() : eventBase_(nullptr) {}
  virtual ~TEventRunnable() {}

  /**
   * Set the event base on this runnable
   */
  void setEventBase(folly::EventBase *eventBase) {
    eventBase_ = eventBase;
  }

  /**
   * Returns the event base this runnable is currently associated with
   *
   * The TEventJobQueue will set this correctly for run and jobComplete
   */
  folly::EventBase *getEventBase() const {
    return eventBase_;
  }

  /**
   * derived classes must implement run.
   */
  virtual void run() = 0;


 protected:

  folly::EventBase *eventBase_;

};


/**
 * A job queue for working in a EventBase driven application.
 *
 * N threads are spawned and begin consuming a notification queue of
 * TEventRunnable's.
 */
class TEventJobQueue {
 private:

  /**
   * A thread that listens for work on a notification queue, and
   * executes a job
   */
  class JobThread:
      public apache::thrift::concurrency::Runnable,
      public folly::NotificationQueue<TEventRunnable*>::Consumer {

  public:
    explicit JobThread(TEventJobQueue *parent) {}
    ~JobThread() override {}

    folly::EventBase *getEventBase() { return &eventBase_; }

    void join() {
      thread_->join();
    }

    folly::NotificationQueue<TEventRunnable*>* getQueue() {
      return &jobQueue_;
    }

    /**
     * Thread main loop
     */
    void run() override {
      // Listen for new work
      startConsuming(&eventBase_, &jobQueue_);
      try {
        eventBase_.loop();
      } catch (const std::exception &ex) {
        LOG(ERROR) << "Unhandled exception in TEventJobQueue: " <<
          ex.what();
      }
      folly::NotificationQueue<TEventRunnable*>::Consumer::stopConsuming();
    }

    /**
     * Hold a reference so we can join this thread later
     */
    void thread(
      std::shared_ptr<apache::thrift::concurrency::Thread> value) override {

      thread_ = value;
      apache::thrift::concurrency::Runnable::thread(value);
    }

    /**
     * Only overridden because the other thread() is overridden.
     */
    std::shared_ptr<apache::thrift::concurrency::Thread> thread() override {
      return apache::thrift::concurrency::Runnable::thread();
    }

   private:
    folly::EventBase eventBase_;
    std::shared_ptr<apache::thrift::concurrency::Thread> thread_;
    folly::NotificationQueue<TEventRunnable*> jobQueue_;

    /**
     * A new runnable arrived - run it!
     */
    void messageAvailable(TEventRunnable*&& runnable) override {
      runnable->setEventBase(&eventBase_);
      runnable->run();
    }
  };

 public:
  explicit TEventJobQueue(uint32_t numThreads) :
      numThreads_(numThreads),
      curThread_(0),
      shouldJoin_(true) {
  }
  // Constructor where setNumThreads must be called later
  explicit TEventJobQueue() :
      numThreads_(0),
      curThread_(0),
      shouldJoin_(true){}

  // Must be called before init() is called
  void setNumThreads(uint32_t numThreads) {
    numThreads_ = numThreads;
  }

  ~TEventJobQueue() {
    shutdown(true, shouldJoin_); // force
  }


  /**
   * Start the worker threads
   */
  bool init(apache::thrift::concurrency::ThreadFactory* argFactory = nullptr) {
    CHECK(numThreads_ > 0);
    apache::thrift::concurrency::PosixThreadFactory defaultFactory;
    defaultFactory.setDetached(false);
    apache::thrift::concurrency::ThreadFactory* factory = &defaultFactory;
    if (argFactory != nullptr) {
      factory = argFactory;
      shouldJoin_ = false;
    }

    threads_.resize(numThreads_);
    for (uint32_t idx = 0; idx < numThreads_; idx++) {
      threads_[idx] = std::shared_ptr<JobThread>(new JobThread(this));
      std::shared_ptr<apache::thrift::concurrency::Thread> thread =
        factory->newThread(threads_[idx]);
      if (!thread) {
        LOG(CRITICAL) << "Cannot create thread with number: " << idx;
        shutdown(false, shouldJoin_);
        return false;
      }
      thread->start();
    }
    return true;
  }

  /**
   * Instruct each thread to stop listening for new work.
   *
   * When force is true, terminate the consumer thread's event loop
   * even if there is work pending.
   *
   * When join is true, wait for all threads to terminate before
   * leaving shutdown.
   */
  void shutdown(bool force = false, bool join = true) {
    for (auto &thread: threads_) {
      thread->getEventBase()->runInEventBaseThread([thread, force]{
          thread->stopConsuming();
          if (force) {
            thread->getEventBase()->terminateLoopSoon();
          }
        });
    }
    if (join) {
      for (auto &thread: threads_) {
        thread->join();
      }
      threads_.clear();
    }
  }

  /**
   * Add some work to the queue
   *
   * You may call enqueueJob from any thread with no performance penalty
   */
  void enqueueJob(TEventRunnable* runnable) {
    // Note: The logic here must be thread safe with no locking.
    // It isn't perfectly fair: on wrap around, curThread_ may skip
    // some queues, but a mod and set could not be performed atomically
    threads_[(curThread_++) % numThreads_]->getQueue()->putMessage(runnable);
  }

  /**
   * Run the given function in all consumer threads.  The function is passed
   * the event base for the consumer thread.
   */
  void runInAllThreads(const std::function<void(folly::EventBase *)>& fn) {
    for (auto &thread: threads_) {
      thread->getEventBase()->runInEventBaseThread(
        std::bind(fn, thread->getEventBase()));
    }
  }

private:
  uint32_t numThreads_;
  std::atomic<uint64_t> curThread_;
  std::vector<std::shared_ptr<JobThread>> threads_;
  bool shouldJoin_;
};

}}}

#endif
