zstd/contrib/pzstd/utils/WorkQueue.h

/**
 * Copyright (c) 2016-present, Facebook, Inc.
 * All rights reserved.
 *
 * This source code is licensed under the BSD-style license found in the
 * LICENSE file in the root directory of this source tree. An additional grant
 * of patent rights can be found in the PATENTS file in the same directory.
 */
#pragma once

#include "utils/Buffer.h"

#include <atomic>
#include <cassert>
#include <condition_variable>
#include <cstddef>
#include <functional>
#include <mutex>
#include <queue>

namespace pzstd {

/// Unbounded thread-safe work queue.
template <typename T>
class WorkQueue {
  // Protects all member variable access
  std::mutex mutex_;
  std::condition_variable cv_;

  std::queue<T> queue_;
  bool done_;

 public:
  /// Constructs an empty work queue.
  WorkQueue() : done_(false) {}

  /**
   * Push an item onto the work queue.  Notify a single thread that work is
   * available.  If `finish()` has been called, do nothing and return false.
   *
   * @param item  Item to push onto the queue.
   * @returns     True upon success, false if `finish()` has been called.  An
   *               item was pushed iff `push()` returns true.
   */
  bool push(T item) {
    {
      std::lock_guard<std::mutex> lock(mutex_);
      if (done_) {
        return false;
      }
      queue_.push(std::move(item));
    }
    cv_.notify_one();
    return true;
  }

  /**
   * Attempts to pop an item off the work queue.  It will block until data is
   * available or `finish()` has been called.
   *
   * @param[out] item  If `pop` returns `true`, it contains the popped item.
   *                    If `pop` returns `false`, it is unmodified.
   * @returns          True upon success.  False if the queue is empty and
   *                    `finish()` has been called.
   */
  bool pop(T& item) {
    std::unique_lock<std::mutex> lock(mutex_);
    while (queue_.empty() && !done_) {
      cv_.wait(lock);
    }
    if (queue_.empty()) {
      assert(done_);
      return false;
    }
    item = std::move(queue_.front());
    queue_.pop();
    return true;
  }

  /**
   * Promise that `push()` won't be called again, so once the queue is empty
   * there will never any more work.
   */
  void finish() {
    {
      std::lock_guard<std::mutex> lock(mutex_);
      assert(!done_);
      done_ = true;
    }
    cv_.notify_all();
  }

  /// Blocks until `finish()` has been called (but the queue may not be empty).
  void waitUntilFinished() {
    std::unique_lock<std::mutex> lock(mutex_);
    while (!done_) {
      cv_.wait(lock);
      // If we were woken by a push, we need to wake a thread waiting on pop().
      if (!done_) {
        lock.unlock();
        cv_.notify_one();
        lock.lock();
      }
    }
  }
};

/// Work queue for `Buffer`s that knows the total number of bytes in the queue.
class BufferWorkQueue {
  WorkQueue<Buffer> queue_;
  std::atomic<std::size_t> size_;

 public:
  BufferWorkQueue() : size_(0) {}

  void push(Buffer buffer) {
    size_.fetch_add(buffer.size());
    queue_.push(std::move(buffer));
  }

  bool pop(Buffer& buffer) {
    bool result = queue_.pop(buffer);
    if (result) {
      size_.fetch_sub(buffer.size());
    }
    return result;
  }

  void finish() {
    queue_.finish();
  }

  /**
   * Blocks until `finish()` has been called.
   *
   * @returns The total number of bytes of all the `Buffer`s currently in the
   *           queue.
   */
  std::size_t size() {
    queue_.waitUntilFinished();
    return size_.load();
  }
};
}