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
5 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