netstack3_base/
counters.rs

1// Copyright 2024 The Fuchsia Authors. All rights reserved.
2// Use of this source code is governed by a BSD-style license that can be
3// found in the LICENSE file.
4
5//! Common counter abstractions.
6
7use core::fmt::Debug;
8use core::sync::atomic::{AtomicU64, Ordering};
9
10use crate::test_only::TestOnlyPartialEq;
11
12/// An atomic counter for packet statistics, e.g. IPv4 packets received.
13#[derive(Debug, Default)]
14pub struct Counter(AtomicU64);
15
16impl Counter {
17    /// Increments the counter value by 1.
18    pub fn increment(&self) {
19        // Use relaxed ordering since we do not use packet counter values to
20        // synchronize other accesses.  See:
21        // https://doc.rust-lang.org/nomicon/atomics.html#relaxed
22        let Self(v) = self;
23        let _: u64 = v.fetch_add(1, Ordering::Relaxed);
24    }
25
26    /// Adds the provided value to the counter.
27    pub fn add(&self, n: u64) {
28        // Use relaxed ordering since we do not use packet counter values to
29        // synchronize other accesses.  See:
30        // https://doc.rust-lang.org/nomicon/atomics.html#relaxed
31        let Self(v) = self;
32        let _: u64 = v.fetch_add(n, Ordering::Relaxed);
33    }
34
35    /// Atomically retrieves the counter value as a `u64`.
36    pub fn get(&self) -> u64 {
37        // Use relaxed ordering since we do not use packet counter values to
38        // synchronize other accesses.  See:
39        // https://doc.rust-lang.org/nomicon/atomics.html#relaxed
40        let Self(v) = self;
41        v.load(Ordering::Relaxed)
42    }
43}
44
45/// A context that stores counters.
46///
47/// `CounterContext` exposes access to counters for observation and debugging.
48pub trait CounterContext<T> {
49    /// Returns a reference to the counters.
50    fn counters(&self) -> &T;
51}
52
53/// A context that provides access to per-resource counters for observation and
54/// debugging.
55pub trait ResourceCounterContext<R, T>: CounterContext<T> {
56    /// Returns a reference to the set of counters on `resource`.
57    fn per_resource_counters<'a>(&'a self, resource: &'a R) -> &'a T;
58
59    /// Increments both the per-resource and stackwide versions of
60    /// the counter returned by the callback.
61    fn increment_both<F: Fn(&T) -> &Counter>(&self, resource: &R, cb: F) {
62        cb(self.per_resource_counters(resource)).increment();
63        cb(self.counters()).increment();
64    }
65
66    /// Adds `value` to both the per-resource and stackwide versions of
67    /// the counter returned by the callback.
68    fn add_both<F: Fn(&T) -> &Counter>(&self, resource: &R, value: u64, cb: F) {
69        cb(self.per_resource_counters(resource)).add(value);
70        cb(self.counters()).add(value);
71    }
72
73    /// Like `add_both` but takes a `usize` value, saturating to `u64::max` if
74    /// conversion fails.
75    fn add_both_usize<F: Fn(&T) -> &Counter>(&self, resource: &R, value: usize, cb: F) {
76        self.add_both(resource, value.try_into().unwrap_or(u64::MAX), cb);
77    }
78}
79
80mod sealed {
81    use super::*;
82
83    /// Used to implement the "sealed trait" pattern.
84    pub trait Sealed {}
85
86    impl Sealed for u64 {}
87    impl Sealed for Counter {}
88}
89
90/// A marker trait to indicate types that may be used as a counter.
91pub trait CounterRepr: sealed::Sealed + Default + Debug + TestOnlyPartialEq {
92    /// Get the held counter value.
93    fn get(&self) -> u64;
94    /// Construct a new counter from the value.
95    fn new(value: u64) -> Self;
96
97    /// Convert one `CounterRepr` into another `CounterRepr`.
98    fn cast<C: CounterRepr>(&self) -> C {
99        C::new(self.get())
100    }
101}
102
103impl CounterRepr for Counter {
104    fn get(&self) -> u64 {
105        self.get()
106    }
107
108    fn new(value: u64) -> Self {
109        Counter(AtomicU64::new(value))
110    }
111}
112
113// Only allow `u64` as a counter in tests.
114#[cfg(any(test, feature = "testutils"))]
115impl CounterRepr for u64 {
116    fn get(&self) -> u64 {
117        *self
118    }
119
120    fn new(value: u64) -> Self {
121        value
122    }
123}
124
125/// The specification of a [`CounterCollection`].
126pub trait CounterCollectionSpec {
127    /// The underlying [`CounterCollection`] for this specification.
128    type CounterCollection<C: CounterRepr>: CounterCollection<Spec = Self>;
129    /// A utility method to transform the underlying counter representation
130    /// within [`Self::CounterCollection`].
131    fn transform<C1: CounterRepr, C2: CounterRepr>(
132        counters: &Self::CounterCollection<C1>,
133    ) -> Self::CounterCollection<C2>;
134}
135
136/// A collection of counters.
137pub trait CounterCollection: Debug + Default + TestOnlyPartialEq {
138    /// The [`CounterCollectionSpec`] associated with this collection.
139    type Spec: CounterCollectionSpec<CounterCollection<Self::Repr> = Self>;
140    /// The counter representation held by this collection.
141    type Repr: CounterRepr;
142    /// A utility method to cast this collection of counters into a different
143    /// underlying counter representation.
144    fn cast<C: CounterRepr>(&self) -> <Self::Spec as CounterCollectionSpec>::CounterCollection<C> {
145        <Self::Spec as CounterCollectionSpec>::transform::<Self::Repr, C>(self)
146    }
147}
netstack3_base/counters.rs

netstack3_base/
counters.rs
