| /* |
| * MemConstrainedWriter experiment. |
| * |
| * Wraps TsFileWriter with explicit two-level memory-budget control: |
| * Level 1 – data budget: flush when M_data exceeds data_budget. |
| * Level 2 – meta budget: rotate file when M_meta exceeds meta_budget. |
| * |
| * Two measurement methods are computed simultaneously and logged to CSV: |
| * Direct – calculate_mem_size_for_all_group() / calculate_meta_mem_size() |
| * Formula – s_data * rows_since_flush / K_flush * b |
| * |
| * Control uses the DIRECT measurements. Formula estimates are logged purely |
| * for comparison (they overestimate M_data because they ignore compression). |
| * |
| * Schema: same 8-FIELD aligned table as write_memory experiment. |
| * |
| * Build (ENABLE_MEM_STAT not required): |
| * cmake -DBUILD_TEST=OFF .. |
| * cmake --build . --target writer_bench |
| * |
| * Usage: |
| * ./writer_bench [total_rows] [mem_limit_mb] [batch_cap] [csv] [base_path] |
| * |
| * Default: |
| * total_rows = 20 000 000 (10 devices x 2 M rows) |
| * mem_limit = 30 MB |
| * batch_cap = 65 536 rows |
| * csv = writer_bench.csv |
| * base_path = /tmp/mc_writer (files: base_path_0.tsfile, _1.tsfile, ...) |
| */ |
| |
| #include <fcntl.h> |
| |
| #include <climits> |
| #include <cstdint> |
| #include <fstream> |
| #include <iomanip> |
| #include <iostream> |
| #include <memory> |
| #include <random> |
| #include <string> |
| #include <vector> |
| |
| #include "common/config/config.h" |
| #include "common/schema.h" |
| #include "common/tablet.h" |
| #include "file/write_file.h" |
| #include "writer/tsfile_writer.h" |
| |
| using namespace storage; |
| using namespace common; |
| |
| // --------------------------------------------------------------------------- |
| // Schema constants (aligned table model) |
| // |
| // s_data = 8 (shared timestamp) + sum_field sizeof(field_type) |
| // = 8 + (4+4+8+8+4+4+8+8) = 56 bytes / row |
| // |
| // b = n_field * 104 + 96 |
| // = 8 * 104 + 96 = 928 bytes / device / flush |
| // --------------------------------------------------------------------------- |
| static const char* kTable = "mem_bench"; |
| static const int kNumDevices = 10; |
| static const int64_t kSData = 56; // formula: bytes per row (aligned) |
| static const int64_t kB = 928; // formula: meta bytes per device per flush |
| |
| static std::string device_name(int i) { return "device_" + std::to_string(i); } |
| |
| // --------------------------------------------------------------------------- |
| // MemSnapshot |
| // --------------------------------------------------------------------------- |
| struct MemSnapshot { |
| int64_t rows_written; |
| std::string event; // "batch" | "flush_before" | "flush_after" | |
| // "rotate_before" | "rotate_after" | "final_close" |
| int file_idx; |
| int flush_count; // K_cur: flushes in the current file |
| |
| // Direct measurements |
| int64_t m_data_direct; |
| int64_t m_meta_direct; |
| |
| // Formula estimates |
| int64_t m_data_formula; |
| int64_t |
| m_meta_formula; // total_flush_count * b (sequential: 1 dev/flush) |
| |
| // Budget reference lines |
| int64_t data_budget; |
| int64_t meta_budget; |
| }; |
| |
| // --------------------------------------------------------------------------- |
| // MemConstrainedWriter |
| // --------------------------------------------------------------------------- |
| class MemConstrainedWriter { |
| public: |
| MemConstrainedWriter(const std::string& base_path, |
| std::shared_ptr<TableSchema> schema, int n_devices, |
| int64_t mem_limit_bytes, |
| int64_t m_init_bytes = 900LL * 1024, |
| double alpha = 0.9) |
| : base_path_(base_path), |
| schema_(std::move(schema)), |
| n_devices_(n_devices), |
| mem_limit_(mem_limit_bytes), |
| m_init_(m_init_bytes), |
| alpha_(alpha), |
| rows_since_flush_(0), |
| total_flush_count_(0), |
| flush_count_cur_(0), |
| file_count_(0), |
| total_rows_(0), |
| writer_(nullptr), |
| write_file_(nullptr) { |
| m_avail_ = mem_limit_ - m_init_; |
| data_budget_ = m_avail_ / 2; |
| meta_budget_ = m_avail_ / 2; |
| |
| // Disable TsFileWriter's built-in auto-flush so we drive it ourselves. |
| g_config_value_.chunk_group_size_threshold_ = INT64_MAX / 2; |
| |
| open_new_file(); |
| } |
| |
| ~MemConstrainedWriter() { cleanup_writer(); } |
| |
| // nrows must match the number of rows written to the tablet. |
| int write_tablet(Tablet& tablet, int64_t nrows) { |
| int ret = writer_->write_table(tablet); |
| if (ret != 0) return ret; |
| |
| rows_since_flush_ += nrows; |
| total_rows_ += nrows; |
| |
| take_snapshot("batch"); |
| |
| // Level 1: flush when direct M_data exceeds data budget |
| int64_t m_data = writer_->calculate_mem_size_for_all_group(); |
| if (m_data > data_budget_) { |
| do_flush(); |
| } |
| return 0; |
| } |
| |
| int close() { |
| take_snapshot("before_close"); |
| // Final flush (if any data remains) |
| if (writer_->calculate_mem_size_for_all_group() > 0) { |
| take_snapshot("flush_before"); |
| writer_->flush(); |
| total_flush_count_++; |
| flush_count_cur_++; |
| rows_since_flush_ = 0; |
| take_snapshot("flush_after"); |
| } |
| writer_->close(); |
| take_snapshot("final_close"); |
| cleanup_writer(); |
| return 0; |
| } |
| |
| const std::vector<MemSnapshot>& snapshots() const { return snapshots_; } |
| |
| // Derived parameters (for printing) |
| int64_t data_budget() const { return data_budget_; } |
| int64_t meta_budget() const { return meta_budget_; } |
| int64_t rows_per_flush_formula() const { return data_budget_ / kSData; } |
| int64_t flushes_per_file_formula() const { |
| // sequential: 1 device per flush, so meta grows by kB per flush |
| int64_t k = meta_budget_ / kB; |
| return k > 0 ? k : 1; |
| } |
| |
| private: |
| // Open a new tsfile and register the table schema |
| void open_new_file() { |
| std::string path = |
| base_path_ + "_" + std::to_string(file_count_) + ".tsfile"; |
| |
| write_file_ = new WriteFile(); |
| int ret = write_file_->create(path.c_str(), |
| O_WRONLY | O_CREAT | O_TRUNC, 0666); |
| if (ret != 0) { |
| std::cerr << "Failed to create file: " << path << " err=" << ret |
| << "\n"; |
| return; |
| } |
| |
| writer_ = new TsFileWriter(); |
| writer_->init(write_file_); |
| writer_->register_table(schema_); |
| |
| rows_since_flush_ = 0; |
| flush_count_cur_ = 0; |
| } |
| |
| void cleanup_writer() { |
| delete writer_; |
| delete write_file_; |
| writer_ = nullptr; |
| write_file_ = nullptr; |
| } |
| |
| void do_flush() { |
| take_snapshot("flush_before"); |
| writer_->flush(); |
| total_flush_count_++; |
| flush_count_cur_++; |
| rows_since_flush_ = 0; |
| take_snapshot("flush_after"); |
| |
| // Level 2: rotate file when direct M_meta exceeds meta budget * alpha |
| int64_t m_meta = writer_->calculate_meta_mem_size(); |
| if (m_meta >= static_cast<int64_t>(meta_budget_ * alpha_)) { |
| rotate_file(); |
| } |
| } |
| |
| void rotate_file() { |
| take_snapshot("rotate_before"); |
| writer_->close(); |
| cleanup_writer(); |
| |
| file_count_++; |
| flush_count_cur_ = 0; |
| open_new_file(); |
| take_snapshot("rotate_after"); |
| } |
| |
| void take_snapshot(const std::string& event) { |
| MemSnapshot s; |
| s.rows_written = total_rows_; |
| s.event = event; |
| s.file_idx = file_count_; |
| s.flush_count = flush_count_cur_; |
| |
| // Direct |
| s.m_data_direct = |
| writer_ ? writer_->calculate_mem_size_for_all_group() : 0; |
| s.m_meta_direct = writer_ ? writer_->calculate_meta_mem_size() : 0; |
| |
| // Formula (sequential mode: each flush = 1 device) |
| s.m_data_formula = kSData * rows_since_flush_; |
| s.m_meta_formula = total_flush_count_ * kB; |
| |
| s.data_budget = data_budget_; |
| s.meta_budget = meta_budget_; |
| |
| snapshots_.push_back(std::move(s)); |
| } |
| |
| std::string base_path_; |
| std::shared_ptr<TableSchema> schema_; |
| [[maybe_unused]] int n_devices_; |
| int64_t mem_limit_; |
| int64_t m_init_; |
| double alpha_; |
| int64_t m_avail_; |
| int64_t data_budget_; |
| int64_t meta_budget_; |
| |
| int64_t rows_since_flush_; |
| int64_t total_flush_count_; // all flushes across all files |
| int flush_count_cur_; // flushes in the current file |
| int file_count_; |
| int64_t total_rows_; |
| |
| TsFileWriter* writer_; |
| WriteFile* write_file_; |
| |
| std::vector<MemSnapshot> snapshots_; |
| }; |
| |
| // --------------------------------------------------------------------------- |
| // main |
| // --------------------------------------------------------------------------- |
| int main(int argc, char* argv[]) { |
| int64_t total_rows = 20000000LL; |
| int64_t mem_limit_mb = 30; |
| uint32_t batch_cap = 65536; |
| std::string csv_path = "writer_bench.csv"; |
| std::string base_path = "/tmp/mc_writer"; |
| |
| if (argc > 1) total_rows = std::atoll(argv[1]); |
| if (argc > 2) mem_limit_mb = std::atoll(argv[2]); |
| if (argc > 3) batch_cap = static_cast<uint32_t>(std::atoi(argv[3])); |
| if (argc > 4) csv_path = argv[4]; |
| if (argc > 5) base_path = argv[5]; |
| |
| int64_t mem_limit = mem_limit_mb * 1024LL * 1024; |
| |
| libtsfile_init(); |
| |
| // ------------------------------------------------------------------ |
| // Build table schema (same as write_memory experiment) |
| // ------------------------------------------------------------------ |
| auto schema = std::make_shared<TableSchema>( |
| std::string(kTable), |
| std::vector<ColumnSchema>{ |
| ColumnSchema("id1", STRING, UNCOMPRESSED, PLAIN, |
| ColumnCategory::TAG), |
| ColumnSchema("id2", STRING, UNCOMPRESSED, PLAIN, |
| ColumnCategory::TAG), |
| ColumnSchema("s1", INT32, SNAPPY, TS_2DIFF, ColumnCategory::FIELD), |
| ColumnSchema("s2", INT32, SNAPPY, TS_2DIFF, ColumnCategory::FIELD), |
| ColumnSchema("s3", INT64, SNAPPY, TS_2DIFF, ColumnCategory::FIELD), |
| ColumnSchema("s4", INT64, SNAPPY, TS_2DIFF, ColumnCategory::FIELD), |
| ColumnSchema("s5", FLOAT, SNAPPY, GORILLA, ColumnCategory::FIELD), |
| ColumnSchema("s6", FLOAT, SNAPPY, GORILLA, ColumnCategory::FIELD), |
| ColumnSchema("s7", DOUBLE, SNAPPY, GORILLA, ColumnCategory::FIELD), |
| ColumnSchema("s8", DOUBLE, SNAPPY, GORILLA, ColumnCategory::FIELD), |
| }); |
| |
| MemConstrainedWriter writer(base_path, schema, kNumDevices, mem_limit); |
| |
| std::cout << "=== MemConstrainedWriter Experiment ===\n" |
| << " total_rows: " << total_rows << "\n" |
| << " mem_limit: " << mem_limit_mb << " MB\n" |
| << " batch_cap: " << batch_cap << " rows\n" |
| << " devices: " << kNumDevices << "\n" |
| << " s_data (formula): " << kSData << " bytes/row\n" |
| << " b (formula): " << kB << " bytes/device/flush\n" |
| << " data_budget: " |
| << writer.data_budget() / (1024 * 1024) << " MB\n" |
| << " meta_budget: " |
| << writer.meta_budget() / (1024 * 1024) << " MB\n" |
| << " F (formula): " << writer.rows_per_flush_formula() |
| << " rows/flush\n" |
| << " K (formula, seq): " << writer.flushes_per_file_formula() |
| << " flushes/file\n" |
| << " csv: " << csv_path << "\n" |
| << " base_path: " << base_path << "\n\n"; |
| |
| // ------------------------------------------------------------------ |
| // Data generation helpers |
| // ------------------------------------------------------------------ |
| std::mt19937_64 rng(42); |
| std::uniform_int_distribution<int32_t> ni(-100, 100); |
| std::uniform_real_distribution<float> nf(-5.0f, 5.0f); |
| std::uniform_real_distribution<double> nd(-0.5, 0.5); |
| |
| auto make_tablet = [](uint32_t n) { |
| return Tablet( |
| kTable, |
| {"id1", "id2", "s1", "s2", "s3", "s4", "s5", "s6", "s7", "s8"}, |
| {STRING, STRING, INT32, INT32, INT64, INT64, FLOAT, FLOAT, DOUBLE, |
| DOUBLE}, |
| {ColumnCategory::TAG, ColumnCategory::TAG, ColumnCategory::FIELD, |
| ColumnCategory::FIELD, ColumnCategory::FIELD, |
| ColumnCategory::FIELD, ColumnCategory::FIELD, |
| ColumnCategory::FIELD, ColumnCategory::FIELD, |
| ColumnCategory::FIELD}, |
| std::max(n, 1u)); |
| }; |
| |
| // ------------------------------------------------------------------ |
| // Write: sequential per-device |
| // ------------------------------------------------------------------ |
| int64_t rows_per_dev = total_rows / kNumDevices; |
| |
| for (int dev = 0; dev < kNumDevices; dev++) { |
| std::string dev_id = device_name(dev); |
| int64_t dev_base = static_cast<int64_t>(dev) * rows_per_dev; |
| |
| for (int64_t off = 0; off < rows_per_dev;) { |
| uint32_t n = static_cast<uint32_t>( |
| std::min<int64_t>(batch_cap, rows_per_dev - off)); |
| auto tablet = make_tablet(n); |
| |
| for (uint32_t i = 0; i < n; i++) { |
| int64_t ts = dev_base + off + i; |
| tablet.add_timestamp(i, ts); |
| tablet.add_value(i, "id1", dev_id.c_str()); |
| tablet.add_value(i, "id2", "tag_b"); |
| tablet.add_value(i, "s1", |
| static_cast<int32_t>(ts % 100000) + ni(rng)); |
| tablet.add_value(i, "s2", |
| static_cast<int32_t>(ts % 50000) + ni(rng)); |
| tablet.add_value(i, "s3", ts + static_cast<int64_t>(ni(rng))); |
| tablet.add_value(i, "s4", |
| ts * 2 + static_cast<int64_t>(ni(rng))); |
| tablet.add_value(i, "s5", |
| static_cast<float>(ts % 10000) + nf(rng)); |
| tablet.add_value(i, "s6", |
| static_cast<float>(ts % 5000) + nf(rng)); |
| tablet.add_value(i, "s7", ts * 1.1 + nd(rng)); |
| tablet.add_value(i, "s8", ts * 0.5 + nd(rng)); |
| } |
| |
| int ret = writer.write_tablet(tablet, static_cast<int64_t>(n)); |
| if (ret != 0) { |
| std::cerr << "write_tablet failed: " << ret << "\n"; |
| return 1; |
| } |
| off += n; |
| } |
| |
| std::cout << " device " << std::setw(2) << dev << " done\n"; |
| } |
| |
| writer.close(); |
| |
| // ------------------------------------------------------------------ |
| // Write CSV |
| // ------------------------------------------------------------------ |
| std::ofstream csv(csv_path); |
| if (!csv.is_open()) { |
| std::cerr << "cannot open csv: " << csv_path << "\n"; |
| return 1; |
| } |
| |
| csv << "rows_written,event,file_idx,flush_count," |
| "m_data_direct,m_meta_direct,m_total_direct," |
| "m_data_formula,m_meta_formula,m_total_formula," |
| "data_budget,meta_budget\n"; |
| |
| for (const auto& s : writer.snapshots()) { |
| int64_t m_total_direct = s.m_data_direct + s.m_meta_direct; |
| int64_t m_total_formula = s.m_data_formula + s.m_meta_formula; |
| csv << s.rows_written << "," << s.event << "," << s.file_idx << "," |
| << s.flush_count << "," << s.m_data_direct << "," << s.m_meta_direct |
| << "," << m_total_direct << "," << s.m_data_formula << "," |
| << s.m_meta_formula << "," << m_total_formula << "," |
| << s.data_budget << "," << s.meta_budget << "\n"; |
| } |
| csv.close(); |
| |
| std::cout << "\nSnapshots: " << writer.snapshots().size() << "\n" |
| << "CSV: " << csv_path << "\n"; |
| |
| libtsfile_destroy(); |
| return 0; |
| } |
