1 // Copyright 2021, The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 // http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14
15 //! A Rust interface for the StatsD pull API.
16
17 use statslog_rust_header::{Atoms, Stat, StatsError};
18 use statspull_bindgen::*;
19 use std::collections::HashMap;
20 use std::convert::TryInto;
21 use std::os::raw::c_void;
22 use std::sync::{LazyLock, Mutex};
23
24 /// The return value of callbacks.
25 pub type StatsPullResult = Vec<Box<dyn Stat>>;
26
27 /// A wrapper for AStatsManager_PullAtomMetadata.
28 /// It calls AStatsManager_PullAtomMetadata_release on drop.
29 pub struct Metadata {
30 metadata: *mut AStatsManager_PullAtomMetadata,
31 }
32
33 impl Metadata {
34 /// Calls AStatsManager_PullAtomMetadata_obtain.
new() -> Self35 pub fn new() -> Self {
36 // Safety: We panic if the memory allocation fails.
37 let metadata = unsafe { AStatsManager_PullAtomMetadata_obtain() };
38 if metadata.is_null() {
39 panic!("Cannot obtain pull atom metadata.");
40 } else {
41 Metadata { metadata }
42 }
43 }
44
45 /// Calls AStatsManager_PullAtomMetadata_setCoolDownMillis.
set_cooldown_millis(&mut self, cooldown_millis: i64)46 pub fn set_cooldown_millis(&mut self, cooldown_millis: i64) {
47 // Safety: Metadata::new ensures that self.metadata is a valid object.
48 unsafe { AStatsManager_PullAtomMetadata_setCoolDownMillis(self.metadata, cooldown_millis) }
49 }
50
51 /// Calls AStatsManager_PullAtomMetadata_getCoolDownMillis.
get_cooldown_millis(&self) -> i6452 pub fn get_cooldown_millis(&self) -> i64 {
53 // Safety: Metadata::new ensures that self.metadata is a valid object.
54 unsafe { AStatsManager_PullAtomMetadata_getCoolDownMillis(self.metadata) }
55 }
56
57 /// Calls AStatsManager_PullAtomMetadata_setTimeoutMillis.
set_timeout_millis(&mut self, timeout_millis: i64)58 pub fn set_timeout_millis(&mut self, timeout_millis: i64) {
59 // Safety: Metadata::new ensures that self.metadata is a valid object.
60 unsafe { AStatsManager_PullAtomMetadata_setTimeoutMillis(self.metadata, timeout_millis) }
61 }
62
63 /// Calls AStatsManager_PullAtomMetadata_getTimeoutMillis.
get_timeout_millis(&self) -> i6464 pub fn get_timeout_millis(&self) -> i64 {
65 // Safety: Metadata::new ensures that self.metadata is a valid object.
66 unsafe { AStatsManager_PullAtomMetadata_getTimeoutMillis(self.metadata) }
67 }
68
69 /// Calls AStatsManager_PullAtomMetadata_setAdditiveFields.
set_additive_fields(&mut self, additive_fields: &mut [i32])70 pub fn set_additive_fields(&mut self, additive_fields: &mut [i32]) {
71 // Safety: Metadata::new ensures that self.metadata is a valid object.
72 unsafe {
73 AStatsManager_PullAtomMetadata_setAdditiveFields(
74 self.metadata,
75 additive_fields.as_mut_ptr(),
76 additive_fields.len().try_into().expect("Cannot convert length to i32"),
77 )
78 }
79 }
80
81 /// Calls AStatsManager_PullAtomMetadata_getAdditiveFields.
get_additive_fields(&self) -> Vec<i32>82 pub fn get_additive_fields(&self) -> Vec<i32> {
83 // Safety: Metadata::new ensures that self.metadata is a valid object.
84 // We call getNumAdditiveFields to ensure we pass getAdditiveFields a large enough array.
85 unsafe {
86 let num_fields = AStatsManager_PullAtomMetadata_getNumAdditiveFields(self.metadata)
87 .try_into()
88 .expect("Cannot convert num additive fields to usize");
89 let mut fields = vec![0; num_fields];
90 AStatsManager_PullAtomMetadata_getAdditiveFields(self.metadata, fields.as_mut_ptr());
91 fields
92 }
93 }
94 }
95
96 impl Drop for Metadata {
drop(&mut self)97 fn drop(&mut self) {
98 // Safety: Metadata::new ensures that self.metadata is a valid object.
99 unsafe { AStatsManager_PullAtomMetadata_release(self.metadata) }
100 }
101 }
102
103 impl Default for Metadata {
default() -> Self104 fn default() -> Self {
105 Self::new()
106 }
107 }
108
109 static COOKIES: LazyLock<Mutex<HashMap<i32, fn() -> StatsPullResult>>> =
110 LazyLock::new(|| Mutex::new(HashMap::new()));
111
112 /// # Safety
113 ///
114 /// `data` must be a valid pointer with no aliases.
callback_wrapper( atom_tag: i32, data: *mut AStatsEventList, _cookie: *mut c_void, ) -> AStatsManager_PullAtomCallbackReturn115 unsafe extern "C" fn callback_wrapper(
116 atom_tag: i32,
117 data: *mut AStatsEventList,
118 _cookie: *mut c_void,
119 ) -> AStatsManager_PullAtomCallbackReturn {
120 if !data.is_null() {
121 let map = COOKIES.lock().unwrap();
122 let cb = map.get(&atom_tag);
123 match cb {
124 None => log::error!("No callback found for {}", atom_tag),
125 Some(cb) => {
126 let stats = cb();
127 let result = stats
128 .iter()
129 // Safety: The caller promises that `data` is valid and unaliased.
130 .map(|stat| stat.add_astats_event(unsafe { &mut *data }))
131 .collect::<Result<Vec<()>, StatsError>>();
132 match result {
133 Ok(_) => {
134 return AStatsManager_PULL_SUCCESS as AStatsManager_PullAtomCallbackReturn
135 }
136 _ => log::error!("Error adding astats events: {:?}", result),
137 }
138 }
139 }
140 }
141 AStatsManager_PULL_SKIP as AStatsManager_PullAtomCallbackReturn
142 }
143
144 /// Rust wrapper for AStatsManager_setPullAtomCallback.
set_pull_atom_callback( atom: Atoms, metadata: Option<&Metadata>, callback: fn() -> StatsPullResult, )145 pub fn set_pull_atom_callback(
146 atom: Atoms,
147 metadata: Option<&Metadata>,
148 callback: fn() -> StatsPullResult,
149 ) {
150 COOKIES.lock().unwrap().insert(atom as i32, callback);
151 let metadata_raw = match metadata {
152 Some(m) => m.metadata,
153 None => std::ptr::null_mut(),
154 };
155 // Safety: We pass a valid function as the callback.
156 unsafe {
157 AStatsManager_setPullAtomCallback(
158 atom as i32,
159 metadata_raw,
160 Some(callback_wrapper),
161 std::ptr::null_mut(),
162 );
163 }
164 }
165
166 /// Rust wrapper for AStatsManager_clearPullAtomCallback.
clear_pull_atom_callback(atom: Atoms)167 pub fn clear_pull_atom_callback(atom: Atoms) {
168 COOKIES.lock().unwrap().remove(&(atom as i32));
169 // Safety: No memory allocations.
170 unsafe { AStatsManager_clearPullAtomCallback(atom as i32) }
171 }
172