Repository: tikv/tikv
File: components/engine_rocks/src/properties.rs (Rust)

Code before:

const PROP_TOTAL_SIZE: &str = "tikv.total_size";
const PROP_SIZE_INDEX: &str = "tikv.size_index";
const PROP_RANGE_INDEX: &str = "tikv.range_index";
pub const DEFAULT_PROP_SIZE_INDEX_DISTANCE: u64 = 4 * 1024 * 1024;
pub const DEFAULT_PROP_KEYS_INDEX_DISTANCE: u64 = 40 * 1024;

/// Caps the expansion factor to avoid unrealistic estimates due to
/// miscalculation or noise.
const TITAN_MAX_EXPANSION_FACTOR: f64 = 5000.0;
/// Number of samples to keep for estimating Titan blob size. With a default
/// metrics flush interval of 10s, this corresponds to 5 minutes.
const TITAN_ESTIMATOR_WINDOW_SIZE: usize = 30;

lazy_static! {
    // Tracks Titan blob compression ratio for raw size estimation.
    pub static ref TITAN_BLOB_SIZE_ESTIMATOR: TitanBlobSizeEstimator =
        TitanBlobSizeEstimator::new();
}

fn get_entry_size(value: &[u8], entry_type: DBEntryType) -> std::result::Result<u64, ()> {
    match entry_type {
        DBEntryType::Put => Ok(value.len() as u64),
        DBEntryType::BlobIndex => match TitanBlobIndex::decode(value) {
            Ok(index) => {
                let estimated_raw_size = TITAN_BLOB_SIZE_ESTIMATOR.estimate(index.blob_size);
                Ok(estimated_raw_size + value.len() as u64)
            }
            Err(_) => Err(()),
        },
        _ => Err(()),
    }
}

/// Tracks recent Titan blob compression stats to estimate raw blob size from
/// compressed size. Maintains a sliding window of recent (raw, compressed)
/// pairs and the largest blob seen.
pub struct TitanBlobSizeEstimator {
    history: Mutex<VecDeque<(u64, u64)>>, // (raw, compressed)
    expansion_factor: AtomicU64,          // total_raw / total_compressed
    max_blob_value: AtomicU64,
}

impl TitanBlobSizeEstimator {
    pub fn new() -> Self {
        Self {
            history: Mutex::new(VecDeque::with_capacity(TITAN_ESTIMATOR_WINDOW_SIZE)),
            expansion_factor: AtomicU64::new(f64::to_bits(1.0)),
            max_blob_value: AtomicU64::new(0),
        }
    }

Review comment (performance):
You already use atomic operations, however, I still worry about the performance impact, since `get_entry_size()` is called for every entry in the compaction job. Would you please do a benchmark? Or change it to a "lighter" implementation, e.g. cache it in the collector?

Code after:

const PROP_TOTAL_SIZE: &str = "tikv.total_size";
const PROP_SIZE_INDEX: &str = "tikv.size_index";
const PROP_RANGE_INDEX: &str = "tikv.range_index";
pub const DEFAULT_PROP_SIZE_INDEX_DISTANCE: u64 = 4 * 1024 * 1024;
pub const DEFAULT_PROP_KEYS_INDEX_DISTANCE: u64 = 40 * 1024;

/// Caps the compaction factor to avoid unrealistic estimates.
pub const TITAN_MAX_COMPACTION_FACTOR: f64 = 5000.0;
const FIVE_MINS_IN_SECONDS: u64 = 5 * 60;

lazy_static! {
    // A global smoother used to estimate the Titan blob compression factor over
    // the last 5 minutes. The window size is 30, roughly matching the number of
    // data points collected in 5 minutes assuming one data point every 10 secs.
    pub static ref TITAN_COMPRESSION_FACTOR_SMOOTHER: Mutex<Smoother<f64, 30, FIVE_MINS_IN_SECONDS, 0>> =
        Mutex::new(Smoother::<f64, 30, FIVE_MINS_IN_SECONDS, 0>::default());
    pub static ref TITAN_COMPRESSION_FACTOR: AtomicU64 = AtomicU64::new(f64::to_bits(1.0));
    pub static ref TITAN_MAX_BLOB_SIZE_SEEN: AtomicU64 = AtomicU64::new(u64::MAX);
}

fn get_entry_size(
    value: &[u8],
    entry_type: DBEntryType,
    titan_compression_factor: f64,
    titan_max_blob_size: u64,
) -> std::result::Result<u64, ()> {
    match entry_type {
        DBEntryType::Put => Ok(value.len() as u64),
        DBEntryType::BlobIndex => match TitanBlobIndex::decode(value) {
            Ok(index) => {
                // Estimate the raw blob size using the Titan compression factor.
                let estimation = (index.blob_size as f64 * titan_compression_factor) as u64;
                let blob_raw_size = estimation.min(titan_max_blob_size).max(index.blob_size);
                Ok(blob_raw_size + value.len() as u64)
            }
            Err(_) => Err(()),
        },
        _ => Err(()),
    }
}

// Deprecated. Only for compatible issue from v2.0 or older version.
#[derive(Debug, Default)]
pub struct SizeProperties {
    pub total_size: u64,
    pub index_handles: IndexHandles,
}

impl SizeProperties {
    pub fn encode(&self) -> UserProperties {