data_batcher.h

Go to the documentation of this file.

/*
 * Copyright 2019 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file 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.
 */

#pragma once

#include <atomic>
#include <chrono>
#include <iostream>
#include <list>
#include <map>
#include <memory>
#include <mutex>
#include <string>
#include <unordered_map>
#include <stdexcept>

#include <dataflow_lite/utils/service.h>

template<typename T>
class DataBatcher : public Service {
public:

  static const size_t kDefaultTriggerSize = SIZE_MAX;
  static const size_t kDefaultMaxBatchSize = 1024;

  // NOLINTNEXTLINE(google-explicit-constructor, hicpp-explicit-conversions)
  DataBatcher(size_t max_allowable_batch_size = DataBatcher::kDefaultMaxBatchSize,
              size_t trigger_size = DataBatcher::kDefaultTriggerSize,
              std::chrono::microseconds try_enqueue_duration = std::chrono::duration_cast<std::chrono::microseconds>(std::chrono::seconds(2))) {

    validateConfigurableSizes(max_allowable_batch_size, trigger_size);

    this->max_allowable_batch_size_.store(max_allowable_batch_size);
    this->trigger_batch_size_.store(trigger_size);
    this->try_enqueue_duration_ = try_enqueue_duration;

    resetBatchedData();
  }

  ~DataBatcher() override = default;

  virtual bool batchData(const T &data_to_batch) {
    std::lock_guard<std::recursive_mutex> lk(mtx);

    this->batched_data_->push_back(data_to_batch);

    // check if we have exceeded the allowed bounds
    auto allowed_max = getMaxAllowableBatchSize();
    if (getCurrentBatchSize() > allowed_max) {
      emptyCollection();
      return false;
    }

    // publish if the size has been configured
    auto mbs = this->getTriggerBatchSize();
    if (mbs != kDefaultTriggerSize && this->batched_data_->size() >= mbs) {
      publishBatchedData(); // don't return publisher success / fail here
    }

    return true;
  }

  size_t getCurrentBatchSize() {
    std::lock_guard<std::recursive_mutex> lk(mtx);

    return this->batched_data_->size();
  }

  void resetBatchedData() {
    std::lock_guard<std::recursive_mutex> lk(mtx);

    this->batched_data_ = std::make_shared<std::list<T>>();
  }

  void setTriggerBatchSize(size_t new_value) {
    validateConfigurableSizes(this->max_allowable_batch_size_, new_value);

    this->trigger_batch_size_.store(new_value);
  }

  size_t getTriggerBatchSize() {
    return this->trigger_batch_size_.load();

  }

  size_t getMaxAllowableBatchSize() {
    return this->max_allowable_batch_size_.load();
  }

  void setMaxAllowableBatchSize(int new_value) {

    validateConfigurableSizes(new_value, this->trigger_batch_size_);

    this->max_allowable_batch_size_.store(new_value);
  }

  void resetTriggerBatchSize() {
    this->trigger_batch_size_.store(kDefaultTriggerSize);
  }

  void setTryEnqueueDuration(std::chrono::microseconds duration) {
    this->try_enqueue_duration_.store(duration);
  }

  std::chrono::microseconds getTryEnqueueDuration() {
    return this->try_enqueue_duration_.load();
  }

  virtual bool publishBatchedData() = 0;

  static void validateConfigurableSizes(size_t batch_max_queue_size, size_t batch_trigger_publish_size) {

    if (0 == batch_max_queue_size || 0 == batch_trigger_publish_size) {
      throw std::invalid_argument("0 is not a valid size");
    }

    if(kDefaultTriggerSize != batch_trigger_publish_size && batch_trigger_publish_size >= batch_max_queue_size) {
      throw std::invalid_argument("batch_trigger_publish_size must be less than batch_max_queue_size");
    }
  }

  bool shutdown() override {
    bool is_shutdown = Service::shutdown();
    std::lock_guard<std::recursive_mutex> lk(mtx);
    this->emptyCollection();  // attempt to write to disk before discarding
    return is_shutdown;
  }

protected:

  virtual void emptyCollection() {
    std::lock_guard<std::recursive_mutex> lk(mtx);
    this->batched_data_->clear();
  }

  std::shared_ptr<std::list<T>> batched_data_;
  mutable std::recursive_mutex mtx;

private:
    std::atomic<size_t> max_allowable_batch_size_{};
    std::atomic<size_t> trigger_batch_size_{};
    std::atomic<std::chrono::microseconds> try_enqueue_duration_{};
};