ark_module/types/behavior.rs
1//! 🦹 Behavior module
2//!
3//! Behavior instances are identified through a densely packed identifier called `BehaviorInstanceId`.
4//!
5//! `BehaviorInstanceId` uniquely identifies a triple of:
6//!
7//! * An external behavior module handle
8//! * A behavior type in this behavior module
9//! * A behavior instance of this behavior type in this module
10//!
11//! Different users of this identifier need different parts:
12//!
13//! * `BehaviorModuleHandle`, while internal to Ark, consists of the high 16 bits of this id,
14//! and can be used to index into an array of behavior modules
15//! * `BehaviorTypeId` consists of the module handle as well as the local type id, that is, the
16//! high 32-bits. It is used to create new behavior instances, through explicit creation or deserialization.
17//! * `ForeignBehaviorInstanceId` consists of the parts that are local to a behavior module. Ark will complete
18//! it together with a `BehaviorModuleHandle` into a `BehaviorInstanceId`.
19//!
20//! /---------- `ForeignBehaviorInstanceId` ------------\
21//! / \
22//! `LocalBehaviorInstanceId` `LocalBehaviorTypeId` `BehaviorModuleHandle`
23//! 0.........................31 32......................47 48.........................63
24//! \ /
25//! \--------------- `BehaviorTypeId` ------------------/
26//! \ /
27//! \--------------------------- `BehaviorInstanceId` ----------------------------/
28
29pub use crate::ErrorCode;
30pub use ark_api::behavior::ForeignBehaviorInstanceId;
31pub use ark_api::behavior::IncomingMessage;
32pub use ark_api::behavior::IncomingMessageAddr;
33pub use ark_api::behavior::LocalBehaviorInstanceId;
34pub use ark_api::behavior::LocalBehaviorTypeId;
35pub use ark_api_ffi::behavior_v0::ActorId;
36pub use ark_api_ffi::behavior_v0::IncomingMessage as IncomingMessageRaw;
37pub use ark_api_ffi::behavior_v0::LocalBehaviorRegistration;
38pub use ark_api_ffi::behavior_v0::LocalModuleRegistration;
39
40/// Ergonomic wrapper around a behavior module's entry points.
41///
42/// See the module level documentation for an explanation on the various identifiers
43/// used in this trait
44pub trait BehaviorEntry {
45 /// Should return a serialized JSON vector of bytes of local registration info.
46 ///
47 /// # Example
48 ///
49 /// ```rust
50 /// use ark_module::behavior::{LocalBehaviorRegistration, LocalModuleRegistration};
51 ///
52 /// fn register(&self) -> Vec<u8> {
53 /// let behavior_info = LocalBehaviorRegistration{
54 /// type_id: 0,
55 /// name: String::from("test123"),
56 /// };
57 ///
58 /// let registration_info = LocalModuleRegistration{
59 /// behavior_infos: vec![behavior_info],
60 /// };
61 ///
62 /// serde_json::to_vec(®istration_info).unwrap()
63 /// }
64 /// ```
65 /// Users can return custom registration data by including `LocalModuleRegistration`
66 /// and tagging it with `#[serde(flatten)]`.
67 ///
68 /// ```rust
69 /// use ark_module::behavior::{LocalBehaviorRegistration, LocalModuleRegistration};
70 ///
71 /// #[derive(serde::Deserialize)]
72 /// pub struct CustomRegistrationInfo {
73 /// #[serde(flatten)]
74 /// pub registration_info: LocalModuleRegistration,
75 /// pub custom_string: String,
76 /// pub custom_uint: u32,
77 /// }
78 ///
79 /// fn register(&self) -> Vec<u8> {
80 /// let behavior_info = LocalBehaviorRegistration{
81 /// type_id: 0,
82 /// name: String::from("test123"),
83 /// };
84 ///
85 /// let registration_info = LocalModuleRegistration{
86 /// behavior_infos: vec![behavior_info],
87 /// };
88 ///
89 /// let custom_registration_info = CustomRegistrationInfo {
90 /// registration_info,
91 /// custom_string: String::from("test456"),
92 /// custom_uint: 1,
93 /// };
94 /// }
95 /// ```
96 fn register(&mut self) -> Vec<u8>;
97
98 /// Persists a whole module after it failed, before an attempt to reload it and re-run it.
99 /// In particular, the module should keep an account of all the cleanups it should run after
100 /// restoring is done.
101 fn persist(&mut self) -> Vec<u8>;
102
103 /// Restores a module serialized with `persist`. The module should run additional cleanups at
104 /// this point.
105 fn restore(&mut self, state: &[u8]);
106
107 /// Construct a new [`LocalBehaviorInstanceId`] given the type id and ron parameters,
108 /// which can then be packed into a [`ForeignBehaviorInstanceId`] and returned here.
109 fn instance_create(
110 &mut self,
111 local_type_id: LocalBehaviorTypeId,
112 ron_params: &str,
113 ) -> ForeignBehaviorInstanceId;
114
115 /// Destroys the provided instance in the behavior module
116 fn instance_destroy(&mut self, instance_id: ForeignBehaviorInstanceId);
117
118 /// Clones the provided instance in the behavior module
119 fn instance_clone(
120 &mut self,
121 instance_id: ForeignBehaviorInstanceId,
122 ) -> ForeignBehaviorInstanceId;
123
124 /// Serialize the instance's state into a binary blob
125 fn instance_persist(&self, instance_id: ForeignBehaviorInstanceId) -> Vec<u8>;
126
127 /// Restore an instance given the provided type id and serialized state
128 fn instance_restore(
129 &mut self,
130 local_type_id: LocalBehaviorTypeId,
131 blob: &[u8],
132 ) -> ForeignBehaviorInstanceId;
133
134 /// Have the provided instance handle the incoming serialized message
135 fn instance_handle_message(&mut self, addr: IncomingMessageAddr, serialized_message: &[u8]);
136
137 /// Process the incoming messages
138 fn instances_handle_messages<'a>(
139 &mut self,
140 messages: impl ExactSizeIterator<Item = IncomingMessage<'a>>,
141 ) {
142 for msg in messages {
143 self.instance_handle_message(msg.addr, msg.serialized_message);
144 }
145 }
146}
147
148/// An iterator over a collection of [`IncomingMessage`]s encoded in a byte buffer
149#[doc(hidden)]
150pub struct IncomingMessageIter<'a> {
151 bytes: &'a [u8],
152 offset: usize,
153 len: usize,
154}
155
156impl<'a> IncomingMessageIter<'a> {
157 #[doc(hidden)]
158 pub fn new(bytes: &'a [u8]) -> Result<Self, ErrorCode> {
159 let mut offset = 0;
160 let len = Self::read_u32_len(bytes, &mut offset).ok_or(ErrorCode::InternalError)?;
161
162 Ok(Self { bytes, offset, len })
163 }
164
165 /// Read a u32 value from our byte stream, but return as usize as we use it for length
166 fn read_u32_len(bytes: &[u8], offset: &mut usize) -> Option<usize> {
167 let result = bytes.get((*offset)..(*offset + 4)).map(|slice| {
168 *offset += 4;
169 u32::from_le_bytes(slice.try_into().unwrap()) as usize
170 });
171
172 result
173 }
174
175 fn read_incoming_msg_ffi(bytes: &[u8], offset: &mut usize) -> Option<IncomingMessageRaw> {
176 const INCOMING_MESSAGE_SIZE: usize = std::mem::size_of::<IncomingMessageRaw>();
177
178 let result = bytes
179 .get((*offset)..(*offset + INCOMING_MESSAGE_SIZE))
180 .map(|slice| {
181 *offset += INCOMING_MESSAGE_SIZE;
182 // SAFETY: Safe due to the layout enforcement of `repr(C)` and being `Copy`
183 unsafe {
184 std::mem::transmute::<[u8; INCOMING_MESSAGE_SIZE], IncomingMessageRaw>(
185 slice.try_into().unwrap(),
186 )
187 }
188 });
189
190 result
191 }
192
193 fn read_incoming_message(bytes: &'a [u8], offset: &mut usize) -> Option<IncomingMessage<'a>> {
194 Self::read_incoming_msg_ffi(bytes, offset).map(|incoming| {
195 *offset += incoming.serialized_message_len as usize;
196 IncomingMessage::from(incoming)
197 })
198 }
199}
200
201impl<'a> Iterator for IncomingMessageIter<'a> {
202 type Item = IncomingMessage<'a>;
203
204 fn next(&mut self) -> Option<Self::Item> {
205 IncomingMessageIter::read_incoming_message(self.bytes, &mut self.offset)
206 }
207}
208
209impl<'a> ExactSizeIterator for IncomingMessageIter<'a> {
210 fn len(&self) -> usize {
211 self.len
212 }
213}
214
215/// By passing a type that implements the [`BehaviorEntry`] trait, this macro will hook up
216/// and route the behavior module's raw entry points to the trait methods.
217#[macro_export]
218macro_rules! impl_behavior_module {
219 ($module:ty) => {
220 use $crate::behavior::BehaviorEntry as _;
221
222 static mut MODULE: Option<$module> = None;
223
224 #[no_mangle]
225 pub fn ark_behavior_register() {
226 $crate::init();
227
228 let mut module = <$module>::default();
229 let behavior_infos = module.register();
230
231 // SAFETY: Sound to access as are no threads in Wasm and this is the first behavior function that will be called in the module
232 let module = unsafe { MODULE.replace(module) };
233
234 ark::core::return_slice(&behavior_infos);
235 }
236
237 #[no_mangle]
238 pub fn ark_behavior_persist() {
239 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
240 let module = unsafe { MODULE.as_mut().unwrap() };
241 let state = module.persist();
242 ark::core::return_slice(&state);
243 }
244
245 #[no_mangle]
246 pub unsafe fn ark_behavior_restore(state_ptr: *const u8, state_size: u32) {
247 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
248 let module = unsafe { MODULE.as_mut().unwrap() };
249 // SAFETY: Usage is kept within lifetime of function call only
250 let state_slice = unsafe { $crate::util::param_byte_slice(state_ptr, state_size) };
251 module.restore(&state_slice);
252 }
253
254 #[no_mangle]
255 pub unsafe fn ark_behavior_instance_create(
256 local_type_id: u32,
257 params_ptr: *const u8,
258 params_len: u32,
259 ) -> u64 {
260 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
261 let module = unsafe { MODULE.as_mut().unwrap() };
262 // SAFETY: Usage is kept within lifetime of function call only
263 let params_slice = unsafe { $crate::util::param_byte_slice(params_ptr, params_len) };
264 // TODO: should we require UTF-8 strings for the parameters, this constraint is not described in the entrypoint definition
265 let params_str = std::str::from_utf8(params_slice).unwrap();
266
267 assert!(local_type_id <= u16::max_value() as u32);
268 let local_type_id = $crate::behavior::LocalBehaviorTypeId(local_type_id as u16);
269 module.instance_create(local_type_id, params_str).0
270 }
271
272 #[no_mangle]
273 pub fn ark_behavior_instance_destroy(instance_id: u64) {
274 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
275 let module = unsafe { MODULE.as_mut().unwrap() };
276 module.instance_destroy($crate::behavior::ForeignBehaviorInstanceId(instance_id));
277 }
278
279 #[no_mangle]
280 pub fn ark_behavior_instance_clone(instance_id: u64) -> u64 {
281 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
282 let module = unsafe { MODULE.as_mut().unwrap() };
283 module
284 .instance_clone($crate::behavior::ForeignBehaviorInstanceId(instance_id))
285 .0
286 }
287
288 #[no_mangle]
289 pub fn ark_behavior_instance_persist(instance_id: u64) {
290 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
291 let module = unsafe { MODULE.as_mut().unwrap() };
292 let bytes =
293 module.instance_persist($crate::behavior::ForeignBehaviorInstanceId(instance_id));
294 ark::core::return_slice(&bytes);
295 }
296
297 #[no_mangle]
298 pub unsafe fn ark_behavior_instance_restore(
299 local_type_id: u32,
300 blob_ptr: *const u8,
301 blob_size: u32,
302 ) -> u64 {
303 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
304 let module = unsafe { MODULE.as_mut().unwrap() };
305 // SAFETY: Usage is kept within lifetime of function call only
306 let blob_slice = unsafe { $crate::util::param_byte_slice(blob_ptr, blob_size) };
307 assert!(local_type_id <= u16::max_value() as u32);
308 let local_type_id = $crate::behavior::LocalBehaviorTypeId(local_type_id as u16);
309
310 module.instance_restore(local_type_id, blob_slice).0
311 }
312
313 #[no_mangle]
314 pub unsafe fn ark_behavior_instance_handle_message(
315 instance_id: u64,
316 actor_id: $crate::behavior::ActorId,
317 message_blob_ptr: *const u8,
318 message_blob_size: u32,
319 ) {
320 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
321 let module = unsafe { MODULE.as_mut().unwrap() };
322 // SAFETY: Usage is kept within lifetime of function call only
323 let message_blob_slice =
324 unsafe { $crate::util::param_byte_slice(message_blob_ptr, message_blob_size) };
325
326 let instance_id = $crate::behavior::ForeignBehaviorInstanceId(instance_id);
327
328 let addr = $crate::behavior::IncomingMessageAddr::from_raw(actor_id, instance_id);
329
330 module.instance_handle_message(addr, message_blob_slice);
331 }
332
333 #[no_mangle]
334 pub unsafe fn ark_behavior_instances_handle_messages(
335 encoded_ptr: *const u8,
336 encoded_len: u32,
337 ) {
338 // SAFETY: Sound to access static as we do not use threads in Wasm, and this is guaranteed to be initialized on startup with `ark_behavior_register` first
339 let module = unsafe { MODULE.as_mut().unwrap() };
340 // SAFETY: Usage is kept within lifetime of function call only
341 let encoded_slice = unsafe { $crate::util::param_byte_slice(encoded_ptr, encoded_len) };
342 if let Ok(iter) = $crate::behavior::IncomingMessageIter::new(encoded_slice) {
343 module.instances_handle_messages(iter);
344 } else {
345 $crate::error!("Invalid encoded `IncomingMessage` buffer");
346 }
347 }
348 };
349}