Skip to main content

ios_core/services/instruments/
energy.rs

1use tokio::io::{AsyncRead, AsyncWrite};
2
3use crate::services::dtx::codec::{DtxConnection, DtxError};
4use crate::services::dtx::primitive_enc::archived_object;
5use crate::services::dtx::types::DtxPayload;
6
7#[derive(Debug, Clone, PartialEq, serde::Serialize)]
8pub struct EnergySample {
9    pub payload: serde_json::Value,
10}
11
12pub struct EnergyMonitorClient<S> {
13    conn: DtxConnection<S>,
14    channel_code: i32,
15    pids: Vec<i32>,
16}
17
18impl<S: AsyncRead + AsyncWrite + Unpin + Send> EnergyMonitorClient<S> {
19    pub async fn connect(stream: S, pids: &[i32]) -> Result<Self, DtxError> {
20        let mut conn = DtxConnection::new(stream);
21        let channel_code = conn.request_channel(super::ENERGY_MONITOR_SVC).await?;
22        let mut client = Self {
23            conn,
24            channel_code,
25            pids: pids.to_vec(),
26        };
27        client.stop_sampling().await?;
28        client.start_sampling().await?;
29        tokio::time::sleep(std::time::Duration::from_millis(500)).await;
30        Ok(client)
31    }
32
33    pub async fn sample(&mut self) -> Result<EnergySample, DtxError> {
34        let response = self
35            .conn
36            .method_call(
37                self.channel_code,
38                "sampleAttributes:forPIDs:",
39                &[
40                    archived_object(crate::proto::nskeyedarchiver_encode::archive_dict(vec![])),
41                    archived_object(encode_pid_array(&self.pids)),
42                ],
43            )
44            .await?;
45        parse_energy_sample(&response.payload)
46    }
47
48    pub async fn stop_sampling(&mut self) -> Result<(), DtxError> {
49        self.conn
50            .method_call_async(
51                self.channel_code,
52                "stopSamplingForPIDs:",
53                &[archived_object(encode_pid_array(&self.pids))],
54            )
55            .await
56    }
57
58    async fn start_sampling(&mut self) -> Result<(), DtxError> {
59        self.conn
60            .method_call_async(
61                self.channel_code,
62                "startSamplingForPIDs:",
63                &[archived_object(encode_pid_array(&self.pids))],
64            )
65            .await
66    }
67}
68
69fn encode_pid_array(pids: &[i32]) -> Vec<u8> {
70    crate::proto::nskeyedarchiver_encode::archive_array(
71        pids.iter()
72            .map(|pid| plist::Value::Integer((*pid as i64).into()))
73            .collect(),
74    )
75}
76
77fn parse_energy_sample(payload: &DtxPayload) -> Result<EnergySample, DtxError> {
78    let payload = match payload {
79        DtxPayload::Response(value) => value.clone(),
80        DtxPayload::Raw(bytes) => super::unarchive_raw_payload(bytes).ok_or_else(|| {
81            DtxError::Protocol("energy response was not a valid archived payload".into())
82        })?,
83        other => {
84            return Err(DtxError::Protocol(format!(
85                "unexpected energy response payload: {other:?}"
86            )))
87        }
88    };
89
90    Ok(EnergySample {
91        payload: super::nsobject_to_json(&payload),
92    })
93}
94
95#[cfg(test)]
96mod tests {
97    use indexmap::IndexMap;
98
99    use super::*;
100    use crate::services::dtx::types::NSObject;
101
102    #[test]
103    fn parses_energy_sample_from_response_dict() {
104        let payload = DtxPayload::Response(NSObject::Dict(IndexMap::from_iter([
105            ("energy".to_string(), NSObject::Double(17.5)),
106            (
107                "processes".to_string(),
108                NSObject::Array(vec![NSObject::Dict(IndexMap::from_iter([(
109                    "pid".to_string(),
110                    NSObject::Int(42),
111                )]))]),
112            ),
113        ])));
114
115        let sample = parse_energy_sample(&payload).unwrap();
116        assert_eq!(sample.payload["energy"], serde_json::json!(17.5));
117        assert_eq!(sample.payload["processes"][0]["pid"], serde_json::json!(42));
118    }
119
120    #[test]
121    fn encodes_pid_array() {
122        let encoded = encode_pid_array(&[12, 34]);
123        let decoded = crate::proto::nskeyedarchiver::unarchive(&encoded).unwrap();
124        let values = decoded.as_array().unwrap();
125        assert_eq!(values[0].as_int(), Some(12));
126        assert_eq!(values[1].as_int(), Some(34));
127    }
128}