30069185152

Committed 24 Jul 2026 05:20AM UTC coverage: 50.687% (-2.0%) from 52.66%

Build # 30069185152

Build Type

Pull #99

github

Committed by

web-flow

Commit Message

Merge 47af36077 into 06bc84ec9

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

Coverage Stats

18062 of 48203 branches covered (37.47%)

Branch coverage included in aggregate %.

1510 of 2205 new or added lines in 59 files covered. (68.48%)

742 existing lines in 114 files now uncovered.

23711 of 34210 relevant lines covered (69.31%)

39853.91 hits per line

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

61.89

/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() {	4,563✔
23	static std::atomic<bool> flag{false};
24	return flag;	4,563✔
25	}
26
27	const ::rocksdb::ReadOptions& read_options() {	13,247✔
28	static const auto options = [] {	115✔
29	::rocksdb::ReadOptions ro;	115✔
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;	115✔
34	ro.adaptive_readahead = true;	115✔
35	return ro;	115✔
36	}();	13,247!
37	return options;	13,243✔
38	}
39
40	const ::rocksdb::WriteOptions& write_options() {	2,245✔
41	static const auto options = [] {	112✔
42	::rocksdb::WriteOptions wo;	112✔
43	wo.disableWAL = true;	112✔
44	return wo;	112✔
45	}();	2,245!
46	return options;	2,248✔
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() {	6✔
67	process_exiting_flag().store(true, std::memory_order_relaxed);	6✔
68	}	6✔
69
70	RocksDatabase::RocksDatabase() = default;	4,575✔
71
72	RocksDatabase::RocksDatabase(const std::string& db_path, OpenMode open_mode) {	3✔
73	open(db_path, open_mode);	3!
74	}	3✔
75
76	RocksDatabase::~RocksDatabase() { close(); }	4,579✔
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() {	1,497✔
100	return cf::ALL;	1,497✔
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() {	23,158✔
107	static std::shared_ptr<::rocksdb::Cache> cache =
108	::rocksdb::NewLRUCache(constants::rocksdb::BLOCK_CACHE_BYTES);	23,158!
109	return cache;	23,157✔
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() {	4,815✔
115	static std::shared_ptr<::rocksdb::WriteBufferManager> mgr =
116	std::make_shared<::rocksdb::WriteBufferManager>(
117	constants::rocksdb::WRITE_BUFFER_BYTES, shared_block_cache());	4,815!
118	return mgr;	4,815✔
119	}
120	} // namespace
121
122	::rocksdb::Options RocksDatabase::default_options() {	4,815✔
123	::rocksdb::Options options;	4,815✔
124	options.create_if_missing = true;	4,815✔
125	options.create_missing_column_families = true;	4,815✔
126	options.allow_concurrent_memtable_write = true;	4,815✔
127	options.enable_pipelined_write = true;	4,815✔
128	options.max_open_files = Env::rocksdb_max_open_files();	4,815!
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;	4,815✔
133	options.max_background_jobs = 8;	4,815✔
134	options.max_subcompactions = 8;	4,815✔
135	options.write_buffer_size = 256 * 1024 * 1024;	4,815✔
136	options.max_write_buffer_number = 4;	4,815✔
137	options.write_buffer_manager = shared_write_buffer_manager();	4,815!
138	return options;	4,815✔
UNCOV 139	}	×
140
141	::rocksdb::ColumnFamilyOptions RocksDatabase::default_column_family_options() {	13,928✔
142	::rocksdb::ColumnFamilyOptions options;	13,928!
143
144	::rocksdb::BlockBasedTableOptions bbt;	13,929✔
145	bbt.block_size = 32 * 1024;	13,929✔
146	bbt.format_version = 7;	13,929✔
147	bbt.index_block_restart_interval = 16;	13,929✔
148	bbt.block_cache = shared_block_cache();	13,929!
149	// Better compression and read CPU for the scanned CFs (fixed-layout blobs).
150	bbt.separate_key_value_in_data_block = true;	13,929✔
151	options.table_factory.reset(::rocksdb::NewBlockBasedTableFactory(bbt));	13,929!
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;	13,929✔
174	options.bottommost_compression = ::rocksdb::kZlibCompression;	13,929✔
175	#endif
176	return options;	27,858✔
177	}	13,929✔
178
179	::rocksdb::ColumnFamilyOptions
180	RocksDatabase::point_lookup_column_family_options() {	9,113✔
181	auto options = default_column_family_options();	9,113!
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;	9,114✔
187	bbt.block_size = constants::rocksdb::POINT_LOOKUP_BLOCK_SIZE;	9,114✔
188	bbt.format_version = 7;	9,114✔
189	bbt.index_block_restart_interval = 16;	9,114✔
190	bbt.block_cache = shared_block_cache();	9,114!
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 =	9,114✔
194	::rocksdb::BlockBasedTableOptions::kInterpolation;
195	bbt.filter_policy.reset(::rocksdb::NewBloomFilterPolicy(	9,114!
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;	9,114✔
200	options.table_factory.reset(::rocksdb::NewBlockBasedTableFactory(bbt));	9,114!
201	return options;	18,228✔
202	}	9,114✔
203
204	bool RocksDatabase::is_point_lookup_cf(std::string_view name) noexcept {	118,462✔
205	return name == cf::HASH_TABLES \|\| name == cf::NAME_DICTIONARY;	118,462✔
206	}
207
208	bool RocksDatabase::open(const std::string& db_path, OpenMode open_mode) {	4,579✔
209	close();	4,579!
210	db_path_ = db_path;	4,579!
211	open_mode_ = open_mode;	4,579✔
212
213	std::error_code ec;	4,579✔
214	if (open_mode_ == OpenMode::ReadWrite) {	4,579✔
215	fs::create_directories(fs::path(db_path_), ec);	917!
216	}
217
218	auto db_options = default_options();	4,578!
219	if (open_mode_ == OpenMode::ReadOnly) {	4,579✔
220	db_options.create_if_missing = false;	3,662✔
221	db_options.create_missing_column_families = false;	3,662✔
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")	7,324!
228	.has_value()) {	3,662✔
229	db_options.max_open_files = -1;	3,662✔
230	}
231	}
232	file_system_ = make_dftracer_file_system();	4,579!
233	env_ = make_dftracer_env(file_system_);	4,578!
234	db_options.env = env_.get();	4,578✔
235	auto cf_options = default_column_family_options();	4,578!
236
237	std::vector<std::string> column_family_names;	4,579✔
238	auto list_status = ::rocksdb::DB::ListColumnFamilies(db_options, db_path_,	4,579✔
239	&column_family_names);	4,579!
240	if (!list_status.ok()) {	4,579!
241	if (open_mode_ == OpenMode::ReadOnly) {	602✔
242	throw DFTUtilsException(	22!
UNCOV 243	ErrorCode::IO, "Failed to list RocksDB column families at '" +	×
244	db_path_ + "': " + list_status.ToString());	44!
245	}
246	column_family_names.reserve(default_column_families().size());	580!
247	for (auto name : default_column_families()) {	15,647!
248	column_family_names.emplace_back(name);	15,068!
249	}
250	} else {
251	if (open_mode_ == OpenMode::ReadWrite) {	3,977✔
252	for (const auto& name : default_column_families()) {	9,099!
253	if (std::find(column_family_names.begin(),	8,762!
254	column_family_names.end(),
255	name) == column_family_names.end()) {	17,524!
256	column_family_names.emplace_back(name);	×
257	}
258	}
259	}
260	}
261
262	std::vector<::rocksdb::ColumnFamilyDescriptor> descriptors;	4,557✔
263	descriptors.reserve(column_family_names.size());	4,557!
264	for (const auto& name : column_family_names) {	123,029✔
265	auto opts = is_point_lookup_cf(name)	118,470✔
266	? point_lookup_column_family_options()
267	: cf_options;	118,468!
268	if (cf_options_override_) {	118,476✔
269	cf_options_override_(name, opts);	117,593!
270	}
271	descriptors.emplace_back(name, opts);	118,477!
272	}	118,479✔
273
274	std::vector<::rocksdb::ColumnFamilyHandle*> handles;	4,556✔
275	::rocksdb::Status status;	4,557✔
276	std::unique_ptr<::rocksdb::DB> opened;	4,557✔
277	if (open_mode_ == OpenMode::ReadOnly) {	4,557✔
278	status = ::rocksdb::DB::OpenForReadOnly(	3,640✔
279	db_options, db_path_, descriptors, &handles, &opened, false);	3,640!
280	} else {
281	status = ::rocksdb::DB::Open(db_options, db_path_, descriptors,	1,834!
282	&handles, &opened);	917✔
283	}
284	if (!status.ok()) {	4,557!
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();	4,557✔
292
293	column_families_.clear();	4,557✔
294	for (std::size_t i = 0; i < descriptors.size(); ++i) {	123,036✔
295	column_families_.emplace(descriptors[i].name, handles[i]);	118,480!
296	}
297
298	return true;	4,557✔
299	}	4,644✔
300
301	void RocksDatabase::close() {	9,158✔
302	if (db_ == nullptr) {	9,158✔
303	column_families_.clear();	4,601✔
304	return;	4,601✔
305	}
306
307	if (process_exiting_flag().load(std::memory_order_relaxed)) {	4,557!
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_) {	123,027✔
317	if (entry.second != nullptr) {	118,475!
318	db_->DestroyColumnFamilyHandle(entry.second);	118,480!
319	entry.second = nullptr;	118,475✔
320	}
321	}
322	column_families_.clear();	4,555✔
323
324	auto* db = db_;	4,557✔
325	db_ = nullptr;	4,557✔
326	static_cast<void>(db->Close());	4,557✔
327	delete db;	4,557!
328	env_.reset();	4,557✔
329	file_system_.reset();	4,557✔
330	db_path_.clear();	4,557✔
331	}
332
333	bool RocksDatabase::is_open() const noexcept { return db_ != nullptr; }	18✔
334
335	bool RocksDatabase::is_read_only() const noexcept {	676✔
336	return open_mode_ == OpenMode::ReadOnly;	676✔
337	}
338
339	const std::string& RocksDatabase::path() const noexcept { return db_path_; }	58✔
340
341	::rocksdb::DB* RocksDatabase::get() const noexcept { return db_; }	×
342
343	::rocksdb::ColumnFamilyHandle* RocksDatabase::column_family_handle(	41,796✔
344	std::string_view column_family) const {
345	const auto name = column_family.empty() ? std::string(cf::DEFAULT)	41,796✔
346	: std::string(column_family);	83,551!
347	const auto it = column_families_.find(name);	41,772!
348	if (it == column_families_.end() \|\| it->second == nullptr) {	41,771!
349	throw DFTUtilsException(ErrorCode::INVALID_ARGUMENT,	×
350	"Unknown RocksDB column family: " + name);	×
351	}
352	return it->second;	83,469✔
353	}	41,706✔
354
355	::rocksdb::Status RocksDatabase::put(std::string_view key,	751✔
356	std::string_view value,
357	std::string_view column_family) {
358	return db_->Put(write_options(), column_family_handle(column_family),	751✔
359	::rocksdb::Slice(key.data(), key.size()),	751✔
360	::rocksdb::Slice(value.data(), value.size()));	2,253!
361	}
362
363	::rocksdb::Status RocksDatabase::get(std::string_view key, std::string* value,	7,915✔
364	std::string_view column_family) const {
365	return db_->Get(read_options(), column_family_handle(column_family),	7,915✔
366	::rocksdb::Slice(key.data(), key.size()), value);	15,846!
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) {	4,544✔
378	cf_options_override_ = std::move(override);	4,544✔
379	}	4,545✔
380
381	::rocksdb::Status RocksDatabase::merge(Batch& batch,	963✔
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),	963✔
386	::rocksdb::Slice(key.data(), key.size()),	963✔
387	::rocksdb::Slice(value.data(), value.size()));	2,894!
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,	26,531✔
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),	26,531✔
410	::rocksdb::Slice(key.data(), key.size()),	26,536✔
411	::rocksdb::Slice(value.data(), value.size()));	79,618!
412	}
413
414	::rocksdb::Status RocksDatabase::del(Batch& batch,	60✔
415	std::string_view column_family,
416	std::string_view key) {
417	return batch.Delete(column_family_handle(column_family),	60✔
418	::rocksdb::Slice(key.data(), key.size()));	120!
419	}
420
421	RocksDatabase::Batch RocksDatabase::begin_batch() const { return Batch(); }	1,499✔
422
423	::rocksdb::Status RocksDatabase::commit_batch(Batch& batch) {	1,499✔
424	return db_->Write(write_options(), &batch);	1,499✔
425	}
426
427	std::unique_ptr<::rocksdb::Iterator> RocksDatabase::new_iterator(	5,345✔
428	std::string_view column_family) const {
429	return std::unique_ptr<::rocksdb::Iterator>(
430	db_->NewIterator(read_options(), column_family_handle(column_family)));	5,345✔
431	}
432
433	::rocksdb::Status RocksDatabase::compact(std::string_view column_family) {	24✔
434	::rocksdb::CompactRangeOptions opts;	24✔
435	opts.max_subcompactions = 8;	24✔
436	return db_->CompactRange(opts, column_family_handle(column_family), nullptr,	24✔
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(	15✔
462	const std::vector<
463	std::pair<std::string_view, const std::vector<std::string>*>>& per_cf) {
464	::rocksdb::IngestExternalFileOptions opts;	15✔
465	opts.move_files = true;	15✔
466	opts.snapshot_consistency = false;	15✔
467	opts.allow_global_seqno = true;	15✔
468	opts.write_global_seqno = false;	15✔
469	opts.allow_blocking_flush = true;	15✔
470
471	std::vector<::rocksdb::IngestExternalFileArg> args;	15✔
472	args.reserve(per_cf.size());	15!
473	for (const auto& [cf_name, files] : per_cf) {	233✔
474	if (!files \|\| files->empty()) continue;	218!
475	::rocksdb::IngestExternalFileArg arg;	218✔
476	arg.column_family = column_family_handle(cf_name);	218!
477	arg.external_files = *files;	218!
478	arg.options = opts;	218✔
479	args.push_back(std::move(arg));	218!
480	}	218✔
481	if (args.empty()) return ::rocksdb::Status::OK();	15!
482	return db_->IngestExternalFiles(args);	15!
483	}	15✔
484
485	} // namespace dftracer::utils::rocksdb

llnl / dftracer-utils / 30069185152

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