Skip to main content

aranya_runtime/
client.rs

1use core::{fmt, iter::DoubleEndedIterator};
2
3use buggy::Bug;
4use tracing::error;
5
6use crate::{
7    Address, CmdId, Command, GraphId, LocatedAddress, PeerCache, Perspective as _, Policy,
8    PolicyError, PolicyStore, Sink, Storage as _, StorageError, StorageProvider, TraversalBuffer,
9    policy::ActionPlacement, storage::Spill,
10};
11
12pub(crate) mod braiding; // exposed for `buffers.rs` (StrandHeap)
13mod buffers;
14pub(crate) mod convergence_map; // exposed for `buffers.rs` (ConvergenceStorage)
15mod session;
16mod transaction;
17
18pub(crate) use buffers::BraidBuffer;
19pub use buffers::RuntimeBuffers;
20
21pub use self::{session::Session, transaction::Transaction};
22
23/// An error returned by the runtime client.
24#[derive(Debug, thiserror::Error)]
25#[non_exhaustive]
26pub enum ClientError {
27    #[error("no such parent: {0}")]
28    NoSuchParent(CmdId),
29    #[error("policy error: {0}")]
30    PolicyError(PolicyError),
31    #[error("storage error: {0}")]
32    StorageError(#[from] StorageError),
33    #[error("init error")]
34    InitError,
35    #[error("not authorized")]
36    NotAuthorized,
37    #[error("could not deserialize session command")]
38    SessionDeserialize,
39    /// Attempted to braid two parallel finalize commands together.
40    ///
41    /// Policy must be designed such that two parallel finalize commands are never produced.
42    ///
43    /// Currently, this is practically an unrecoverable error. You must wipe all graphs containing
44    /// the "bad" finalize command and resync from the "good" clients. Otherwise, your network will
45    /// split into two separate graph states which can never successfully sync.
46    #[error("found parallel finalize commands during braid")]
47    ParallelFinalize,
48    #[error("concurrent transaction usage")]
49    ConcurrentTransaction,
50    #[error(transparent)]
51    Bug(#[from] Bug),
52}
53
54impl From<PolicyError> for ClientError {
55    fn from(error: PolicyError) -> Self {
56        match error {
57            PolicyError::Check => Self::NotAuthorized,
58            _ => Self::PolicyError(error),
59        }
60    }
61}
62
63/// Keeps track of client graph state.
64///
65/// - `PS` should be an implementation of [`PolicyStore`].
66/// - `SP` should be an implementation of [`StorageProvider`].
67pub struct ClientState<PS, SP> {
68    policy_store: PS,
69    provider: SP,
70}
71
72// Manual Debug impl to exclude `buffers` (large, not useful in debug output).
73impl<PS: fmt::Debug, SP: fmt::Debug> fmt::Debug for ClientState<PS, SP> {
74    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
75        f.debug_struct("ClientState")
76            .field("policy_store", &self.policy_store)
77            .field("provider", &self.provider)
78            .finish_non_exhaustive()
79    }
80}
81
82impl<PS, SP> ClientState<PS, SP> {
83    /// Creates a `ClientState`.
84    pub const fn new(policy_store: PS, provider: SP) -> Self {
85        Self {
86            policy_store,
87            provider,
88        }
89    }
90
91    /// Provide access to the [`StorageProvider`].
92    pub fn provider(&mut self) -> &mut SP {
93        &mut self.provider
94    }
95}
96
97impl<PS, SP> ClientState<PS, SP>
98where
99    PS: PolicyStore,
100    SP: StorageProvider,
101{
102    /// Create a new graph (AKA Team). This graph will start with the initial policy
103    /// provided which must be compatible with the policy store PS. The `payload` is the initial
104    /// init message that will bootstrap the graph facts. Effects produced when processing
105    /// the payload are emitted to the sink.
106    pub fn new_graph(
107        &mut self,
108        policy_data: &[u8],
109        action: <PS::Policy as Policy>::Action<'_>,
110        sink: &mut impl Sink<PS::Effect>,
111    ) -> Result<GraphId, ClientError> {
112        let policy_id = self.policy_store.add_policy(policy_data)?;
113        let policy = self.policy_store.get_policy(policy_id)?;
114
115        let mut perspective = self.provider.new_perspective(policy_id);
116        sink.begin();
117        policy
118            .call_action(action, &mut perspective, sink, ActionPlacement::OnGraph)
119            .inspect_err(|_| sink.rollback())?;
120        sink.commit();
121
122        let (graph_id, _) = self.provider.new_storage(perspective)?;
123
124        Ok(graph_id)
125    }
126
127    /// Remove a graph (AKA Team). The graph commands will be removed from storage.
128    pub fn remove_graph(&mut self, graph_id: GraphId) -> Result<(), ClientError> {
129        self.provider.remove_storage(graph_id)?;
130
131        Ok(())
132    }
133
134    /// Commit the [`Transaction`] to storage, after merging all temporary heads.
135    ///
136    /// `make_spill` is called whenever an internal braid or convergence map
137    /// needs byte-addressable overflow storage. It is typically
138    /// `|| LibcSpill::new(&dir)` for production callers or `MemSpill::new`
139    /// for tests.
140    ///
141    /// Returns whether any new commands were added.
142    pub fn commit<F, MS>(
143        &mut self,
144        trx: Transaction<SP, PS>,
145        sink: &mut impl Sink<PS::Effect>,
146        buffers: &mut RuntimeBuffers<SP::Segment>,
147        make_spill: MS,
148    ) -> Result<bool, ClientError>
149    where
150        F: Spill,
151        MS: Fn() -> Result<F, StorageError>,
152    {
153        trx.commit::<F, MS>(
154            &mut self.provider,
155            &mut self.policy_store,
156            sink,
157            buffers,
158            &make_spill,
159        )
160    }
161
162    /// Add commands to the transaction, writing the results to
163    /// `sink`.
164    ///
165    /// `make_spill` is called whenever an internal braid or convergence map
166    /// needs byte-addressable overflow storage. See [`Self::commit`].
167    ///
168    /// Returns the number of commands that were added.
169    pub fn add_commands<F, MS>(
170        &mut self,
171        trx: &mut Transaction<SP, PS>,
172        sink: &mut impl Sink<PS::Effect>,
173        commands: &[impl Command],
174        buffers: &mut RuntimeBuffers<SP::Segment>,
175        make_spill: MS,
176    ) -> Result<usize, ClientError>
177    where
178        F: Spill,
179        MS: Fn() -> Result<F, StorageError>,
180    {
181        trx.add_commands::<F, MS>(
182            commands,
183            &mut self.provider,
184            &mut self.policy_store,
185            sink,
186            buffers,
187            &make_spill,
188        )
189    }
190
191    pub fn update_heads<I>(
192        &mut self,
193        graph_id: GraphId,
194        addrs: I,
195        request_heads: &mut PeerCache,
196        buffer: &mut TraversalBuffer,
197    ) -> Result<(), ClientError>
198    where
199        I: IntoIterator<Item = Address>,
200        I::IntoIter: DoubleEndedIterator,
201    {
202        let storage = self.provider.get_storage(graph_id)?;
203
204        // Commands in sync messages are always ancestor-first (lower max_cut to higher max_cut).
205        // Reverse the iterator to process highest max_cut first, which allows us to skip ancestors
206        // since if a command is an ancestor of one we've already added, we don't need to add it.
207        for address in addrs.into_iter().rev() {
208            if let Some(loc) = storage.get_location(address, buffer)? {
209                request_heads.add_command(
210                    storage,
211                    LocatedAddress {
212                        id: address.id,
213                        segment: loc.segment,
214                        max_cut: address.max_cut,
215                    },
216                    buffer,
217                )?;
218            } else {
219                error!(
220                    "UPDATE_HEADS: Address {:?} does NOT exist in storage, skipping (should not happen if command was successfully added)",
221                    address
222                );
223            }
224        }
225
226        Ok(())
227    }
228
229    /// Returns the address of the head of the graph.
230    pub fn head_address(&mut self, graph_id: GraphId) -> Result<Address, ClientError> {
231        let storage = self.provider.get_storage(graph_id)?;
232        let address = storage.get_head_address()?;
233        Ok(address)
234    }
235
236    /// Performs an `action`, writing the results to `sink`.
237    pub fn action(
238        &mut self,
239        graph_id: GraphId,
240        sink: &mut impl Sink<PS::Effect>,
241        action: <PS::Policy as Policy>::Action<'_>,
242    ) -> Result<(), ClientError> {
243        let storage = self.provider.get_storage(graph_id)?;
244
245        let head = storage.get_head()?;
246
247        let mut perspective = storage.get_linear_perspective(head)?;
248
249        let policy_id = perspective.policy();
250        let policy = self.policy_store.get_policy(policy_id)?;
251
252        // No need to checkpoint the perspective since it is only for this action.
253        // Must checkpoint once we add action transactions.
254
255        sink.begin();
256        match policy.call_action(action, &mut perspective, sink, ActionPlacement::OnGraph) {
257            Ok(()) => {
258                let segment = storage.write(perspective)?;
259                storage.commit(segment)?;
260                sink.commit();
261                Ok(())
262            }
263            Err(e) => {
264                sink.rollback();
265                Err(e.into())
266            }
267        }
268    }
269}
270
271impl<PS, SP> ClientState<PS, SP>
272where
273    SP: StorageProvider,
274{
275    /// Create a new [`Transaction`], used to receive [`Command`]s when syncing.
276    pub fn transaction(&mut self, graph_id: GraphId) -> Transaction<SP, PS> {
277        Transaction::new(graph_id)
278    }
279
280    /// Create an ephemeral [`Session`] associated with this client.
281    pub fn session(&mut self, graph_id: GraphId) -> Result<Session<SP, PS>, ClientError> {
282        Session::new(&mut self.provider, graph_id)
283    }
284
285    /// Checks if a command with the given address exists in the specified graph.
286    ///
287    /// Returns `true` if the command exists, `false` if it doesn't exist or the graph doesn't exist.
288    /// This method is used to determine if we need to sync when a hello message is received.
289    pub fn command_exists(
290        &mut self,
291        graph_id: GraphId,
292        address: Address,
293        buffer: &mut TraversalBuffer,
294    ) -> bool {
295        let Ok(storage) = self.provider.get_storage(graph_id) else {
296            // Graph doesn't exist
297            return false;
298        };
299        storage
300            .get_location(address, buffer)
301            .unwrap_or(None)
302            .is_some()
303    }
304}