vinted-prometheus-client-mmap 1.2.1

data/ext/fast_mmaped_file_rs/src/map.rs

@@ -0,0 +1,492 @@
+use hashbrown::hash_map::RawEntryMut;
+use hashbrown::HashMap;
+use magnus::{exception::*, Error, RArray};
+use std::hash::{BuildHasher, Hash, Hasher};
+use std::mem::size_of;
+
+use crate::error::MmapError;
+use crate::file_entry::{BorrowedData, EntryData, EntryMetadata, FileEntry};
+use crate::file_info::FileInfo;
+use crate::raw_entry::RawEntry;
+use crate::util::read_u32;
+use crate::Result;
+use crate::{err, HEADER_SIZE};
+
+/// A HashMap of JSON strings and their associated metadata.
+/// Used to print metrics in text format.
+///
+/// The map key is the entry's JSON string and an optional pid string. The latter
+/// allows us to have multiple entries on the map for multiple pids using the
+/// same string.
+#[derive(Default, Debug)]
+pub struct EntryMap(HashMap<EntryData, EntryMetadata>);
+
+impl EntryMap {
+    /// Construct a new EntryMap.
+    pub fn new() -> Self {
+        Self(HashMap::new())
+    }
+
+    /// Given a list of files, read each one into memory and parse the metrics it contains.
+    pub fn aggregate_files(&mut self, list_of_files: RArray) -> magnus::error::Result<()> {
+        // Pre-allocate the `HashMap` and validate we don't OOM. The C implementation
+        // ignores allocation failures here. We perform this check to avoid potential
+        // panics. We assume ~1,000 entries per file, so 72 KiB allocated per file.
+        self.0
+            .try_reserve(list_of_files.len() * 1024)
+            .map_err(|_| {
+                err!(
+                    no_mem_error(),
+                    "Couldn't allocate for {} memory",
+                    size_of::<FileEntry>() * list_of_files.len() * 1024
+                )
+            })?;
+
+        // We expect file sizes between 4KiB and 4MiB. Pre-allocate 16KiB to reduce reallocations
+        // a bit.
+        let mut buf = Vec::new();
+        buf.try_reserve(16_384)
+            .map_err(|_| err!(no_mem_error(), "Couldn't allocate for {} memory", 16_384))?;
+
+        for item in list_of_files.each() {
+            let params = RArray::from_value(item?).expect("file list was not a Ruby Array");
+            if params.len() != 4 {
+                return Err(err!(
+                    arg_error(),
+                    "wrong number of arguments {} instead of 4",
+                    params.len()
+                ));
+            }
+
+            let params = params.to_value_array::<4>()?;
+
+            let mut file_info = FileInfo::open_from_params(&params)?;
+            file_info.read_from_file(&mut buf)?;
+            self.process_buffer(file_info, &buf)?;
+        }
+        Ok(())
+    }
+
+    /// Consume the `EntryMap` and convert the key/value into`FileEntry`
+    /// objects, sorting them by their JSON strings.
+    pub fn into_sorted(self) -> Result<Vec<FileEntry>> {
+        let mut sorted = Vec::new();
+
+        // To match the behavior of the C version, pre-allocate the entries
+        // and check for allocation failure. Generally idiomatic Rust would
+        // `collect` the iterator into a new `Vec` in place, but this panics
+        // if it can't allocate and we want to continue execution in that
+        // scenario.
+        if sorted.try_reserve_exact(self.0.len()).is_err() {
+            return Err(MmapError::OutOfMemory(
+                self.0.len() * size_of::<FileEntry>(),
+            ));
+        }
+
+        sorted.extend(
+            self.0
+                .into_iter()
+                .map(|(data, meta)| FileEntry { data, meta }),
+        );
+
+        sorted.sort_unstable_by(|x, y| x.data.cmp(&y.data));
+
+        Ok(sorted)
+    }
+
+    /// Check if the `EntryMap` already contains the JSON string.
+    /// If yes, update the associated value, if not insert the
+    /// entry into the map.
+    pub fn merge_or_store(&mut self, data: BorrowedData, meta: EntryMetadata) -> Result<()> {
+        // Manually hash the `BorrowedData` and perform an equality check on the
+        // key. This allows us to perform the comparison without allocating a
+        // new `EntryData` that may not be needed.
+        let mut state = self.0.hasher().build_hasher();
+        data.hash(&mut state);
+        let hash = state.finish();
+
+        match self.0.raw_entry_mut().from_hash(hash, |k| k == &data) {
+            RawEntryMut::Vacant(entry) => {
+                // Allocate a new `EntryData` as the JSON/pid combination is
+                // not present in the map.
+                let owned = EntryData::try_from(data)?;
+                entry.insert(owned, meta);
+            }
+            RawEntryMut::Occupied(mut entry) => {
+                let existing = entry.get_mut();
+                existing.merge(&meta);
+            }
+        }
+
+        Ok(())
+    }
+
+    /// Parse metrics data from a `.db` file and store in the `EntryMap`.
+    fn process_buffer(&mut self, file_info: FileInfo, source: &[u8]) -> Result<()> {
+        if source.len() < HEADER_SIZE {
+            // Nothing to read, OK.
+            return Ok(());
+        }
+
+        // CAST: no-op on 32-bit, widening on 64-bit.
+        let used = read_u32(source, 0)? as usize;
+
+        if used > source.len() {
+            return Err(MmapError::PromParsing(format!(
+                "source file {} corrupted, used {used} > file size {}",
+                file_info.path.display(),
+                source.len()
+            )));
+        }
+
+        let mut pos = HEADER_SIZE;
+
+        while pos + size_of::<u32>() < used {
+            let raw_entry = RawEntry::from_slice(&source[pos..used])?;
+
+            if pos + raw_entry.total_len() > used {
+                return Err(MmapError::PromParsing(format!(
+                    "source file {} corrupted, used {used} < stored data length {}",
+                    file_info.path.display(),
+                    pos + raw_entry.total_len()
+                )));
+            }
+
+            let meta = EntryMetadata::new(&raw_entry, &file_info)?;
+            let data = BorrowedData::new(&raw_entry, &file_info, meta.is_pid_significant())?;
+
+            self.merge_or_store(data, meta)?;
+
+            pos += raw_entry.total_len();
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod test {
+    use magnus::Symbol;
+    use std::mem;
+
+    use super::*;
+    use crate::file_entry::FileEntry;
+    use crate::testhelper::{self, TestFile};
+
+    impl EntryData {
+        /// A helper function for tests to convert owned data to references.
+        fn as_borrowed(&self) -> BorrowedData {
+            BorrowedData {
+                json: &self.json,
+                pid: self.pid.as_deref(),
+            }
+        }
+    }
+
+    #[test]
+    fn test_into_sorted() {
+        let _cleanup = unsafe { magnus::embed::init() };
+        let ruby = magnus::Ruby::get().unwrap();
+        crate::init(&ruby).unwrap();
+
+        let entries = vec![
+            FileEntry {
+                data: EntryData {
+                    json: "zzzzzz".to_string(),
+                    pid: Some("worker-0_0".to_string()),
+                },
+                meta: EntryMetadata {
+                    multiprocess_mode: Symbol::new("max"),
+                    type_: Symbol::new("gauge"),
+                    value: 1.0,
+                },
+            },
+            FileEntry {
+                data: EntryData {
+                    json: "zzz".to_string(),
+                    pid: Some("worker-0_0".to_string()),
+                },
+                meta: EntryMetadata {
+                    multiprocess_mode: Symbol::new("max"),
+                    type_: Symbol::new("gauge"),
+                    value: 1.0,
+                },
+            },
+            FileEntry {
+                data: EntryData {
+                    json: "zzzaaa".to_string(),
+                    pid: Some("worker-0_0".to_string()),
+                },
+                meta: EntryMetadata {
+                    multiprocess_mode: Symbol::new("max"),
+                    type_: Symbol::new("gauge"),
+                    value: 1.0,
+                },
+            },
+            FileEntry {
+                data: EntryData {
+                    json: "aaa".to_string(),
+                    pid: Some("worker-0_0".to_string()),
+                },
+                meta: EntryMetadata {
+                    multiprocess_mode: Symbol::new("max"),
+                    type_: Symbol::new("gauge"),
+                    value: 1.0,
+                },
+            },
+            FileEntry {
+                data: EntryData {
+                    json: "ooo".to_string(),
+                    pid: Some("worker-1_0".to_string()),
+                },
+                meta: EntryMetadata {
+                    multiprocess_mode: Symbol::new("all"),
+                    type_: Symbol::new("gauge"),
+                    value: 1.0,
+                },
+            },
+            FileEntry {
+                data: EntryData {
+                    json: "ooo".to_string(),
+                    pid: Some("worker-0_0".to_string()),
+                },
+                meta: EntryMetadata {
+                    multiprocess_mode: Symbol::new("all"),
+                    type_: Symbol::new("gauge"),
+                    value: 1.0,
+                },
+            },
+        ];
+
+        let mut map = EntryMap::new();
+
+        for entry in entries {
+            map.0.insert(entry.data, entry.meta);
+        }
+
+        let result = map.into_sorted();
+        assert!(result.is_ok());
+        let sorted = result.unwrap();
+        assert_eq!(sorted.len(), 6);
+        assert_eq!(sorted[0].data.json, "aaa");
+        assert_eq!(sorted[1].data.json, "ooo");
+        assert_eq!(sorted[1].data.pid.as_deref(), Some("worker-0_0"));
+        assert_eq!(sorted[2].data.json, "ooo");
+        assert_eq!(sorted[2].data.pid.as_deref(), Some("worker-1_0"));
+        assert_eq!(sorted[3].data.json, "zzz");
+        assert_eq!(sorted[4].data.json, "zzzaaa");
+        assert_eq!(sorted[5].data.json, "zzzzzz");
+    }
+
+    #[test]
+    fn test_merge_or_store() {
+        let _cleanup = unsafe { magnus::embed::init() };
+        let ruby = magnus::Ruby::get().unwrap();
+        crate::init(&ruby).unwrap();
+
+        let key = "foobar";
+
+        let starting_entry = FileEntry {
+            data: EntryData {
+                json: key.to_string(),
+                pid: Some("worker-0_0".to_string()),
+            },
+            meta: EntryMetadata {
+                multiprocess_mode: Symbol::new("all"),
+                type_: Symbol::new("gauge"),
+                value: 1.0,
+            },
+        };
+
+        let matching_entry = FileEntry {
+            data: EntryData {
+                json: key.to_string(),
+                pid: Some("worker-0_0".to_string()),
+            },
+            meta: EntryMetadata {
+                multiprocess_mode: Symbol::new("all"),
+                type_: Symbol::new("gauge"),
+                value: 5.0,
+            },
+        };
+
+        let same_key_different_worker = FileEntry {
+            data: EntryData {
+                json: key.to_string(),
+                pid: Some("worker-1_0".to_string()),
+            },
+            meta: EntryMetadata {
+                multiprocess_mode: Symbol::new("all"),
+                type_: Symbol::new("gauge"),
+                value: 100.0,
+            },
+        };
+
+        let unmatched_entry = FileEntry {
+            data: EntryData {
+                json: "another key".to_string(),
+                pid: Some("worker-0_0".to_string()),
+            },
+            meta: EntryMetadata {
+                multiprocess_mode: Symbol::new("all"),
+                type_: Symbol::new("gauge"),
+                value: 1.0,
+            },
+        };
+
+        let mut map = EntryMap::new();
+
+        map.0
+            .insert(starting_entry.data.clone(), starting_entry.meta.clone());
+
+        let matching_borrowed = matching_entry.data.as_borrowed();
+        map.merge_or_store(matching_borrowed, matching_entry.meta)
+            .unwrap();
+
+        assert_eq!(
+            5.0,
+            map.0.get(&starting_entry.data).unwrap().value,
+            "value updated"
+        );
+        assert_eq!(1, map.0.len(), "no entry added");
+
+        let same_key_different_worker_borrowed = same_key_different_worker.data.as_borrowed();
+        map.merge_or_store(
+            same_key_different_worker_borrowed,
+            same_key_different_worker.meta,
+        )
+        .unwrap();
+
+        assert_eq!(
+            5.0,
+            map.0.get(&starting_entry.data).unwrap().value,
+            "value unchanged"
+        );
+
+        assert_eq!(2, map.0.len(), "additional entry added");
+
+        let unmatched_entry_borrowed = unmatched_entry.data.as_borrowed();
+        map.merge_or_store(unmatched_entry_borrowed, unmatched_entry.meta)
+            .unwrap();
+
+        assert_eq!(
+            5.0,
+            map.0.get(&starting_entry.data).unwrap().value,
+            "value unchanged"
+        );
+        assert_eq!(3, map.0.len(), "entry added");
+    }
+
+    #[test]
+    fn test_process_buffer() {
+        struct TestCase {
+            name: &'static str,
+            json: &'static [&'static str],
+            values: &'static [f64],
+            used: Option<u32>,
+            expected_ct: usize,
+            expected_err: Option<MmapError>,
+        }
+
+        let _cleanup = unsafe { magnus::embed::init() };
+        let ruby = magnus::Ruby::get().unwrap();
+        crate::init(&ruby).unwrap();
+
+        let tc = vec![
+            TestCase {
+                name: "single entry",
+                json: &[
+                    r#"["first_family","first_name",["label_a","label_b"],["value_a","value_b"]]"#,
+                ],
+                values: &[1.0],
+                used: None,
+                expected_ct: 1,
+                expected_err: None,
+            },
+            TestCase {
+                name: "multiple entries",
+                json: &[
+                    r#"["first_family","first_name",["label_a","label_b"],["value_a","value_b"]]"#,
+                    r#"["second_family","first_name",["label_a","label_b"],["value_a","value_b"]]"#,
+                ],
+                values: &[1.0, 2.0],
+                used: None,
+                expected_ct: 2,
+                expected_err: None,
+            },
+            TestCase {
+                name: "empty",
+                json: &[],
+                values: &[],
+                used: None,
+                expected_ct: 0,
+                expected_err: None,
+            },
+            TestCase {
+                name: "used too long",
+                json: &[
+                    r#"["first_family","first_name",["label_a","label_b"],["value_a","value_b"]]"#,
+                ],
+                values: &[1.0],
+                used: Some(9999),
+                expected_ct: 0,
+                expected_err: Some(MmapError::PromParsing(String::new())),
+            },
+            TestCase {
+                name: "used too short",
+                json: &[
+                    r#"["first_family","first_name",["label_a","label_b"],["value_a","value_b"]]"#,
+                ],
+                values: &[1.0],
+                used: Some(15),
+                expected_ct: 0,
+                expected_err: Some(MmapError::out_of_bounds(88, 7)),
+            },
+        ];
+
+        for case in tc {
+            let name = case.name;
+
+            let input_bytes = testhelper::entries_to_db(case.json, case.values, case.used);
+
+            let TestFile {
+                file,
+                path,
+                dir: _dir,
+            } = TestFile::new(&input_bytes);
+
+            let info = FileInfo {
+                file,
+                path,
+                len: case.json.len(),
+                multiprocess_mode: Symbol::new("max"),
+                type_: Symbol::new("gauge"),
+                pid: "worker-1".to_string(),
+            };
+
+            let mut map = EntryMap::new();
+            let result = map.process_buffer(info, &input_bytes);
+
+            assert_eq!(case.expected_ct, map.0.len(), "test case: {name} - count");
+
+            if let Some(expected_err) = case.expected_err {
+                // Validate we have the right enum type for the error. Error
+                // messages contain the temp dir path and can't be predicted
+                // exactly.
+                assert_eq!(
+                    mem::discriminant(&expected_err),
+                    mem::discriminant(&result.unwrap_err()),
+                    "test case: {name} - failure"
+                );
+            } else {
+                assert_eq!(Ok(()), result, "test case: {name} - success");
+
+                assert_eq!(
+                    case.json.len(),
+                    map.0.len(),
+                    "test case: {name} - all entries captured"
+                );
+            }
+        }
+    }
+}

data/ext/fast_mmaped_file_rs/src/metrics.proto

@@ -0,0 +1,153 @@
+// Copyright 2013 Prometheus Team
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+// http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+syntax = "proto2";
+
+package io.prometheus.client;
+option java_package = "io.prometheus.client";
+option go_package = "github.com/prometheus/client_model/go;io_prometheus_client";
+
+import "google/protobuf/timestamp.proto";
+
+message LabelPair {
+  optional string name  = 1;
+  optional string value = 2;
+}
+
+enum MetricType {
+  // COUNTER must use the Metric field "counter".
+  COUNTER = 0;
+  // GAUGE must use the Metric field "gauge".
+  GAUGE = 1;
+  // SUMMARY must use the Metric field "summary".
+  SUMMARY = 2;
+  // UNTYPED must use the Metric field "untyped".
+  UNTYPED = 3;
+  // HISTOGRAM must use the Metric field "histogram".
+  HISTOGRAM = 4;
+  // GAUGE_HISTOGRAM must use the Metric field "histogram".
+  GAUGE_HISTOGRAM = 5;
+}
+
+message Gauge {
+  optional double value = 1;
+}
+
+message Counter {
+  optional double   value    = 1;
+  optional Exemplar exemplar = 2;
+
+  optional google.protobuf.Timestamp  created_timestamp = 3;
+}
+
+message Quantile {
+  optional double quantile = 1;
+  optional double value    = 2;
+}
+
+message Summary {
+  optional uint64   sample_count = 1;
+  optional double   sample_sum   = 2;
+  repeated Quantile quantile     = 3;
+
+  optional google.protobuf.Timestamp  created_timestamp = 4;
+}
+
+message Untyped {
+  optional double value = 1;
+}
+
+message Histogram {
+  optional uint64 sample_count       = 1;
+  optional double sample_count_float = 4; // Overrides sample_count if > 0.
+  optional double sample_sum         = 2;
+  // Buckets for the conventional histogram.
+  repeated Bucket bucket             = 3; // Ordered in increasing order of upper_bound, +Inf bucket is optional.
+
+  optional google.protobuf.Timestamp created_timestamp = 15;
+
+  // Everything below here is for native histograms (also known as sparse histograms).
+  // Native histograms are an experimental feature without stability guarantees.
+
+  // schema defines the bucket schema. Currently, valid numbers are -4 <= n <= 8.
+  // They are all for base-2 bucket schemas, where 1 is a bucket boundary in each case, and
+  // then each power of two is divided into 2^n logarithmic buckets.
+  // Or in other words, each bucket boundary is the previous boundary times 2^(2^-n).
+  // In the future, more bucket schemas may be added using numbers < -4 or > 8.
+  optional sint32 schema           = 5;
+  optional double zero_threshold   = 6; // Breadth of the zero bucket.
+  optional uint64 zero_count       = 7; // Count in zero bucket.
+  optional double zero_count_float = 8; // Overrides sb_zero_count if > 0.
+
+  // Negative buckets for the native histogram.
+  repeated BucketSpan negative_span =  9;
+  // Use either "negative_delta" or "negative_count", the former for
+  // regular histograms with integer counts, the latter for float
+  // histograms.
+  repeated sint64 negative_delta    = 10; // Count delta of each bucket compared to previous one (or to zero for 1st bucket).
+  repeated double negative_count    = 11; // Absolute count of each bucket.
+
+  // Positive buckets for the native histogram.
+  // Use a no-op span (offset 0, length 0) for a native histogram without any
+  // observations yet and with a zero_threshold of 0. Otherwise, it would be
+  // indistinguishable from a classic histogram.
+  repeated BucketSpan positive_span = 12;
+  // Use either "positive_delta" or "positive_count", the former for
+  // regular histograms with integer counts, the latter for float
+  // histograms.
+  repeated sint64 positive_delta    = 13; // Count delta of each bucket compared to previous one (or to zero for 1st bucket).
+  repeated double positive_count    = 14; // Absolute count of each bucket.
+}
+
+// A Bucket of a conventional histogram, each of which is treated as
+// an individual counter-like time series by Prometheus.
+message Bucket {
+  optional uint64   cumulative_count       = 1; // Cumulative in increasing order.
+  optional double   cumulative_count_float = 4; // Overrides cumulative_count if > 0.
+  optional double   upper_bound            = 2; // Inclusive.
+  optional Exemplar exemplar               = 3;
+}
+
+// A BucketSpan defines a number of consecutive buckets in a native
+// histogram with their offset. Logically, it would be more
+// straightforward to include the bucket counts in the Span. However,
+// the protobuf representation is more compact in the way the data is
+// structured here (with all the buckets in a single array separate
+// from the Spans).
+message BucketSpan {
+  optional sint32 offset = 1; // Gap to previous span, or starting point for 1st span (which can be negative).
+  optional uint32 length = 2; // Length of consecutive buckets.
+}
+
+message Exemplar {
+  repeated LabelPair                 label     = 1;
+  optional double                    value     = 2;
+  optional google.protobuf.Timestamp timestamp = 3; // OpenMetrics-style.
+}
+
+message Metric {
+  repeated LabelPair label        = 1;
+  optional Gauge     gauge        = 2;
+  optional Counter   counter      = 3;
+  optional Summary   summary      = 4;
+  optional Untyped   untyped      = 5;
+  optional Histogram histogram    = 7;
+  optional int64     timestamp_ms = 6;
+}
+
+message MetricFamily {
+  optional string     name   = 1;
+  optional string     help   = 2;
+  optional MetricType type   = 3;
+  repeated Metric     metric = 4;
+}