libstats/pull_rust/stats_pull.rs

// 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: 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::<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) }
}