rs-zero 0.2.6

Rust-first microservice framework inspired by go-zero engineering practices
Documentation 
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
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103

use std::collections::BTreeMap;

use pyroscope::{
    backend::{BackendConfig, PprofConfig, pprof_backend},
    pyroscope::PyroscopeAgentBuilder,
};

use crate::{
    core::{CoreError, CoreResult},
    profiling::{ProfilingConfig, ProfilingHandle, PyroscopeAgentConfig, PyroscopeConfig},
};

const SPY_NAME: &str = "pyroscope-rs";

/// Normalizes Pyroscope configuration before building or starting an agent.
pub fn build_pyroscope_agent_config(config: PyroscopeConfig) -> CoreResult<PyroscopeAgentConfig> {
    let endpoint = config.endpoint.trim().trim_end_matches('/').to_string();
    if endpoint.is_empty() {
        return Err(CoreError::Profiling(
            "pyroscope endpoint is empty".to_string(),
        ));
    }
    let service_name = config.service_name.trim().to_string();
    if service_name.is_empty() {
        return Err(CoreError::Profiling(
            "pyroscope service name is empty".to_string(),
        ));
    }
    let tags = normalize_tags(config.tags);
    let sample_rate = if config.sample_rate == 0 {
        100
    } else {
        config.sample_rate
    };

    Ok(PyroscopeAgentConfig {
        endpoint,
        service_name,
        sample_rate,
        tags,
        shutdown_timeout: config.shutdown_timeout,
    })
}

/// Starts profiling when enabled, otherwise returns a no-op handle.
pub fn start_profiling(config: ProfilingConfig) -> CoreResult<ProfilingHandle> {
    if !config.enabled {
        return Ok(ProfilingHandle::disabled());
    }
    start_pyroscope(config.pyroscope.unwrap_or_default())
}

/// Starts a Pyroscope agent with the pprof-rs backend.
pub fn start_pyroscope(config: PyroscopeConfig) -> CoreResult<ProfilingHandle> {
    let normalized = build_pyroscope_agent_config(config)?;
    let backend = pprof_backend(
        PprofConfig {
            sample_rate: normalized.sample_rate,
        },
        BackendConfig {
            report_thread_id: false,
            report_thread_name: true,
            report_pid: true,
        },
    );
    let tag_pairs = normalized
        .tags
        .iter()
        .map(|(key, value)| (key.as_str(), value.as_str()))
        .collect::<Vec<_>>();
    let agent = PyroscopeAgentBuilder::new(
        &normalized.endpoint,
        &normalized.service_name,
        normalized.sample_rate,
        SPY_NAME,
        env!("CARGO_PKG_VERSION"),
        backend,
    )
    .tags(tag_pairs)
    .build()
    .map_err(|error| CoreError::Profiling(error.to_string()))?
    .start()
    .map_err(|error| CoreError::Profiling(error.to_string()))?;

    Ok(ProfilingHandle::pyroscope(
        agent,
        normalized.shutdown_timeout,
    ))
}

fn normalize_tags(tags: BTreeMap<String, String>) -> BTreeMap<String, String> {
    tags.into_iter()
        .filter_map(|(key, value)| {
            let key = key.trim();
            let value = value.trim();
            if key.is_empty() || value.is_empty() {
                None
            } else {
                Some((key.to_string(), value.to_string()))
            }
        })
        .collect()
}