Add pipelined & parallel compression optimization (#6262)
Summary: This PR adds support for pipelined & parallel compression optimization for `BlockBasedTableBuilder`. This optimization makes block building, block compression and block appending a pipeline, and uses multiple threads to accelerate block compression. Users can set `CompressionOptions::parallel_threads` greater than 1 to enable compression parallelism. Pull Request resolved: https://github.com/facebook/rocksdb/pull/6262 Reviewed By: ajkr Differential Revision: D20651306 fbshipit-source-id: 62125590a9c15b6d9071def9dc72589c1696a4cbmain
Ziyue Yang
4 years ago
committed by
Facebook GitHub Bot
parent
719c0f91bf
commit
03a781a90c
@ -0,0 +1,149 @@ |
||||
// Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
|
||||
// This source code is licensed under both the GPLv2 (found in the
|
||||
// COPYING file in the root directory) and Apache 2.0 License
|
||||
// (found in the LICENSE.Apache file in the root directory).
|
||||
|
||||
/*
|
||||
* Copyright (c) 2016-present, Facebook, Inc. |
||||
* All rights reserved. |
||||
* |
||||
* This source code is licensed under both the BSD-style license (found in the |
||||
* LICENSE file in the root directory of this source tree) and the GPLv2 (found |
||||
* in the COPYING file in the root directory of this source tree). |
||||
*/ |
||||
#pragma once |
||||
|
||||
#include <atomic> |
||||
#include <cassert> |
||||
#include <condition_variable> |
||||
#include <cstddef> |
||||
#include <cstddef> |
||||
#include <functional> |
||||
#include <mutex> |
||||
#include <queue> |
||||
|
||||
namespace ROCKSDB_NAMESPACE { |
||||
|
||||
/// Unbounded thread-safe work queue.
|
||||
//
|
||||
// This file is an excerpt from Facebook's zstd repo at
|
||||
// https://github.com/facebook/zstd/. The relevant file is
|
||||
// contrib/pzstd/utils/WorkQueue.h.
|
||||
|
||||
template <typename T> |
||||
class WorkQueue { |
||||
// Protects all member variable access
|
||||
std::mutex mutex_; |
||||
std::condition_variable readerCv_; |
||||
std::condition_variable writerCv_; |
||||
std::condition_variable finishCv_; |
||||
|
||||
std::queue<T> queue_; |
||||
bool done_; |
||||
std::size_t maxSize_; |
||||
|
||||
// Must have lock to call this function
|
||||
bool full() const { |
||||
if (maxSize_ == 0) { |
||||
return false; |
||||
} |
||||
return queue_.size() >= maxSize_; |
||||
} |
||||
|
||||
public: |
||||
/**
|
||||
* Constructs an empty work queue with an optional max size. |
||||
* If `maxSize == 0` the queue size is unbounded. |
||||
* |
||||
* @param maxSize The maximum allowed size of the work queue. |
||||
*/ |
||||
WorkQueue(std::size_t maxSize = 0) : done_(false), maxSize_(maxSize) {} |
||||
|
||||
/**
|
||||
* 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. |
||||
* If `push()` returns false, then `item` has not been copied from. |
||||
* |
||||
* @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. |
||||
*/ |
||||
template <typename U> |
||||
bool push(U&& item) { |
||||
{ |
||||
std::unique_lock<std::mutex> lock(mutex_); |
||||
while (full() && !done_) { |
||||
writerCv_.wait(lock); |
||||
} |
||||
if (done_) { |
||||
return false; |
||||
} |
||||
queue_.push(std::forward<U>(item)); |
||||
} |
||||
readerCv_.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_) { |
||||
readerCv_.wait(lock); |
||||
} |
||||
if (queue_.empty()) { |
||||
assert(done_); |
||||
return false; |
||||
} |
||||
item = queue_.front(); |
||||
queue_.pop(); |
||||
} |
||||
writerCv_.notify_one(); |
||||
return true; |
||||
} |
||||
|
||||
/**
|
||||
* Sets the maximum queue size. If `maxSize == 0` then it is unbounded. |
||||
* |
||||
* @param maxSize The new maximum queue size. |
||||
*/ |
||||
void setMaxSize(std::size_t maxSize) { |
||||
{ |
||||
std::lock_guard<std::mutex> lock(mutex_); |
||||
maxSize_ = maxSize; |
||||
} |
||||
writerCv_.notify_all(); |
||||
} |
||||
|
||||
/**
|
||||
* 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; |
||||
} |
||||
readerCv_.notify_all(); |
||||
writerCv_.notify_all(); |
||||
finishCv_.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_) { |
||||
finishCv_.wait(lock); |
||||
} |
||||
} |
||||
}; |
||||
} |
||||
@ -0,0 +1,268 @@ |
||||
// Copyright (c) 2011-present, Facebook, Inc. All rights reserved.
|
||||
// This source code is licensed under both the GPLv2 (found in the
|
||||
// COPYING file in the root directory) and Apache 2.0 License
|
||||
// (found in the LICENSE.Apache file in the root directory).
|
||||
|
||||
/*
|
||||
* Copyright (c) 2016-present, Facebook, Inc. |
||||
* All rights reserved. |
||||
* |
||||
* This source code is licensed under both the BSD-style license (found in the |
||||
* LICENSE file in the root directory of this source tree) and the GPLv2 (found |
||||
* in the COPYING file in the root directory of this source tree). |
||||
*/ |
||||
#include "util/work_queue.h" |
||||
|
||||
#include <gtest/gtest.h> |
||||
#include <iostream> |
||||
#include <memory> |
||||
#include <mutex> |
||||
#include <thread> |
||||
#include <vector> |
||||
|
||||
namespace ROCKSDB_NAMESPACE { |
||||
|
||||
// Unit test for work_queue.h.
|
||||
//
|
||||
// This file is an excerpt from Facebook's zstd repo at
|
||||
// https://github.com/facebook/zstd/. The relevant file is
|
||||
// contrib/pzstd/utils/test/WorkQueueTest.cpp.
|
||||
|
||||
struct Popper { |
||||
WorkQueue<int>* queue; |
||||
int* results; |
||||
std::mutex* mutex; |
||||
|
||||
void operator()() { |
||||
int result; |
||||
while (queue->pop(result)) { |
||||
std::lock_guard<std::mutex> lock(*mutex); |
||||
results[result] = result; |
||||
} |
||||
} |
||||
}; |
||||
|
||||
TEST(WorkQueue, SingleThreaded) { |
||||
WorkQueue<int> queue; |
||||
int result; |
||||
|
||||
queue.push(5); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(5, result); |
||||
|
||||
queue.push(1); |
||||
queue.push(2); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(1, result); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(2, result); |
||||
|
||||
queue.push(1); |
||||
queue.push(2); |
||||
queue.finish(); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(1, result); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(2, result); |
||||
EXPECT_FALSE(queue.pop(result)); |
||||
|
||||
queue.waitUntilFinished(); |
||||
} |
||||
|
||||
TEST(WorkQueue, SPSC) { |
||||
WorkQueue<int> queue; |
||||
const int max = 100; |
||||
|
||||
for (int i = 0; i < 10; ++i) { |
||||
queue.push(i); |
||||
} |
||||
|
||||
std::thread thread([&queue, max] { |
||||
int result; |
||||
for (int i = 0;; ++i) { |
||||
if (!queue.pop(result)) { |
||||
EXPECT_EQ(i, max); |
||||
break; |
||||
} |
||||
EXPECT_EQ(i, result); |
||||
} |
||||
}); |
||||
|
||||
std::this_thread::yield(); |
||||
for (int i = 10; i < max; ++i) { |
||||
queue.push(i); |
||||
} |
||||
queue.finish(); |
||||
|
||||
thread.join(); |
||||
} |
||||
|
||||
TEST(WorkQueue, SPMC) { |
||||
WorkQueue<int> queue; |
||||
std::vector<int> results(50, -1); |
||||
std::mutex mutex; |
||||
std::vector<std::thread> threads; |
||||
for (int i = 0; i < 5; ++i) { |
||||
threads.emplace_back(Popper{&queue, results.data(), &mutex}); |
||||
} |
||||
|
||||
for (int i = 0; i < 50; ++i) { |
||||
queue.push(i); |
||||
} |
||||
queue.finish(); |
||||
|
||||
for (auto& thread : threads) { |
||||
thread.join(); |
||||
} |
||||
|
||||
for (int i = 0; i < 50; ++i) { |
||||
EXPECT_EQ(i, results[i]); |
||||
} |
||||
} |
||||
|
||||
TEST(WorkQueue, MPMC) { |
||||
WorkQueue<int> queue; |
||||
std::vector<int> results(100, -1); |
||||
std::mutex mutex; |
||||
std::vector<std::thread> popperThreads; |
||||
for (int i = 0; i < 4; ++i) { |
||||
popperThreads.emplace_back(Popper{&queue, results.data(), &mutex}); |
||||
} |
||||
|
||||
std::vector<std::thread> pusherThreads; |
||||
for (int i = 0; i < 2; ++i) { |
||||
auto min = i * 50; |
||||
auto max = (i + 1) * 50; |
||||
pusherThreads.emplace_back([&queue, min, max] { |
||||
for (int j = min; j < max; ++j) { |
||||
queue.push(j); |
||||
} |
||||
}); |
||||
} |
||||
|
||||
for (auto& thread : pusherThreads) { |
||||
thread.join(); |
||||
} |
||||
queue.finish(); |
||||
|
||||
for (auto& thread : popperThreads) { |
||||
thread.join(); |
||||
} |
||||
|
||||
for (int i = 0; i < 100; ++i) { |
||||
EXPECT_EQ(i, results[i]); |
||||
} |
||||
} |
||||
|
||||
TEST(WorkQueue, BoundedSizeWorks) { |
||||
WorkQueue<int> queue(1); |
||||
int result; |
||||
queue.push(5); |
||||
queue.pop(result); |
||||
queue.push(5); |
||||
queue.pop(result); |
||||
queue.push(5); |
||||
queue.finish(); |
||||
queue.pop(result); |
||||
EXPECT_EQ(5, result); |
||||
} |
||||
|
||||
TEST(WorkQueue, BoundedSizePushAfterFinish) { |
||||
WorkQueue<int> queue(1); |
||||
int result; |
||||
queue.push(5); |
||||
std::thread pusher([&queue] { queue.push(6); }); |
||||
// Dirtily try and make sure that pusher has run.
|
||||
std::this_thread::sleep_for(std::chrono::seconds(1)); |
||||
queue.finish(); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(5, result); |
||||
EXPECT_FALSE(queue.pop(result)); |
||||
|
||||
pusher.join(); |
||||
} |
||||
|
||||
TEST(WorkQueue, SetMaxSize) { |
||||
WorkQueue<int> queue(2); |
||||
int result; |
||||
queue.push(5); |
||||
queue.push(6); |
||||
queue.setMaxSize(1); |
||||
std::thread pusher([&queue] { queue.push(7); }); |
||||
// Dirtily try and make sure that pusher has run.
|
||||
std::this_thread::sleep_for(std::chrono::seconds(1)); |
||||
queue.finish(); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(5, result); |
||||
EXPECT_TRUE(queue.pop(result)); |
||||
EXPECT_EQ(6, result); |
||||
EXPECT_FALSE(queue.pop(result)); |
||||
|
||||
pusher.join(); |
||||
} |
||||
|
||||
TEST(WorkQueue, BoundedSizeMPMC) { |
||||
WorkQueue<int> queue(10); |
||||
std::vector<int> results(200, -1); |
||||
std::mutex mutex; |
||||
std::cerr << "Creating popperThreads" << std::endl; |
||||
std::vector<std::thread> popperThreads; |
||||
for (int i = 0; i < 4; ++i) { |
||||
popperThreads.emplace_back(Popper{&queue, results.data(), &mutex}); |
||||
} |
||||
|
||||
std::cerr << "Creating pusherThreads" << std::endl; |
||||
std::vector<std::thread> pusherThreads; |
||||
for (int i = 0; i < 2; ++i) { |
||||
auto min = i * 100; |
||||
auto max = (i + 1) * 100; |
||||
pusherThreads.emplace_back([&queue, min, max] { |
||||
for (int j = min; j < max; ++j) { |
||||
queue.push(j); |
||||
} |
||||
}); |
||||
} |
||||
|
||||
std::cerr << "Joining pusherThreads" << std::endl; |
||||
for (auto& thread : pusherThreads) { |
||||
thread.join(); |
||||
} |
||||
std::cerr << "Finishing queue" << std::endl; |
||||
queue.finish(); |
||||
|
||||
std::cerr << "Joining popperThreads" << std::endl; |
||||
for (auto& thread : popperThreads) { |
||||
thread.join(); |
||||
} |
||||
|
||||
std::cerr << "Inspecting results" << std::endl; |
||||
for (int i = 0; i < 200; ++i) { |
||||
EXPECT_EQ(i, results[i]); |
||||
} |
||||
} |
||||
|
||||
TEST(WorkQueue, FailedPush) { |
||||
WorkQueue<int> queue; |
||||
EXPECT_TRUE(queue.push(1)); |
||||
queue.finish(); |
||||
EXPECT_FALSE(queue.push(1)); |
||||
} |
||||
|
||||
TEST(WorkQueue, FailedPop) { |
||||
WorkQueue<int> queue; |
||||
int x = 5; |
||||
EXPECT_TRUE(queue.push(x)); |
||||
queue.finish(); |
||||
x = 0; |
||||
EXPECT_TRUE(queue.pop(x)); |
||||
EXPECT_EQ(5, x); |
||||
EXPECT_FALSE(queue.pop(x)); |
||||
EXPECT_EQ(5, x); |
||||
} |
||||
|
||||
} // namespace ROCKSDB_NAMESPACE
|
||||
|
||||
int main(int argc, char** argv) { |
||||
::testing::InitGoogleTest(&argc, argv); |
||||
return RUN_ALL_TESTS(); |
||||
} |
||||
Loading…
Reference in new issue