producer_consumer_queue.h

// Copyright 2014 Google LLC. All rights reserved.
//
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file or at
// https://developers.google.com/open-source/licenses/bsd
 
#ifndef PACKAGER_MEDIA_BASE_PRODUCER_CONSUMER_QUEUE_H_
#define PACKAGER_MEDIA_BASE_PRODUCER_CONSUMER_QUEUE_H_
 
#include <chrono>
#include <deque>
 
#include <absl/log/check.h>
#include <absl/log/log.h>
#include <absl/strings/str_format.h>
#include <absl/synchronization/mutex.h>
#include <absl/time/time.h>
 
#include <packager/macros/classes.h>
#include <packager/status.h>
 
namespace shaka {
namespace media {
 
static const size_t kUnlimitedCapacity = 0u;
static const int64_t kInfiniteTimeout = -1;
 
template <class T>
class ProducerConsumerQueue {
 public:
  explicit ProducerConsumerQueue(size_t capacity);
 
  ProducerConsumerQueue(size_t capacity, size_t starting_pos);
 
  ~ProducerConsumerQueue();
 
  Status Push(const T& element, int64_t timeout_ms);
 
  Status Pop(T* element, int64_t timeout_ms);
 
  Status Peek(size_t pos, T* element, int64_t timeout_ms);
 
  void Stop() {
    absl::MutexLock lock(&mutex_);
    stop_requested_ = true;
    not_empty_cv_.SignalAll();
    not_full_cv_.SignalAll();
    new_element_cv_.SignalAll();
  }
 
  bool Empty() const {
    absl::MutexLock lock(&mutex_);
    return q_.empty();
  }
 
  size_t Size() const {
    absl::MutexLock lock(&mutex_);
    return q_.size();
  }
 
  size_t HeadPos() const {
    absl::MutexLock lock(&mutex_);
    return head_pos_;
  }
 
  size_t TailPos() const {
    absl::MutexLock lock(&mutex_);
    return head_pos_ + q_.size() - 1;
  }
 
  bool Stopped() const {
    absl::MutexLock lock(&mutex_);
    return stop_requested_;
  }
 
 private:
  // Move head_pos_ to center on pos.
  void SlideHeadOnCenter(size_t pos);
 
  const size_t capacity_;  // Maximum number of elements; zero means unlimited.
 
  mutable absl::Mutex mutex_;
  size_t head_pos_ ABSL_GUARDED_BY(mutex_);  // Head position.
  std::deque<T> q_
      ABSL_GUARDED_BY(mutex_);  // Internal queue holding the elements.
  absl::CondVar not_empty_cv_ ABSL_GUARDED_BY(mutex_);
  absl::CondVar not_full_cv_ ABSL_GUARDED_BY(mutex_);
  absl::CondVar new_element_cv_ ABSL_GUARDED_BY(mutex_);
  bool stop_requested_
      ABSL_GUARDED_BY(mutex_);  // True after Stop has been called.
 
