Skip to main content

keyhog_profile/
session.rs

1use crate::collector::SnapshotCollector;
2use crate::hardware::HardwareSession;
3use crate::resources::{resource_usage, state_measurements, ProcessResourceCollector};
4use crate::runtime::{ContextGuard, Runtime};
5use crate::schema::{
6    ResourceSample, ResourceSnapshot, RunIdentity, RunProfile, RunState, StateTransition,
7    PROFILE_SCHEMA,
8};
9use crate::system::SystemSession;
10use std::fmt;
11use std::time::Instant;
12
13/// Reserved error type for profile-session initialization failures.
14#[derive(Clone, Copy, Debug, Eq, PartialEq)]
15pub struct SessionActive;
16
17impl fmt::Display for SessionActive {
18    fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
19        formatter.write_str("a KeyHog profile session could not be initialized")
20    }
21}
22
23impl std::error::Error for SessionActive {}
24
25/// One causal profiling session with isolated owned metric storage.
26pub struct Session {
27    identity: Option<RunIdentity>,
28    runtime: Runtime,
29    context: Option<ContextGuard>,
30    started: Instant,
31    resources_at_start: ResourceSnapshot,
32    resource_collector: ProcessResourceCollector,
33    transitions: Vec<StateTransition>,
34    resource_samples: Vec<ResourceSample>,
35    hardware: Option<HardwareSession>,
36    system: Option<SystemSession>,
37    finished: bool,
38}
39
40impl Session {
41    /// Start a fresh isolated session.
42    pub fn start(identity: RunIdentity) -> Result<Self, SessionActive> {
43        let started = Instant::now();
44        let runtime = Runtime::new_at(started);
45        let context = runtime.enter();
46        let mut resource_collector = ProcessResourceCollector::new();
47        let resources_at_start = resource_collector.sample();
48        Ok(Self {
49            identity: Some(identity),
50            runtime,
51            context: Some(context),
52            started,
53            resources_at_start,
54            resource_collector,
55            transitions: vec![StateTransition {
56                version: crate::schema::STATE_TRANSITION_VERSION,
57                state: RunState::Created,
58                elapsed_ns: 0,
59            }],
60            resource_samples: vec![ResourceSample {
61                version: crate::schema::RESOURCE_SAMPLE_VERSION,
62                state: RunState::Created,
63                elapsed_ns: 0,
64                snapshot: resources_at_start,
65            }],
66            hardware: Some(HardwareSession::new()),
67            system: Some(SystemSession::new()),
68            finished: false,
69        })
70    }
71
72    /// Clone the runtime handle for propagation to a worker or async task.
73    pub fn runtime(&self) -> Runtime {
74        self.runtime.clone()
75    }
76
77    /// Mutate run identity before the session is finalized.
78    pub fn identity_mut(&mut self) -> &mut RunIdentity {
79        self.identity
80            .as_mut()
81            .expect("unfinished profile owns identity")
82    }
83
84    /// Record an explicit macro state transition.
85    pub fn transition(&mut self, state: RunState) {
86        let elapsed_ns = u64::try_from(self.started.elapsed().as_nanos()).unwrap_or(u64::MAX);
87        self.transitions.push(StateTransition {
88            version: crate::schema::STATE_TRANSITION_VERSION,
89            state,
90            elapsed_ns,
91        });
92        self.resource_samples.push(ResourceSample {
93            version: crate::schema::RESOURCE_SAMPLE_VERSION,
94            state,
95            elapsed_ns,
96            snapshot: self.resource_collector.sample(),
97        });
98        if let Some(hardware) = self.hardware.as_mut() {
99            hardware.transition_sample();
100        }
101    }
102
103    /// Finish the session and return its complete structured record.
104    pub fn finish(mut self, status: RunState) -> RunProfile {
105        self.transition(status);
106        let wall = self.started.elapsed();
107        let finish_resources = self
108            .resource_samples
109            .last()
110            .map_or(self.resources_at_start, |sample| sample.snapshot);
111        if let Some(resident_bytes) = finish_resources.resident_bytes {
112            self.runtime
113                .set_gauge(crate::GaugeId::ResidentMemory, resident_bytes);
114        }
115        if let Some(virtual_bytes) = finish_resources.virtual_bytes {
116            self.runtime
117                .set_gauge(crate::GaugeId::VirtualMemory, virtual_bytes);
118        }
119        if let Some(thread_count) = finish_resources.thread_count {
120            self.runtime
121                .set_gauge(crate::GaugeId::ProcessThreads, thread_count);
122        }
123        if let (Some(start_ms), Some(finish_ms)) = (
124            self.resources_at_start.cpu_time_ms,
125            finish_resources.cpu_time_ms,
126        ) {
127            self.runtime.add_counter(
128                crate::CounterId::ProcessCpuTime,
129                finish_ms.saturating_sub(start_ms),
130            );
131        }
132        self.context.take();
133        let (input_bytes, input_units) = self.runtime.take_session_input_totals();
134        let (derived_decoder_bytes, backend_dispatched_bytes) =
135            self.runtime.take_session_workload_totals();
136        let stages = self.runtime.take_session_stage_measurements();
137        let resource_samples = std::mem::take(&mut self.resource_samples);
138        let states = state_measurements(&self.transitions, &resource_samples);
139        let resources = resource_usage(
140            self.resources_at_start,
141            finish_resources,
142            wall,
143            &resource_samples,
144        );
145        let wall_ns = u64::try_from(wall.as_nanos()).unwrap_or(u64::MAX);
146        let hardware_session = self
147            .hardware
148            .take()
149            .expect("unfinished profile owns hardware");
150        let mut collectors = vec![self.resource_collector.capability()];
151        collectors.extend(hardware_session.capabilities());
152        let hardware = hardware_session.finish_evidence(wall_ns, &self.runtime);
153        let system_session = self.system.take().expect("unfinished profile owns system");
154        collectors.extend(system_session.capabilities());
155        let system = system_session.finish_evidence(
156            &self.runtime,
157            &finish_resources,
158            input_bytes,
159            derived_decoder_bytes,
160        );
161        let profile = RunProfile {
162            version: crate::schema::RUN_PROFILE_VERSION,
163            schema: PROFILE_SCHEMA.to_string(),
164            identity: self
165                .identity
166                .take()
167                .expect("unfinished profile owns identity"),
168            status,
169            wall_time_ns: wall_ns,
170            input_bytes,
171            input_units,
172            workload: crate::schema::WorkloadMeasurements::measured(
173                derived_decoder_bytes,
174                backend_dispatched_bytes,
175            ),
176            stages,
177            transitions: std::mem::take(&mut self.transitions),
178            states,
179            collectors,
180            resource_samples,
181            resources,
182            hardware,
183            system,
184        };
185        self.finished = true;
186        profile
187    }
188}
189
190impl Drop for Session {
191    fn drop(&mut self) {
192        if !self.finished {
193            self.context.take();
194        }
195    }
196}