summaryrefslogtreecommitdiff
path: root/libstats/pull_rust/stats_pull.rs
blob: 174125e9af6516299b6b8149513a09b071492fcc (plain)
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
// 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) }
}