  DISALLOW_COPY_AND_ASSIGN(ProducerConsumerQueue);
};
 
// Implementations of non-inline functions.
template <class T>
ProducerConsumerQueue<T>::ProducerConsumerQueue(size_t capacity)
    : capacity_(capacity),
      head_pos_(0),
      stop_requested_(false) {}
 
template <class T>
ProducerConsumerQueue<T>::ProducerConsumerQueue(size_t capacity,
                                                size_t starting_pos)
    : capacity_(capacity),
      head_pos_(starting_pos),
      stop_requested_(false) {
}
 
template <class T>
ProducerConsumerQueue<T>::~ProducerConsumerQueue() {}
 
template <class T>
Status ProducerConsumerQueue<T>::Push(const T& element, int64_t timeout_ms) {
  absl::MutexLock lock(&mutex_);
  bool woken = false;
 
  // Check for queue shutdown.
  if (stop_requested_)
    return Status(error::STOPPED, "");
 
  auto start = std::chrono::steady_clock::now();
  auto timeout_delta = std::chrono::milliseconds(timeout_ms);
 
  if (capacity_) {
    while (q_.size() == capacity_) {
      if (timeout_ms < 0) {
        // Wait forever, or until Stop.
        not_full_cv_.Wait(&mutex_);
      } else {
        auto elapsed = std::chrono::steady_clock::now() - start;
        if (elapsed < timeout_delta) {
          // Wait with timeout, or until Stop.
          not_full_cv_.WaitWithTimeout(
              &mutex_, absl::FromChrono(timeout_delta - elapsed));
        } else {
          // We're through waiting.
          return Status(error::TIME_OUT, "Time out on pushing.");
        }
      }
      // Re-check for queue shutdown after waking from Wait.
      if (stop_requested_)
        return Status(error::STOPPED, "");
 
      woken = true;
    }
    DCHECK_LT(q_.size(), capacity_);
  }
 
  // Signal consumer to proceed if we are going to create some elements.
  if (q_.empty())
    not_empty_cv_.Signal();
  new_element_cv_.Signal();
 
  q_.push_back(element);
 
  // Signal other producers if we just acquired more capacity.
  if (woken && q_.size() != capacity_)
    not_full_cv_.Signal();
  return Status::OK;
}
 
template <class T>
Status ProducerConsumerQueue<T>::Pop(T* element, int64_t timeout_ms) {
  absl::MutexLock lock(&mutex_);
  bool woken = false;
 
  auto start = std::chrono::steady_clock::now();
  auto timeout_delta = std::chrono::milliseconds(timeout_ms);
 
  while (q_.empty()) {
    if (stop_requested_)
      return Status(error::STOPPED, "");
 
    if (timeout_ms < 0) {
      // Wait forever, or until Stop.
      not_empty_cv_.Wait(&mutex_);
    } else {
      auto elapsed = std::chrono::steady_clock::now() - start;
      if (elapsed < timeout_delta) {
        // Wait with timeout, or until Stop.
        not_empty_cv_.WaitWithTimeout(
            &mutex_, absl::FromChrono(timeout_delta - elapsed));
      } else {
        // We're through waiting.
        return Status(error::TIME_OUT, "Time out on popping.");
      }
    }
    woken = true;
  }
 
  // Signal producer to proceed if we are going to create some capacity.
  if (q_.size() == capacity_)
    not_full_cv_.Signal();
 
  *element = q_.front();
  q_.pop_front();
  ++head_pos_;
 
  // Signal other consumers if we have more elements.
  if (woken && !q_.empty())
    not_empty_cv_.Signal();
  return Status::OK;
}
 
template <class T>
Status ProducerConsumerQueue<T>::Peek(size_t pos,
                                      T* element,
                                      int64_t timeout_ms) {
  absl::MutexLock lock(&mutex_);
  if (pos < head_pos_) {
    return Status(error::INVALID_ARGUMENT,
                  absl::StrFormat("pos (%zu) is too small; head is at %zu.",
                                  pos, head_pos_));
  }
 
  bool woken = false;
 
  auto start = std::chrono::steady_clock::now();
  auto timeout_delta = std::chrono::milliseconds(timeout_ms);
 
  // Move head to create some space (move the sliding window centered @ pos).
  SlideHeadOnCenter(pos);
 
  while (pos >= head_pos_ + q_.size()) {
    if (stop_requested_)
      return Status(error::STOPPED, "");
 
    if (timeout_ms < 0) {
      // Wait forever, or until Stop.
      new_element_cv_.Wait(&mutex_);
    } else {
      auto elapsed = std::chrono::steady_clock::now() - start;
      if (elapsed < timeout_delta) {
        // Wait with timeout, or until Stop.
        new_element_cv_.WaitWithTimeout(
            &mutex_, absl::FromChrono(timeout_delta - elapsed));
      } else {
        // We're through waiting.
        return Status(error::TIME_OUT, "Time out on peeking.");
      }
    }
    // Move head to create some space (move the sliding window centered @ pos).
    SlideHeadOnCenter(pos);
    woken = true;
  }
 
  *element = q_[pos - head_pos_];
 
  // Signal other consumers if we have more elements.
  if (woken && !q_.empty())
    new_element_cv_.Signal();
  return Status::OK;
}
 
template <class T>
void ProducerConsumerQueue<T>::SlideHeadOnCenter(size_t pos) {
  mutex_.AssertHeld();
 
  if (capacity_) {
    // Signal producer to proceed if we are going to create some capacity.
    if (q_.size() == capacity_ && pos > head_pos_ + capacity_ / 2)
      not_full_cv_.Signal();
 
    while (!q_.empty() && pos > head_pos_ + capacity_ / 2) {
      ++head_pos_;
      q_.pop_front();
    }
  }
}
 
}  // namespace media
}  // namespace shaka
 
#endif  // PACKAGER_MEDIA_BASE_PRODUCER_CONSUMER_QUEUE_H_