flaky_component/
lib.rs

1// Copyright (c) 2023 Graphcore Ltd. All rights reserved.
2
3//! This is an example component that will randomly drop data being passed
4//! through it.
5//!
6//! The `main.rs` in this folder shows how it can be used.
7//!
8//! # Ports
9//!
10//! This component has two ports
11//!  - One [input port](gwr_engine::port::InPort): `rx`
12//!  - One [output put port](gwr_engine::port::OutPort): `tx`
13
14// ANCHOR: use
15
16/// The `RefCell` allows the engine to be able to access state immutably as
17/// well as mutably.
18use std::cell::RefCell;
19/// The `Rc` part of the standard library brings in types used for thread
20/// synchronisation.
21use std::rc::Rc;
22
23use async_trait::async_trait;
24use gwr_components::{connect_tx, port_rx, take_option};
25use gwr_engine::engine::Engine;
26use gwr_engine::port::{InPort, OutPort, PortStateResult};
27use gwr_engine::traits::{Runnable, SimObject};
28use gwr_engine::types::{SimError, SimResult};
29use gwr_model_builder::EntityDisplay;
30/// The gwr_track library provides tracing/logging features.
31use gwr_track::entity::Entity;
32use gwr_track::trace;
33/// Random library is just used by this component to implement its drop
34/// decisions.
35use rand::rngs::StdRng;
36use rand::{RngCore, SeedableRng};
37
38// ANCHOR_END: use
39
40// ANCHOR: struct
41
42/// The overall structure for this compoment.
43///
44/// Note that in this example it is a *Generic* type in that it can be used in
45/// a simulation of any type - as long as that type implements the `SimObject`
46/// trait.
47///
48/// The `fmt::Display` trait is used when converting a component to a string
49/// for logging/printing using "{}". Simply pass through to the entity. This can
50/// be hand-written, but the `EntityDisplay` derive writes this automatically.
51#[derive(EntityDisplay)]
52pub struct Flaky<T>
53where
54    T: SimObject,
55{
56    /// Every component should include an Entity that defines where in the
57    /// overall simulation hierarchy it is. The Entity is also used to
58    /// filter logging.
59    pub entity: Rc<Entity>,
60
61    /// Store the ratio at which packets should be dropped.
62    drop_ratio: f64,
63
64    /// Random number generator used for deciding when to drop. Note that it is
65    /// wrapped in a [RefCell] which allows it to be used mutably in the `put()`
66    /// function despite the fact that the struct will be immutable.
67    rng: RefCell<StdRng>,
68
69    /// Rx port to which to send any data that hasn't been dropped.
70    /// Again, needs to be wrapped in the [RefCell] to allow it to be changed
71    /// when components are actually connected.
72    ///
73    /// Note: It is also wrapped in an [Option] so that it can be taken out in
74    /// the `run()` function.
75    rx: RefCell<Option<InPort<T>>>,
76
77    /// Tx port to which to send any data that hasn't been dropped.
78    ///
79    /// Note: It is also wrapped in an [Option] so that it can be taken out in
80    /// the `run()` function.
81    tx: RefCell<Option<OutPort<T>>>,
82}
83// ANCHOR_END: struct
84
85// ANCHOR: implFlaky
86
87/// The next thing to do is define the generic functions for the new component.
88impl<T> Flaky<T>
89where
90    T: SimObject,
91{
92    /// In this case, the `new()` function creates the component from the
93    /// parameters provided.
94    pub fn new_and_register(
95        engine: &Engine,
96        parent: &Rc<Entity>,
97        name: &str,
98        drop_ratio: f64,
99        seed: u64,
100    ) -> Result<Rc<Self>, SimError> {
101        // The entity needs to be created first because it is shared between the state
102        // and the component itself.
103        let entity = Entity::new(parent, name);
104
105        // Because it is shared it needs to be wrapped in an Rc
106        let entity = Rc::new(entity);
107
108        let rx = InPort::new(&entity, "rx");
109        let tx = OutPort::new(&entity, "tx");
110        // Finally, the top-level struct is created and wrapped in an Rc.
111        let rc_self = Rc::new(Self {
112            entity,
113            drop_ratio,
114            rng: RefCell::new(StdRng::seed_from_u64(seed)),
115            rx: RefCell::new(Some(rx)),
116            tx: RefCell::new(Some(tx)),
117        });
118        engine.register(rc_self.clone());
119        Ok(rc_self)
120    }
121
122    /// This provides the `InPort` to which you can connect
123    pub fn port_rx(&self) -> PortStateResult<T> {
124        // The `port_rx!` macro is the most consise way to access the rx port state
125        // when wrapped in `RefCell<Option<>>`.
126        port_rx!(self.rx, state)
127    }
128
129    /// The ports of this component are effectively defined by the functions
130    /// this component exposes. In this case, the `connect_port_tx` shows
131    /// that this component has an TX port which should be connected to an RX
132    /// port.
133    pub fn connect_port_tx(&self, port_state: PortStateResult<T>) -> SimResult {
134        // Because the State is immutable then we use the `connect_tx!` macro
135        // in order to simplify the setup when wrapped in `RefCell<Option<>>`.
136        connect_tx!(self.tx, connect ; port_state)
137    }
138
139    /// Return the next random u32
140    ///
141    /// This is wrapped in a separate function to hide the interior mutation
142    fn next_u32(&self) -> u32 {
143        self.rng.borrow_mut().next_u32()
144    }
145}
146
147#[async_trait(?Send)]
148impl<T> Runnable for Flaky<T>
149where
150    T: SimObject,
151{
152    async fn run(&self) -> SimResult {
153        let rx = take_option!(self.rx);
154        let tx = take_option!(self.tx);
155
156        loop {
157            // Receive a value from the input
158            let value = rx.get()?.await;
159
160            let next_u32 = self.next_u32();
161            let ratio = next_u32 as f64 / u32::MAX as f64;
162            if ratio > self.drop_ratio {
163                // Only pass on a percentage of the data
164                tx.put(value)?.await;
165            } else {
166                // Let the user know this value has been dropped.
167                trace!(self.entity ; "drop {}", value);
168            }
169        }
170    }
171}
172// ANCHOR_END: implFlaky