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 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173
|
// Copyright 2021, The Android Open Source Project
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//! A Rust interface for the StatsD pull API.
use lazy_static::lazy_static;
use statslog_rust_header::{Atoms, Stat, StatsError};
use statspull_bindgen::*;
use std::collections::HashMap;
use std::convert::TryInto;
use std::os::raw::c_void;
use std::sync::Mutex;
/// The return value of callbacks.
pub type StatsPullResult = Vec<Box<dyn Stat>>;
/// A wrapper for AStatsManager_PullAtomMetadata.
/// It calls AStatsManager_PullAtomMetadata_release on drop.
pub struct Metadata {
metadata: *mut AStatsManager_PullAtomMetadata,
}
impl Metadata {
/// Calls AStatsManager_PullAtomMetadata_obtain.
pub fn new() -> Self {
// Safety: We panic if the memory allocation fails.
let metadata = unsafe { AStatsManager_PullAtomMetadata_obtain() };
if metadata.is_null() {
panic!("Cannot obtain pull atom metadata.");
} else {
Metadata { metadata }
}
}
/// Calls AStatsManager_PullAtomMetadata_setCoolDownMillis.
pub fn set_cooldown_millis(&mut self, cooldown_millis: i64) {
// Safety: Metadata::new ensures that self.metadata is a valid object.
unsafe { AStatsManager_PullAtomMetadata_setCoolDownMillis(self.metadata, cooldown_millis) }
}
/// Calls AStatsManager_PullAtomMetadata_getCoolDownMillis.
pub fn get_cooldown_millis(&self) -> i64 {
// Safety: Metadata::new ensures that self.metadata is a valid object.
unsafe { AStatsManager_PullAtomMetadata_getCoolDownMillis(self.metadata) }
}
/// Calls AStatsManager_PullAtomMetadata_setTimeoutMillis.
pub fn set_timeout_millis(&mut self, timeout_millis: i64) {
// Safety: Metadata::new ensures that self.metadata is a valid object.
unsafe { AStatsManager_PullAtomMetadata_setTimeoutMillis(self.metadata, timeout_millis) }
}
/// Calls AStatsManager_PullAtomMetadata_getTimeoutMillis.
pub fn get_timeout_millis(&self) -> i64 {
// Safety: Metadata::new ensures that self.metadata is a valid object.
unsafe { AStatsManager_PullAtomMetadata_getTimeoutMillis(self.metadata) }
}
/// Calls AStatsManager_PullAtomMetadata_setAdditiveFields.
pub fn set_additive_fields(&mut self, additive_fields: &mut [i32]) {
// Safety: Metadata::new ensures that self.metadata is a valid object.
unsafe {
AStatsManager_PullAtomMetadata_setAdditiveFields(
self.metadata,
additive_fields.as_mut_ptr(),
additive_fields.len().try_into().expect("Cannot convert length to i32"),
)
}
}
/// Calls AStatsManager_PullAtomMetadata_getAdditiveFields.
pub fn get_additive_fields(&self) -> Vec<i32> {
// Safety: Metadata::new ensures that self.metadata is a valid object.
// We call getNumAdditiveFields to ensure we pass getAdditiveFields a large enough array.
unsafe {
let num_fields = AStatsManager_PullAtomMetadata_getNumAdditiveFields(self.metadata)
.try_into()
.expect("Cannot convert num additive fields to usize");
let mut fields = vec![0; num_fields];
AStatsManager_PullAtomMetadata_getAdditiveFields(self.metadata, fields.as_mut_ptr());
fields
}
}
}
impl Drop for Metadata {
fn drop(&mut self) {
// Safety: Metadata::new ensures that self.metadata is a valid object.
unsafe { AStatsManager_PullAtomMetadata_release(self.metadata) }
}
}
impl Default for Metadata {
fn default() -> Self {
Self::new()
}
}
lazy_static! {
static ref COOKIES: Mutex<HashMap<i32, fn() -> StatsPullResult>> = Mutex::new(HashMap::new());
}
/// # Safety
///
/// `data` must be a valid pointer with no aliases.
unsafe extern "C" fn callback_wrapper(
atom_tag: i32,
data: *mut AStatsEventList,
_cookie: *mut c_void,
) -> AStatsManager_PullAtomCallbackReturn {
if !data.is_null() {
let map = COOKIES.lock().unwrap();
let cb = map.get(&atom_tag);
match cb {
None => log::error!("No callback found for {}", atom_tag),
Some(cb) => {
let stats = cb();
let result = stats
.iter()
// Safety: The caller promises that `data` is valid and unaliased.
.map(|stat| stat.add_astats_event(unsafe { &mut *data }))
.collect::<Result<Vec<()>, StatsError>>();
match result {
Ok(_) => {
return AStatsManager_PULL_SUCCESS as AStatsManager_PullAtomCallbackReturn
}
_ => log::error!("Error adding astats events: {:?}", result),
}
}
}
}
AStatsManager_PULL_SKIP as AStatsManager_PullAtomCallbackReturn
}
/// Rust wrapper for AStatsManager_setPullAtomCallback.
pub fn set_pull_atom_callback(
atom: Atoms,
metadata: Option<&Metadata>,
callback: fn() -> StatsPullResult,
) {
COOKIES.lock().unwrap().insert(atom as i32, callback);
let metadata_raw = match metadata {
Some(m) => m.metadata,
None => std::ptr::null_mut(),
};
// Safety: We pass a valid function as the callback.
unsafe {
AStatsManager_setPullAtomCallback(
atom as i32,
metadata_raw,
Some(callback_wrapper),
std::ptr::null_mut(),
);
}
}
/// Rust wrapper for AStatsManager_clearPullAtomCallback.
pub fn clear_pull_atom_callback(atom: Atoms) {
COOKIES.lock().unwrap().remove(&(atom as i32));
// Safety: No memory allocations.
unsafe { AStatsManager_clearPullAtomCallback(atom as i32) }
}
|