1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78

use std::time::Duration;

use tikv_jemalloc_ctl::epoch;
use tikv_jemalloc_ctl::epoch_mib;
use tikv_jemalloc_ctl::stats;

use metrics::counter;
use metrics::describe_counter;
use metrics::describe_gauge;
use metrics::gauge;
use metrics::Unit;

mod error;
pub use error::Error;

pub struct MetricRecorder {
	epoch: epoch_mib,
	active: stats::active_mib,
	allocated: stats::allocated_mib,
	mapped: stats::mapped_mib,
	metadata: stats::metadata_mib,
	resident: stats::resident_mib,
	retained: stats::retained_mib,
}

impl MetricRecorder {
	pub fn new() -> Result<Self, Error> {
		Ok(Self{
			epoch: epoch::mib()?,
			active: stats::active::mib()?,
			allocated: stats::allocated::mib()?,
			mapped: stats::mapped::mib()?,
			metadata: stats::metadata::mib()?,
			resident: stats::resident::mib()?,
			retained: stats::retained::mib()?,
		})
	}

	fn _poll(&self) -> Result<(), Error> {
		self.epoch.advance()?;
		gauge!("jemalloc_active_bytes").set(self.active.read()? as f64);
		counter!("jemalloc_allocated_bytes").absolute(self.allocated.read()? as u64);
		gauge!("jemalloc_mapped_bytes").set(self.mapped.read()? as f64);
		gauge!("jemalloc_metadata_bytes").set(self.metadata.read()? as f64);
		gauge!("jemalloc_resident_bytes").set(self.resident.read()? as f64);
		gauge!("jemalloc_retained_bytes").set(self.retained.read()? as f64);
		Ok(())
	}

	#[inline]
	pub fn poll(&self) {
		if let Err(error) = self._poll() {
			tracing::warn!(%error, "Failed to poll jemalloc stats");
		}
	}

	pub fn start(self) -> tokio::task::JoinHandle<()> {
		tokio::task::spawn(async move {
			let mut interval = tokio::time::interval(Duration::from_secs(10));
			loop {
				self.poll();
				interval.tick().await;
			}
		})
	}
}

pub fn init() -> Result<tokio::task::JoinHandle<()>, Error> {
	describe_gauge!("jemalloc_active_bytes", Unit::Bytes, "Total number of bytes in active pages allocated by the process");
	describe_counter!("jemalloc_allocated_bytes", Unit::Bytes, "Total number of bytes allocated by the process");
	describe_gauge!("jemalloc_mapped_bytes", Unit::Bytes, "Total number of bytes in active extents mapped by the allocator");
	describe_gauge!("jemalloc_metadata_bytes", Unit::Bytes, "Total number of bytes dedicated to jemalloc metadata");
	describe_gauge!("jemalloc_resident_bytes", Unit::Bytes, "Total number of bytes in physically resident data pages mapped by the allocator");
	describe_gauge!("jemalloc_retained_bytes", Unit::Bytes, "Total number of bytes in virtual memory mappings that were retained rather than being returned to the operating system");
	Ok(MetricRecorder::new()?.start())
}