Batch merges into one rocksdb write per batch
Instead of one rocksdb write per increment, the increments of the indexing
threads are collected in a WriteBatch and written once per batch. The batch
size can be set with COLLOCATORDB_BATCH_SIZE (default 65536).
The batch is swapped out under the lock and written without it, so the other
threads keep counting while the write is in flight. With the lock held across
the write, an indexer with several threads was slower than with one, because
the threads queued up on the lock and rocksdb never saw concurrent writers
to batch into one.
A read writes out what is still in the batch first, so it sees every
increment that came before it. A cancelled write is repeated as before.
Change-Id: Id534c541a23dfc5b74435c6a88cc001ab6d842b9
diff --git a/src/collocatordb.cc b/src/collocatordb.cc
index 79d038d..f327c4b 100644
--- a/src/collocatordb.cc
+++ b/src/collocatordb.cc
@@ -15,9 +15,11 @@
#include <cstdint>
#include <iostream>
#include <memory>
+#include <mutex>
#include <rocksdb/merge_operator.h>
#include <rocksdb/slice_transform.h>
#include <rocksdb/version.h>
+#include "rocksdb/write_batch.h"
#include <sstream> // for ostringstream
#include <string>
#include <thread>
@@ -298,12 +300,14 @@
uint64_t freq;
};
+static uint64_t env_size(const char *name, uint64_t fallback);
+
class CollocatorDB {
WriteOptions merge_option_; // for merge
// to repeat a write that rocksdb cancelled, rather than lose the count
WriteOptions blocking_merge_option_;
- std::atomic<uint64_t> stalled_writes_{0};
- std::atomic<uint64_t> failed_writes_{0};
+ mutable std::atomic<uint64_t> stalled_writes_{0};
+ mutable std::atomic<uint64_t> failed_writes_{0};
char _one[sizeof(uint64_t)]{};
Slice _one_slice;
vector<VocabEntry> _vocab;
@@ -311,6 +315,55 @@
uint64_t sentences = 0;
float avg_window_size = 8.0;
+ /* Single-key merges are collected in a batch, so that a long indexing run
+ does one rocksdb write per batch instead of one per collocation pair.
+ rocksdb inserts the merge operands of a key in order anyway, which is why
+ the collection point has a mutex: it is the same serialization rocksdb
+ would impose on the writes.
+
+ They are mutable so that the const read methods can write out what is
+ still buffered: a read has to see every increment that came before it. */
+ mutable std::mutex batch_mutex_;
+ mutable WriteBatch batch_;
+ size_t batch_target_ = 65536;
+
+ /* Writes out a batch of merge operands, retrying a cancelled write the same
+ way merge_one() used to, before the merges were batched.
+
+ Takes a batch that merge_one() or flush() took out of the shared batch,
+ and is called without batch_mutex_: the write is the slow part, and the
+ other threads must be able to keep accumulating while it is in flight.
+ Holding the mutex across the write made the indexer run slower with many
+ threads than with one, because the threads queued up on the mutex and
+ rocksdb never saw concurrent writers to batch into one. */
+ void write_batch(WriteBatch &to_write) const {
+ if (to_write.Count() == 0)
+ return;
+ Status s = db_->Write(merge_option_, &to_write);
+ if (s.ok())
+ return;
+ if (s.IsIncomplete()) {
+ ++stalled_writes_;
+ s = db_->Write(blocking_merge_option_, &to_write);
+ if (s.ok())
+ return;
+ }
+ if (failed_writes_++ == 0)
+ std::cerr << "collocatordb: cannot write, counts are lost: "
+ << s.ToString() << std::endl;
+ }
+
+ /* A read has to see the increments that are still in the batch, so the read
+ methods call this first. */
+ void flush() const {
+ WriteBatch to_write;
+ {
+ std::lock_guard<std::mutex> lock(batch_mutex_);
+ std::swap(batch_, to_write);
+ }
+ write_batch(to_write);
+ }
+
protected:
std::shared_ptr<DB> db_;
@@ -373,6 +426,7 @@
// mapped to a rocksdb Get
bool get(const std::string &key, uint64_t *value) {
+ flush();
std::string str;
auto s = db_->Get(get_option_, key, &str);
@@ -407,18 +461,14 @@
and then the increment is simply gone. Such a write was not applied, so
repeating it blocking cannot count twice. */
void merge_one(const Slice &key) {
- Status s = db_->Merge(merge_option_, key, _one_slice);
- if (s.ok())
- return;
- if (s.IsIncomplete()) {
- ++stalled_writes_;
- s = db_->Merge(blocking_merge_option_, key, _one_slice);
- if (s.ok())
- return;
+ WriteBatch to_write;
+ {
+ std::lock_guard<std::mutex> lock(batch_mutex_);
+ batch_.Merge(key, _one_slice);
+ if (batch_.Count() >= (int)batch_target_)
+ std::swap(batch_, to_write);
}
- if (failed_writes_++ == 0)
- std::cerr << "collocatordb: cannot write, counts are lost: " << s.ToString()
- << std::endl;
+ write_batch(to_write);
}
virtual void inc(const std::string &key) {
@@ -476,6 +526,7 @@
void close() {
if (!db_)
return;
+ flush();
if (stalled_writes_ > 0)
std::cerr << "collocatordb: repeated " << stalled_writes_
<< " writes that rocksdb had cancelled" << std::endl;
@@ -497,6 +548,7 @@
CollocatorDB::CollocatorDB(const char *db_name,
bool read_only = false) {
+ batch_target_ = (size_t)env_size("COLLOCATORDB_BATCH_SIZE", 65536);
// merge_option_.sync = true;
if (read_only)
db_ = OpenDbForRead(strdup(db_name));
@@ -666,6 +718,7 @@
CollocatorIterator *
CollocatorDB::SeekIterator(uint64_t w1, uint64_t w2, int8_t dist) const {
+ flush();
ReadOptions options;
options.prefix_same_as_start = true;
char prefixc[sizeof(uint64_t)];