30058987329

Committed 24 Jul 2026 01:26AM UTC coverage: 52.76% (+0.1%) from 52.66%

Build # 30058987329

Build Type

Pull #99

github

Committed by

web-flow

Commit Message

Merge d78ce6d20 into 06bc84ec9

Pull Request Pull Request #99: Support CM time_metric (NS/MS/SEC/US) across reader and viz

Coverage Stats

41934 of 102825 branches covered (40.78%)

Branch coverage included in aggregate %.

2278 of 2903 new or added lines in 61 files covered. (78.47%)

144 existing lines in 14 files now uncovered.

37072 of 46920 relevant lines covered (79.01%)

59158.32 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

62.67

/src/dftracer/utils/core/rocksdb/database.cpp

#include <dftracer/utils/core/common/constants.h>
#include <dftracer/utils/core/common/error.h>
#include <dftracer/utils/core/common/filesystem.h>
#include <dftracer/utils/core/env.h>
#include <dftracer/utils/core/rocksdb/database.h>
#include <dftracer/utils/core/rocksdb/filesystem.h>
#include <rocksdb/filter_policy.h>
#include <rocksdb/slice.h>
#include <rocksdb/table.h>
#include <rocksdb/write_buffer_manager.h>

#include <algorithm>
#include <atomic>
#include <cstdlib>
#include <stdexcept>
#include <utility>

