diff --git a/libstats/pull_rust/Android.bp b/libstats/pull_rust/Android.bp new file mode 100644 index 000000000..3660199f4 --- /dev/null +++ b/libstats/pull_rust/Android.bp @@ -0,0 +1,55 @@ +// +// Copyright (C) 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. +// + +rust_bindgen { + name: "libstatspull_bindgen", + wrapper_src: "statslog.h", + crate_name: "statspull_bindgen", + source_stem: "bindings", + bindgen_flags: [ + "--size_t-is-usize", + "--whitelist-function=AStatsEventList_addStatsEvent", + "--whitelist-function=AStatsEvent_.*", + "--whitelist-function=AStatsManager_.*", + "--whitelist-var=AStatsManager_.*", + ], + target: { + android: { + shared_libs: [ + "libstatspull", + "libstatssocket", + ], + }, + host: { + static_libs: [ + "libstatspull", + "libstatssocket", + ], + }, + }, +} + +rust_library { + name: "libstatspull_rust", + crate_name: "statspull_rust", + srcs: ["stats_pull.rs"], + rustlibs: [ + "liblazy_static", + "liblog_rust", + "libstatslog_rust_header", + "libstatspull_bindgen", + ], +} diff --git a/libstats/pull_rust/stats_pull.rs b/libstats/pull_rust/stats_pull.rs new file mode 100644 index 000000000..174125e9a --- /dev/null +++ b/libstats/pull_rust/stats_pull.rs @@ -0,0 +1,170 @@ +// 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>; + +/// 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 Vec) { + // 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 { + // 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 StatsPullResult>> = Mutex::new(HashMap::new()); +} + +// Safety: We store our callbacks in the global so they are valid. +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() + .map(|stat| stat.add_astats_event(&mut *data)) + .collect::, 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) } +} diff --git a/libstats/pull_rust/statslog.h b/libstats/pull_rust/statslog.h new file mode 100644 index 000000000..983fb7b72 --- /dev/null +++ b/libstats/pull_rust/statslog.h @@ -0,0 +1,3 @@ +#pragma once + +#include "stats_pull_atom_callback.h"