Program Listing for File threadsafe_batching_pool.h

↰ Return to documentation for file (src/translator/threadsafe_batching_pool.h)

/* Thread-safe wrapper around BatchingPool or AggregateBatchingPool, made generic with templates. */
#ifndef SRC_BERGAMOT_THREADSAFE_BATCHING_POOL_H_
#define SRC_BERGAMOT_THREADSAFE_BATCHING_POOL_H_

#include <condition_variable>
#include <mutex>

#include "aggregate_batching_pool.h"
#include "batching_pool.h"
#include "common/options.h"
#include "definitions.h"
#include "translation_model.h"

namespace marian {
namespace bergamot {


template <class BatchingPoolType>
class ThreadsafeBatchingPool {
 public:
  template <class... Args>
  ThreadsafeBatchingPool(Args &&...args);
  ~ThreadsafeBatchingPool();

  template <class... Args>
  void enqueueRequest(Args &&...args);

  template <class... Args>
  size_t generateBatch(Args &&...args);

  // Removes any pending requests from the batching pool.
  void clear();

  // Signals shut down of batching pool. After this no new requests can be enqueued,
  // but all enqueued requests will be processed. To prevent this from happening,
  // call `clear()` before `shutdown()`.
  void shutdown();

 private:
  BatchingPoolType backend_;

  // Number of sentences in backend_;
  size_t enqueued_;

  // Are we shutting down?
  bool shutdown_;

  // Lock on this object.
  std::mutex mutex_;

  // Signaled when there are sentences to translate.
  std::condition_variable work_;
};

}  // namespace bergamot
}  // namespace marian

#define SRC_BERGAMOT_THREADSAFE_BATCHING_POOL_IMPL
#include "threadsafe_batching_pool.cpp"
#undef SRC_BERGAMOT_THREADSAFE_BATCHING_POOL_IMPL

#endif  // SRC_BERGAMOT_THREADSAFE_BATCHING_POOL_H_