namespace dftracer::utils::rocksdb {

namespace {

std::atomic<bool>& process_exiting_flag() {
    static std::atomic<bool> flag{false};
    return flag;
}

const ::rocksdb::ReadOptions& read_options() {
    static const auto options = [] {
        ::rocksdb::ReadOptions ro;
        // Prefetch data blocks asynchronously ahead of the scan cursor (our FS
        // implements ReadAsync); adaptive_readahead grows the window as the
        // scan runs so short range reads do not over-read.
        ro.async_io = true;
        ro.adaptive_readahead = true;
        return ro;
    }();
    return options;
}

const ::rocksdb::WriteOptions& write_options() {
    static const auto options = [] {
        ::rocksdb::WriteOptions wo;
        wo.disableWAL = true;
        return wo;
    }();
    return options;
}

void cleanup_failed_open(::rocksdb::DB*& db,
                         std::vector<::rocksdb::ColumnFamilyHandle*>& handles) {
    if (db != nullptr) {
        for (auto* handle : handles) {
            if (handle != nullptr) {
                db->DestroyColumnFamilyHandle(handle);
            }
        }
        static_cast<void>(db->Close());
        delete db;
        db = nullptr;
    }
    handles.clear();
}

}  // namespace

void mark_process_exiting_for_rocksdb() {
    process_exiting_flag().store(true, std::memory_order_relaxed);
}

RocksDatabase::RocksDatabase() = default;

RocksDatabase::RocksDatabase(const std::string& db_path, OpenMode open_mode) {
    open(db_path, open_mode);
}

RocksDatabase::~RocksDatabase() { close(); }

RocksDatabase::RocksDatabase(RocksDatabase&& other) noexcept
    : db_path_(std::move(other.db_path_)),
      open_mode_(other.open_mode_),
      file_system_(std::move(other.file_system_)),
      env_(std::move(other.env_)),
      db_(std::exchange(other.db_, nullptr)),
      column_families_(std::move(other.column_families_)) {}

RocksDatabase& RocksDatabase::operator=(RocksDatabase&& other) noexcept {
    if (this != &other) {
        close();
        db_path_ = std::move(other.db_path_);
        open_mode_ = other.open_mode_;
        file_system_ = std::move(other.file_system_);
        env_ = std::move(other.env_);
        db_ = std::exchange(other.db_, nullptr);
        column_families_ = std::move(other.column_families_);
    }
    return *this;
}

const decltype(cf::ALL)& RocksDatabase::default_column_families() {
    return cf::ALL;
}

namespace {
/// One cache for the process. Left unset, every column family gets its own
/// 8 MB default, which no realistic working set fits in.
std::shared_ptr<::rocksdb::Cache>& shared_block_cache() {
    static std::shared_ptr<::rocksdb::Cache> cache =
        ::rocksdb::NewLRUCache(constants::rocksdb::BLOCK_CACHE_BYTES);
    return cache;
}

/// Caps total memtable memory across every column family and every DB in the
/// process; without it each CF reserves its own write buffers.
std::shared_ptr<::rocksdb::WriteBufferManager>& shared_write_buffer_manager() {
    static std::shared_ptr<::rocksdb::WriteBufferManager> mgr =
        std::make_shared<::rocksdb::WriteBufferManager>(
            constants::rocksdb::WRITE_BUFFER_BYTES, shared_block_cache());
    return mgr;
}
}  // namespace

::rocksdb::Options RocksDatabase::default_options() {
    ::rocksdb::Options options;
    options.create_if_missing = true;
    options.create_missing_column_families = true;
    options.allow_concurrent_memtable_write = true;
    options.enable_pipelined_write = true;
    options.max_open_files = Env::rocksdb_max_open_files();
    // Skip the synchronous per-SST stats pass on open; expensive over Lustre.
    // (open_files_async is not usable here: it background-opens SSTs by name
    // and races our ingest/compaction file churn, hitting ENOENT.)
    options.skip_stats_update_on_db_open = true;
    options.max_background_jobs = 8;
    options.max_subcompactions = 8;
    options.write_buffer_size = 256 * 1024 * 1024;
    options.max_write_buffer_number = 4;
    options.write_buffer_manager = shared_write_buffer_manager();
    return options;
}

::rocksdb::ColumnFamilyOptions RocksDatabase::default_column_family_options() {
    ::rocksdb::ColumnFamilyOptions options;

    ::rocksdb::BlockBasedTableOptions bbt;
    bbt.block_size = 32 * 1024;
    bbt.format_version = 7;
    bbt.index_block_restart_interval = 16;
    bbt.block_cache = shared_block_cache();
    // Better compression and read CPU for the scanned CFs (fixed-layout blobs).
    bbt.separate_key_value_in_data_block = true;
    options.table_factory.reset(::rocksdb::NewBlockBasedTableFactory(bbt));

#ifdef DFTRACER_UTILS_ENABLE_ZSTD
    options.compression = ::rocksdb::kZSTD;
    options.compression_opts.level = constants::rocksdb::ZSTD_COMPRESSION_LEVEL;
    options.compression_opts.max_dict_bytes =
        constants::rocksdb::ZSTD_MAX_DICT_BYTES;
    options.compression_opts.zstd_max_train_bytes =
        constants::rocksdb::ZSTD_MAX_TRAIN_BYTES;
    options.compression_opts.enabled = true;
    options.bottommost_compression = ::rocksdb::kZSTD;
    options.bottommost_compression_opts.level =
        constants::rocksdb::ZSTD_COMPRESSION_LEVEL;
    options.bottommost_compression_opts.max_dict_bytes =
        constants::rocksdb::ZSTD_MAX_DICT_BYTES;
    options.bottommost_compression_opts.zstd_max_train_bytes =
        constants::rocksdb::ZSTD_MAX_TRAIN_BYTES;
    options.bottommost_compression_opts.enabled = true;
#elif defined(DFTRACER_UTILS_ENABLE_LZ4)
    options.compression = ::rocksdb::kLZ4Compression;
    options.bottommost_compression = ::rocksdb::kZlibCompression;
#else
    options.compression = ::rocksdb::kZlibCompression;
    options.bottommost_compression = ::rocksdb::kZlibCompression;
#endif
    return options;
}

::rocksdb::ColumnFamilyOptions
RocksDatabase::point_lookup_column_family_options() {
    auto options = default_column_family_options();

    // Hash keys are uniformly distributed, so every SST's range covers almost
    // every key: without a filter a get searches all of them. Large blocks
    // also mean reading and decompressing 32 KB to return a few dozen bytes.
    ::rocksdb::BlockBasedTableOptions bbt;
    bbt.block_size = constants::rocksdb::POINT_LOOKUP_BLOCK_SIZE;
    bbt.format_version = 7;
    bbt.index_block_restart_interval = 16;
    bbt.block_cache = shared_block_cache();
    // Content-hash keys are uniformly distributed: interpolation search beats
    // binary. Read-time only, works on any SST.
    bbt.index_block_search_type =
        ::rocksdb::BlockBasedTableOptions::kInterpolation;
    bbt.filter_policy.reset(::rocksdb::NewBloomFilterPolicy(
        constants::rocksdb::BLOOM_BITS_PER_KEY, false));
    // Left in the table reader rather than the block cache: a scan streams
    // enough data blocks through the cache to evict the filters it needs.
    bbt.cache_index_and_filter_blocks = false;
    options.table_factory.reset(::rocksdb::NewBlockBasedTableFactory(bbt));
    return options;
}

bool RocksDatabase::is_point_lookup_cf(std::string_view name) noexcept {
    return name == cf::HASH_TABLES || name == cf::NAME_DICTIONARY;
}

bool RocksDatabase::open(const std::string& db_path, OpenMode open_mode) {
    close();
    db_path_ = db_path;
    open_mode_ = open_mode;

    std::error_code ec;
    if (open_mode_ == OpenMode::ReadWrite) {
        fs::create_directories(fs::path(db_path_), ec);
    }

    auto db_options = default_options();
    if (open_mode_ == OpenMode::ReadOnly) {
        db_options.create_if_missing = false;
        db_options.create_missing_column_families = false;
        // Read paths do heavy point lookups (e.g. dfanalyzer hash resolution);
        // with a small table cache the SST readers get evicted and every lookup
        // re-opens the file to re-read its filter/index blocks, which thrashes.
        // There is no write-side fd pressure here, so keep every SST open
        // unless an explicit env cap is set.
        if (!Env::get<int>("DFTRACER_UTILS_ROCKSDB_MAX_OPEN_FILES")
                 .has_value()) {
            db_options.max_open_files = -1;
        }
    }
    file_system_ = make_dftracer_file_system();
    env_ = make_dftracer_env(file_system_);
    db_options.env = env_.get();
    auto cf_options = default_column_family_options();

    std::vector<std::string> column_family_names;
    auto list_status = ::rocksdb::DB::ListColumnFamilies(db_options, db_path_,
                                                         &column_family_names);
    if (!list_status.ok()) {
        if (open_mode_ == OpenMode::ReadOnly) {
            throw DFTUtilsException(
                ErrorCode::IO, "Failed to list RocksDB column families at '" +
                                   db_path_ + "': " + list_status.ToString());
        }
        column_family_names.reserve(default_column_families().size());
        for (auto name : default_column_families()) {
            column_family_names.emplace_back(name);
        }
    } else {
        if (open_mode_ == OpenMode::ReadWrite) {
            for (const auto& name : default_column_families()) {
                if (std::find(column_family_names.begin(),
                              column_family_names.end(),
                              name) == column_family_names.end()) {
                    column_family_names.emplace_back(name);
                }
            }
        }
    }

    std::vector<::rocksdb::ColumnFamilyDescriptor> descriptors;
    descriptors.reserve(column_family_names.size());
    for (const auto& name : column_family_names) {
        auto opts = is_point_lookup_cf(name)
                        ? point_lookup_column_family_options()
                        : cf_options;
        if (cf_options_override_) {
            cf_options_override_(name, opts);
        }
        descriptors.emplace_back(name, opts);
    }

    std::vector<::rocksdb::ColumnFamilyHandle*> handles;
    ::rocksdb::Status status;
    std::unique_ptr<::rocksdb::DB> opened;
    if (open_mode_ == OpenMode::ReadOnly) {
        status = ::rocksdb::DB::OpenForReadOnly(
            db_options, db_path_, descriptors, &handles, &opened, false);
    } else {
        status = ::rocksdb::DB::Open(db_options, db_path_, descriptors,
                                     &handles, &opened);
    }
    if (!status.ok()) {
        ::rocksdb::DB* raw = opened.release();
        cleanup_failed_open(raw, handles);
        throw DFTUtilsException(ErrorCode::IO, "Failed to open RocksDB at '" +
                                                   db_path_ +
                                                   "': " + status.ToString());
    }
    db_ = opened.release();

    column_families_.clear();
    for (std::size_t i = 0; i < descriptors.size(); ++i) {
        column_families_.emplace(descriptors[i].name, handles[i]);
    }

    return true;
}

void RocksDatabase::close() {
    if (db_ == nullptr) {
        column_families_.clear();
        return;
    }

    if (process_exiting_flag().load(std::memory_order_relaxed)) {
        db_ = nullptr;
        column_families_.clear();
        env_.reset();
        file_system_.reset();
        db_path_.clear();
        return;
    }

    for (auto& entry : column_families_) {
        if (entry.second != nullptr) {
            db_->DestroyColumnFamilyHandle(entry.second);
            entry.second = nullptr;
        }
    }
    column_families_.clear();

    auto* db = db_;
    db_ = nullptr;
    static_cast<void>(db->Close());
    delete db;
    env_.reset();
    file_system_.reset();
    db_path_.clear();
}

bool RocksDatabase::is_open() const noexcept { return db_ != nullptr; }

bool RocksDatabase::is_read_only() const noexcept {
    return open_mode_ == OpenMode::ReadOnly;
}

const std::string& RocksDatabase::path() const noexcept { return db_path_; }

::rocksdb::DB* RocksDatabase::get() const noexcept { return db_; }

::rocksdb::ColumnFamilyHandle* RocksDatabase::column_family_handle(
    std::string_view column_family) const {
    const auto name = column_family.empty() ? std::string(cf::DEFAULT)
                                            : std::string(column_family);
    const auto it = column_families_.find(name);
    if (it == column_families_.end() || it->second == nullptr) {
        throw DFTUtilsException(ErrorCode::INVALID_ARGUMENT,
                                "Unknown RocksDB column family: " + name);
    }
    return it->second;
}

::rocksdb::Status RocksDatabase::put(std::string_view key,
                                     std::string_view value,
                                     std::string_view column_family) {
    return db_->Put(write_options(), column_family_handle(column_family),
                    ::rocksdb::Slice(key.data(), key.size()),
                    ::rocksdb::Slice(value.data(), value.size()));
}

::rocksdb::Status RocksDatabase::get(std::string_view key, std::string* value,
                                     std::string_view column_family) const {
    return db_->Get(read_options(), column_family_handle(column_family),
                    ::rocksdb::Slice(key.data(), key.size()), value);
}

::rocksdb::Status RocksDatabase::merge(std::string_view key,
                                       std::string_view value,
                                       std::string_view column_family) {
    return db_->Merge(write_options(), column_family_handle(column_family),
                      ::rocksdb::Slice(key.data(), key.size()),
                      ::rocksdb::Slice(value.data(), value.size()));
}

void RocksDatabase::set_cf_options_override(CfOptionsOverride override) {
    cf_options_override_ = std::move(override);
}

::rocksdb::Status RocksDatabase::merge(Batch& batch,
                                       std::string_view column_family,
                                       std::string_view key,
                                       std::string_view value) {
    return batch.Merge(column_family_handle(column_family),
                       ::rocksdb::Slice(key.data(), key.size()),
                       ::rocksdb::Slice(value.data(), value.size()));
}

::rocksdb::Status RocksDatabase::del(std::string_view key,
                                     std::string_view column_family) {
    return db_->Delete(write_options(), column_family_handle(column_family),
                       ::rocksdb::Slice(key.data(), key.size()));
}

::rocksdb::Status RocksDatabase::delete_range(std::string_view begin_key,
                                              std::string_view end_key,
                                              std::string_view column_family) {
    return db_->DeleteRange(
        write_options(), column_family_handle(column_family),
        ::rocksdb::Slice(begin_key.data(), begin_key.size()),
        ::rocksdb::Slice(end_key.data(), end_key.size()));
}

::rocksdb::Status RocksDatabase::put(Batch& batch,
                                     std::string_view column_family,
                                     std::string_view key,
                                     std::string_view value) {
    return batch.Put(column_family_handle(column_family),
                     ::rocksdb::Slice(key.data(), key.size()),
                     ::rocksdb::Slice(value.data(), value.size()));
}

::rocksdb::Status RocksDatabase::del(Batch& batch,
                                     std::string_view column_family,
                                     std::string_view key) {
    return batch.Delete(column_family_handle(column_family),
                        ::rocksdb::Slice(key.data(), key.size()));
}

RocksDatabase::Batch RocksDatabase::begin_batch() const { return Batch(); }

::rocksdb::Status RocksDatabase::commit_batch(Batch& batch) {
    return db_->Write(write_options(), &batch);
}

std::unique_ptr<::rocksdb::Iterator> RocksDatabase::new_iterator(
    std::string_view column_family) const {
    return std::unique_ptr<::rocksdb::Iterator>(
        db_->NewIterator(read_options(), column_family_handle(column_family)));
}

::rocksdb::Status RocksDatabase::compact(std::string_view column_family) {
    ::rocksdb::CompactRangeOptions opts;
    opts.max_subcompactions = 8;
    return db_->CompactRange(opts, column_family_handle(column_family), nullptr,
                             nullptr);
}

::rocksdb::Status RocksDatabase::ingest_external_files(
    std::string_view column_family,
    const std::vector<std::string>& external_files, bool ingest_behind) {
    if (external_files.empty()) {
        return ::rocksdb::Status::OK();
    }
    ::rocksdb::IngestExternalFileOptions opts;
    // Rename same-FS staged SSTs instead of copying; RocksDB copies on failure.
    opts.move_files = true;
    // Offline bulk build with no live readers, so skip the snapshot/seqno sync.
    opts.snapshot_consistency = false;
    opts.allow_global_seqno = true;
    // Keep the assigned seqno in the manifest instead of rewriting it into each
    // SST (format_version 5 supports this), avoiding a per-SST write + fsync.
    opts.write_global_seqno = false;
    opts.allow_blocking_flush = true;
    opts.ingest_behind = ingest_behind;
    return db_->IngestExternalFile(column_family_handle(column_family),
                                   external_files, opts);
}

::rocksdb::Status RocksDatabase::ingest_external_files_multi(
    const std::vector<
        std::pair<std::string_view, const std::vector<std::string>*>>& per_cf) {
    ::rocksdb::IngestExternalFileOptions opts;
    opts.move_files = true;
    opts.snapshot_consistency = false;
    opts.allow_global_seqno = true;
    opts.write_global_seqno = false;
    opts.allow_blocking_flush = true;

    std::vector<::rocksdb::IngestExternalFileArg> args;
    args.reserve(per_cf.size());
    for (const auto& [cf_name, files] : per_cf) {
        if (!files || files->empty()) continue;
        ::rocksdb::IngestExternalFileArg arg;
        arg.column_family = column_family_handle(cf_name);
        arg.external_files = *files;
        arg.options = opts;
        args.push_back(std::move(arg));
    }
    if (args.empty()) return ::rocksdb::Status::OK();
    return db_->IngestExternalFiles(args);
}

}  // namespace dftracer::utils::rocksdb

