Merge "Add Rust interface for the pull API." am: 177f93f9cb am: 5b2b727eb9

Original change: https://android-review.googlesource.com/c/platform/system/core/+/1658885

Change-Id: I3888d5649ce9390299735736d0279d28eb6c4201

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",
+    ],
+}

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<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 Vec<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) }
+}

libstats/pull_rust/statslog.h

@@ -0,0 +1,3 @@
+#pragma once
+
+#include "stats_pull_atom_callback.h"