1	#include <dftracer/utils/core/common/constants.h>
2	#include <dftracer/utils/core/common/error.h>
3	#include <dftracer/utils/core/common/filesystem.h>
4	#include <dftracer/utils/core/env.h>
5	#include <dftracer/utils/core/rocksdb/database.h>
6	#include <dftracer/utils/core/rocksdb/filesystem.h>
7	#include <rocksdb/filter_policy.h>
8	#include <rocksdb/slice.h>
9	#include <rocksdb/table.h>
10	#include <rocksdb/write_buffer_manager.h>
11
12	#include <algorithm>
13	#include <atomic>
14	#include <cstdlib>
15	#include <stdexcept>
16	#include <utility>
17
18	namespace dftracer::utils::rocksdb {
19
20	namespace {
21
22	std::atomic<bool>& process_exiting_flag() {	9,156✔
23	static std::atomic<bool> flag{false};
24	return flag;	9,156✔
25	}
26
27	const ::rocksdb::ReadOptions& read_options() {	26,363✔
28	static const auto options = [] {	13,411!
29	::rocksdb::ReadOptions ro;	230✔
30	// Prefetch data blocks asynchronously ahead of the scan cursor (our FS
31	// implements ReadAsync); adaptive_readahead grows the window as the
32	// scan runs so short range reads do not over-read.
33	ro.async_io = true;	230✔
34	ro.adaptive_readahead = true;	230✔
35	return ro;	230✔
36	}();	13,297!
37	return options;	26,365✔
38	}
39
40	const ::rocksdb::WriteOptions& write_options() {	4,489✔
41	static const auto options = [] {	2,472!
42	::rocksdb::WriteOptions wo;	224✔
43	wo.disableWAL = true;	224✔
44	return wo;	224✔
45	}();	2,241!
46	return options;	4,489✔
47	}
48
49	void cleanup_failed_open(::rocksdb::DB*& db,	×
50	std::vector<::rocksdb::ColumnFamilyHandle*>& handles) {
51	if (db != nullptr) {	×
52	for (auto* handle : handles) {	×
53	if (handle != nullptr) {	×
54	db->DestroyColumnFamilyHandle(handle);	×
55	}
56	}
57	static_cast<void>(db->Close());	×
58	delete db;	×
59	db = nullptr;	×
60	}
61	handles.clear();	×
62	}	×
63
64	} // namespace
65
66	void mark_process_exiting_for_rocksdb() {	12✔
67	process_exiting_flag().store(true, std::memory_order_relaxed);	12✔
68	}	12✔
69
70	RocksDatabase::RocksDatabase() = default;	23,083✔
71
72	RocksDatabase::RocksDatabase(const std::string& db_path, OpenMode open_mode) {	15✔
73	open(db_path, open_mode);	6!
74	}	9✔
75
76	RocksDatabase::~RocksDatabase() { close(); }	13,825!
77
78	RocksDatabase::RocksDatabase(RocksDatabase&& other) noexcept	×
79	: db_path_(std::move(other.db_path_)),	×
80	open_mode_(other.open_mode_),	×
81	file_system_(std::move(other.file_system_)),	×
82	env_(std::move(other.env_)),	×
83	db_(std::exchange(other.db_, nullptr)),	×
84	column_families_(std::move(other.column_families_)) {}	×
85
86	RocksDatabase& RocksDatabase::operator=(RocksDatabase&& other) noexcept {	×
87	if (this != &other) {	×
88	close();	×
89	db_path_ = std::move(other.db_path_);	×
90	open_mode_ = other.open_mode_;	×
91	file_system_ = std::move(other.file_system_);	×
92	env_ = std::move(other.env_);	×
93	db_ = std::exchange(other.db_, nullptr);	×
94	column_families_ = std::move(other.column_families_);	×
95	}
96	return *this;	×
97	}
98
99	const decltype(cf::ALL)& RocksDatabase::default_column_families() {	2,994✔
100	return cf::ALL;	2,994✔
101	}
102
103	namespace {
104	/// One cache for the process. Left unset, every column family gets its own
105	/// 8 MB default, which no realistic working set fits in.
106	std::shared_ptr<::rocksdb::Cache>& shared_block_cache() {	46,466✔
107	static std::shared_ptr<::rocksdb::Cache> cache =	23,563!
108	::rocksdb::NewLRUCache(constants::rocksdb::BLOCK_CACHE_BYTES);	23,133!
109	return cache;	46,466✔
110	}
111
112	/// Caps total memtable memory across every column family and every DB in the
113	/// process; without it each CF reserves its own write buffers.
114	std::shared_ptr<::rocksdb::WriteBufferManager>& shared_write_buffer_manager() {	9,659✔
115	static std::shared_ptr<::rocksdb::WriteBufferManager> mgr =	4,987!
116	std::make_shared<::rocksdb::WriteBufferManager>(	115!
117	constants::rocksdb::WRITE_BUFFER_BYTES, shared_block_cache());	4,902!
118	return mgr;	9,659✔
119	}
120	} // namespace
121
122	::rocksdb::Options RocksDatabase::default_options() {	9,659✔
123	::rocksdb::Options options;	9,659✔
124	options.create_if_missing = true;	9,659✔
125	options.create_missing_column_families = true;	9,659✔
126	options.allow_concurrent_memtable_write = true;	9,659✔
127	options.enable_pipelined_write = true;	9,659✔
128	options.max_open_files = Env::rocksdb_max_open_files();	9,659!
129	// Skip the synchronous per-SST stats pass on open; expensive over Lustre.
130	// (open_files_async is not usable here: it background-opens SSTs by name
131	// and races our ingest/compaction file churn, hitting ENOENT.)
132	options.skip_stats_update_on_db_open = true;	9,659✔
133	options.max_background_jobs = 8;	9,659✔
134	options.max_subcompactions = 8;	9,659✔
135	options.write_buffer_size = 256 * 1024 * 1024;	9,659✔
136	options.max_write_buffer_number = 4;	9,659✔
137	options.write_buffer_manager = shared_write_buffer_manager();	9,659!
138	return options;	9,659✔
139	}	4,872!
140
141	::rocksdb::ColumnFamilyOptions RocksDatabase::default_column_family_options() {	27,948✔
142	::rocksdb::ColumnFamilyOptions options;	27,948!
143
144	::rocksdb::BlockBasedTableOptions bbt;	27,948✔
145	bbt.block_size = 32 * 1024;	27,948✔
146	bbt.format_version = 7;	27,948✔
147	bbt.index_block_restart_interval = 16;	27,948✔
148	bbt.block_cache = shared_block_cache();	27,948!
149	// Better compression and read CPU for the scanned CFs (fixed-layout blobs).
150	bbt.separate_key_value_in_data_block = true;	27,948✔
151	options.table_factory.reset(::rocksdb::NewBlockBasedTableFactory(bbt));	27,948!
152
153	#ifdef DFTRACER_UTILS_ENABLE_ZSTD
154	options.compression = ::rocksdb::kZSTD;
155	options.compression_opts.level = constants::rocksdb::ZSTD_COMPRESSION_LEVEL;
156	options.compression_opts.max_dict_bytes =
157	constants::rocksdb::ZSTD_MAX_DICT_BYTES;
158	options.compression_opts.zstd_max_train_bytes =
159	constants::rocksdb::ZSTD_MAX_TRAIN_BYTES;
160	options.compression_opts.enabled = true;
161	options.bottommost_compression = ::rocksdb::kZSTD;
162	options.bottommost_compression_opts.level =
163	constants::rocksdb::ZSTD_COMPRESSION_LEVEL;
164	options.bottommost_compression_opts.max_dict_bytes =
165	constants::rocksdb::ZSTD_MAX_DICT_BYTES;
166	options.bottommost_compression_opts.zstd_max_train_bytes =
167	constants::rocksdb::ZSTD_MAX_TRAIN_BYTES;
168	options.bottommost_compression_opts.enabled = true;
169	#elif defined(DFTRACER_UTILS_ENABLE_LZ4)
170	options.compression = ::rocksdb::kLZ4Compression;
171	options.bottommost_compression = ::rocksdb::kZlibCompression;
172	#else
173	options.compression = ::rocksdb::kZlibCompression;	27,948✔
174	options.bottommost_compression = ::rocksdb::kZlibCompression;	27,948✔
175	#endif
176	return options;	41,793✔
177	}	27,948!
178
179	::rocksdb::ColumnFamilyOptions
180	RocksDatabase::point_lookup_column_family_options() {	18,287✔
181	auto options = default_column_family_options();	18,287!
182
183	// Hash keys are uniformly distributed, so every SST's range covers almost
184	// every key: without a filter a get searches all of them. Large blocks
185	// also mean reading and decompressing 32 KB to return a few dozen bytes.
186	::rocksdb::BlockBasedTableOptions bbt;	18,288✔
187	bbt.block_size = constants::rocksdb::POINT_LOOKUP_BLOCK_SIZE;	18,288✔
188	bbt.format_version = 7;	18,288✔
189	bbt.index_block_restart_interval = 16;	18,288✔
190	bbt.block_cache = shared_block_cache();	18,288!
191	// Content-hash keys are uniformly distributed: interpolation search beats
192	// binary. Read-time only, works on any SST.
193	bbt.index_block_search_type =	18,288✔
194	::rocksdb::BlockBasedTableOptions::kInterpolation;
195	bbt.filter_policy.reset(::rocksdb::NewBloomFilterPolicy(	18,288!
196	constants::rocksdb::BLOOM_BITS_PER_KEY, false));
197	// Left in the table reader rather than the block cache: a scan streams
198	// enough data blocks through the cache to evict the filters it needs.
199	bbt.cache_index_and_filter_blocks = false;	18,288✔
200	options.table_factory.reset(::rocksdb::NewBlockBasedTableFactory(bbt));	18,288!
201	return options;	27,346✔
202	}	18,288!
203
204	bool RocksDatabase::is_point_lookup_cf(std::string_view name) noexcept {	237,733✔
205	return name == cf::HASH_TABLES \|\| name == cf::NAME_DICTIONARY;	237,733✔
206	}
207
208	bool RocksDatabase::open(const std::string& db_path, OpenMode open_mode) {	9,188✔
209	close();	9,188!
210	db_path_ = db_path;	9,188!
211	open_mode_ = open_mode;	9,187✔
212
213	std::error_code ec;	9,187✔
214	if (open_mode_ == OpenMode::ReadWrite) {	9,188✔
215	fs::create_directories(fs::path(db_path_), ec);	1,856!
216	}	917✔
217
218	auto db_options = default_options();	9,188!
219	if (open_mode_ == OpenMode::ReadOnly) {	9,188✔
220	db_options.create_if_missing = false;	7,354✔
221	db_options.create_missing_column_families = false;	7,354✔
222	// Read paths do heavy point lookups (e.g. dfanalyzer hash resolution);
223	// with a small table cache the SST readers get evicted and every lookup
224	// re-opens the file to re-read its filter/index blocks, which thrashes.
225	// There is no write-side fd pressure here, so keep every SST open
226	// unless an explicit env cap is set.
227	if (!Env::get<int>("DFTRACER_UTILS_ROCKSDB_MAX_OPEN_FILES")	10,988!
228	.has_value()) {	7,354✔
229	db_options.max_open_files = -1;	7,354✔
230	}	3,720✔
231	}	3,720✔
232	file_system_ = make_dftracer_file_system();	9,188!
233	env_ = make_dftracer_env(file_system_);	9,188!
234	db_options.env = env_.get();	9,188✔
235	auto cf_options = default_column_family_options();	9,188!
236
237	std::vector<std::string> column_family_names;	9,188✔
238	auto list_status = ::rocksdb::DB::ListColumnFamilies(db_options, db_path_,	9,188!
239	&column_family_names);	4,551!
240	if (!list_status.ok()) {	9,188!
241	if (open_mode_ == OpenMode::ReadOnly) {	1,204✔
242	throw DFTUtilsException(	66!
243	ErrorCode::IO, "Failed to list RocksDB column families at '" +	22!
244	db_path_ + "': " + list_status.ToString());	66!
245	}
246	column_family_names.reserve(default_column_families().size());	1,160!
247	for (auto name : default_column_families()) {	31,314✔
248	column_family_names.emplace_back(name);	30,155!
249	}
250	} else {	580✔
251	if (open_mode_ == OpenMode::ReadWrite) {	7,984✔
252	for (const auto& name : default_column_families()) {	18,198✔
253	if (std::find(column_family_names.begin(),	26,285!
254	column_family_names.end(),	8,762✔
255	name) == column_family_names.end()) {	26,286!
256	column_family_names.emplace_back(name);	×
257	}
258	}
259	}	337✔
260	}
261
262	std::vector<::rocksdb::ColumnFamilyDescriptor> descriptors;	9,144✔
263	descriptors.reserve(column_family_names.size());	9,144!
264	for (const auto& name : column_family_names) {	246,877✔
265	auto opts = is_point_lookup_cf(name)	357,721✔
266	? point_lookup_column_family_options()	9,230!
267	: cf_options;	228,505!
268	if (cf_options_override_) {	237,739✔
269	cf_options_override_(name, opts);	235,971!
270	}	119,106✔
271	descriptors.emplace_back(name, opts);	237,738!
272	}	237,734✔
273
274	std::vector<::rocksdb::ColumnFamilyHandle*> handles;	9,144✔
275	::rocksdb::Status status;	9,143!
276	std::unique_ptr<::rocksdb::DB> opened;	9,144✔
277	if (open_mode_ == OpenMode::ReadOnly) {	9,144✔
278	status = ::rocksdb::DB::OpenForReadOnly(	7,310!
279	db_options, db_path_, descriptors, &handles, &opened, false);	7,310!
280	} else {	3,698✔
281	status = ::rocksdb::DB::Open(db_options, db_path_, descriptors,	2,751!
282	&handles, &opened);	917✔
283	}
284	if (!status.ok()) {	9,144!
NEW 285	::rocksdb::DB* raw = opened.release();	×
NEW 286	cleanup_failed_open(raw, handles);	×
287	throw DFTUtilsException(ErrorCode::IO, "Failed to open RocksDB at '" +	×
288	db_path_ +	×
289	"': " + status.ToString());	×
290	}
291	db_ = opened.release();	9,144✔
292
293	column_families_.clear();	9,144✔
294	for (std::size_t i = 0; i < descriptors.size(); ++i) {	246,887✔
295	column_families_.emplace(descriptors[i].name, handles[i]);	237,743!
296	}	119,990✔
297
298	return true;	4,529✔
299	}	9,276✔
300
301	void RocksDatabase::close() {	18,376✔
302	if (db_ == nullptr) {	18,376✔
303	column_families_.clear();	9,232✔
304	return;	9,232✔
305	}
306
307	if (process_exiting_flag().load(std::memory_order_relaxed)) {	9,144!
308	db_ = nullptr;	×
309	column_families_.clear();	×
310	env_.reset();	×
311	file_system_.reset();	×
312	db_path_.clear();	×
313	return;	×
314	}
315
316	for (auto& entry : column_families_) {	246,888✔
317	if (entry.second != nullptr) {	237,742✔
318	db_->DestroyColumnFamilyHandle(entry.second);	237,742!
319	entry.second = nullptr;	237,744✔
320	}	119,990✔
321	}
322	column_families_.clear();	9,144✔
323
324	auto* db = db_;	9,144✔
325	db_ = nullptr;	9,144✔
326	static_cast<void>(db->Close());	9,144✔
327	delete db;	9,144✔
328	env_.reset();	9,144✔
329	file_system_.reset();	9,144✔
330	db_path_.clear();	9,144✔
331	}	9,274✔
332
333	bool RocksDatabase::is_open() const noexcept { return db_ != nullptr; }	36✔
334
335	bool RocksDatabase::is_read_only() const noexcept {	1,298✔
336	return open_mode_ == OpenMode::ReadOnly;	1,298✔
337	}
338
339	const std::string& RocksDatabase::path() const noexcept { return db_path_; }	116✔
340
341	::rocksdb::DB* RocksDatabase::get() const noexcept { return db_; }	×
342
343	::rocksdb::ColumnFamilyHandle* RocksDatabase::column_family_handle(	83,488✔
344	std::string_view column_family) const {
345	const auto name = column_family.empty() ? std::string(cf::DEFAULT)	83,488!
346	: std::string(column_family);	125,258!
347	const auto it = column_families_.find(name);	83,502✔
348	if (it == column_families_.end() \|\| it->second == nullptr) {	83,458!
349	throw DFTUtilsException(ErrorCode::INVALID_ARGUMENT,	×
350	"Unknown RocksDB column family: " + name);	×
351	}
352	return it->second;	125,205✔
353	}	83,504✔
354
355	::rocksdb::Status RocksDatabase::put(std::string_view key,	1,502✔
356	std::string_view value,
357	std::string_view column_family) {
358	return db_->Put(write_options(), column_family_handle(column_family),	2,253✔
359	::rocksdb::Slice(key.data(), key.size()),	1,502✔
360	::rocksdb::Slice(value.data(), value.size()));	3,004!
361	}
362
363	::rocksdb::Status RocksDatabase::get(std::string_view key, std::string* value,	15,737✔
364	std::string_view column_family) const {
365	return db_->Get(read_options(), column_family_handle(column_family),	23,604✔
366	::rocksdb::Slice(key.data(), key.size()), value);	23,643!
367	}
368
369	::rocksdb::Status RocksDatabase::merge(std::string_view key,	×
370	std::string_view value,
371	std::string_view column_family) {
372	return db_->Merge(write_options(), column_family_handle(column_family),	×
373	::rocksdb::Slice(key.data(), key.size()),	×
374	::rocksdb::Slice(value.data(), value.size()));	×
375	}
376
377	void RocksDatabase::set_cf_options_override(CfOptionsOverride override) {	9,119✔
378	cf_options_override_ = std::move(override);	9,119✔
379	}	9,120✔
380
381	::rocksdb::Status RocksDatabase::merge(Batch& batch,	1,925✔
382	std::string_view column_family,
383	std::string_view key,
384	std::string_view value) {
385	return batch.Merge(column_family_handle(column_family),	2,892✔
386	::rocksdb::Slice(key.data(), key.size()),	1,926✔
387	::rocksdb::Slice(value.data(), value.size()));	3,837!
388	}
389
390	::rocksdb::Status RocksDatabase::del(std::string_view key,	×
391	std::string_view column_family) {
392	return db_->Delete(write_options(), column_family_handle(column_family),	×
393	::rocksdb::Slice(key.data(), key.size()));	×
394	}
395
396	::rocksdb::Status RocksDatabase::delete_range(std::string_view begin_key,	×
397	std::string_view end_key,
398	std::string_view column_family) {
399	return db_->DeleteRange(	×
400	write_options(), column_family_handle(column_family),
401	::rocksdb::Slice(begin_key.data(), begin_key.size()),	×
402	::rocksdb::Slice(end_key.data(), end_key.size()));	×
403	}
404
405	::rocksdb::Status RocksDatabase::put(Batch& batch,	53,069✔
406	std::string_view column_family,
407	std::string_view key,
408	std::string_view value) {
409	return batch.Put(column_family_handle(column_family),	79,616✔
410	::rocksdb::Slice(key.data(), key.size()),	53,084✔
411	::rocksdb::Slice(value.data(), value.size()));	106,157!
412	}
413
414	::rocksdb::Status RocksDatabase::del(Batch& batch,	120✔
415	std::string_view column_family,
416	std::string_view key) {
417	return batch.Delete(column_family_handle(column_family),	180✔
418	::rocksdb::Slice(key.data(), key.size()));	180!
419	}
420
421	RocksDatabase::Batch RocksDatabase::begin_batch() const { return Batch(); }	2,994✔
422
423	::rocksdb::Status RocksDatabase::commit_batch(Batch& batch) {	2,986✔
424	return db_->Write(write_options(), &batch);	2,986✔
425	}
426
427	std::unique_ptr<::rocksdb::Iterator> RocksDatabase::new_iterator(	10,619✔
428	std::string_view column_family) const {
429	return std::unique_ptr<::rocksdb::Iterator>(	5,317✔
430	db_->NewIterator(read_options(), column_family_handle(column_family)));	10,619✔
431	}
432
433	::rocksdb::Status RocksDatabase::compact(std::string_view column_family) {	48✔
434	::rocksdb::CompactRangeOptions opts;	48✔
435	opts.max_subcompactions = 8;	48✔
436	return db_->CompactRange(opts, column_family_handle(column_family), nullptr,	48✔
437	nullptr);	48!
438	}
439
UNCOV 440	::rocksdb::Status RocksDatabase::ingest_external_files(	×
441	std::string_view column_family,
442	const std::vector<std::string>& external_files, bool ingest_behind) {
UNCOV 443	if (external_files.empty()) {	×
UNCOV 444	return ::rocksdb::Status::OK();	×
445	}
UNCOV 446	::rocksdb::IngestExternalFileOptions opts;	×
447	// Rename same-FS staged SSTs instead of copying; RocksDB copies on failure.
NEW 448	opts.move_files = true;	×
449	// Offline bulk build with no live readers, so skip the snapshot/seqno sync.
NEW 450	opts.snapshot_consistency = false;	×
UNCOV 451	opts.allow_global_seqno = true;	×
452	// Keep the assigned seqno in the manifest instead of rewriting it into each
453	// SST (format_version 5 supports this), avoiding a per-SST write + fsync.
NEW 454	opts.write_global_seqno = false;	×
UNCOV 455	opts.allow_blocking_flush = true;	×
UNCOV 456	opts.ingest_behind = ingest_behind;	×
UNCOV 457	return db_->IngestExternalFile(column_family_handle(column_family),	×
UNCOV 458	external_files, opts);	×
459	}
460
461	::rocksdb::Status RocksDatabase::ingest_external_files_multi(	30✔
462	const std::vector<
463	std::pair<std::string_view, const std::vector<std::string>*>>& per_cf) {
464	::rocksdb::IngestExternalFileOptions opts;	30✔
465	opts.move_files = true;	30✔
466	opts.snapshot_consistency = false;	30✔
467	opts.allow_global_seqno = true;	30✔
468	opts.write_global_seqno = false;	30✔
469	opts.allow_blocking_flush = true;	30✔
470
471	std::vector<::rocksdb::IngestExternalFileArg> args;	30✔
472	args.reserve(per_cf.size());	30!
473	for (const auto& [cf_name, files] : per_cf) {	466✔
474	if (!files \|\| files->empty()) continue;	436!
475	::rocksdb::IngestExternalFileArg arg;	436✔
476	arg.column_family = column_family_handle(cf_name);	436!
477	arg.external_files = *files;	436!
478	arg.options = opts;	436✔
479	args.push_back(std::move(arg));	436!
480	}	436✔
481	if (args.empty()) return ::rocksdb::Status::OK();	30!
482	return db_->IngestExternalFiles(args);	30!
483	}	30✔
484
485	} // namespace dftracer::utils::rocksdb

llnl / dftracer-utils / 30058987329

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous