Skip to main content

hataori_runtime/
action.rs

1use crate::error::{ActionError, RuntimeError};
2use hataori_runtime_foundation::protocol::ActionId;
3use std::{
4    any::Any,
5    collections::BTreeMap,
6    panic::{catch_unwind, AssertUnwindSafe},
7    sync::Arc,
8};
9
10pub type Segments = Vec<Vec<u8>>;
11
12pub trait WireValue: Sized + Send + 'static {
13    const SCHEMA_ID: u64;
14
15    fn encode(self) -> Result<Segments, ActionError>;
16    fn decode(segments: Segments) -> Result<Self, ActionError>;
17}
18
19pub trait Action: WireValue {
20    const ID: u128;
21    type Output: WireValue;
22}
23
24pub(crate) enum ActionValue {
25    Encoded(Segments),
26    Typed(Box<dyn Any + Send>),
27}
28
29type Handler = dyn Fn(Segments) -> Result<Segments, ActionError> + Send + Sync;
30type LocalHandler =
31    dyn Fn(Box<dyn Any + Send>) -> Result<Box<dyn Any + Send>, ActionError> + Send + Sync;
32
33#[derive(Clone)]
34pub(crate) struct RegisteredAction {
35    handler: Arc<Handler>,
36    local_handler: Option<Arc<LocalHandler>>,
37}
38
39impl std::fmt::Debug for RegisteredAction {
40    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
41        formatter.write_str("RegisteredAction { .. }")
42    }
43}
44
45impl RegisteredAction {
46    pub(crate) fn new(
47        handler: impl Fn(Segments) -> Result<Segments, ActionError> + Send + Sync + 'static,
48    ) -> Self {
49        Self {
50            handler: Arc::new(handler),
51            local_handler: None,
52        }
53    }
54
55    pub(crate) fn execute(&self, input: ActionValue) -> Result<ActionValue, ActionError> {
56        match catch_unwind(AssertUnwindSafe(|| match input {
57            ActionValue::Encoded(segments) => (self.handler)(segments).map(ActionValue::Encoded),
58            ActionValue::Typed(value) => self
59                .local_handler
60                .as_ref()
61                .ok_or_else(|| ActionError::codec("action has no typed local handler"))?(
62                value
63            )
64            .map(ActionValue::Typed),
65        })) {
66            Ok(result) => result,
67            Err(_) => Err(ActionError::Panic),
68        }
69    }
70
71    pub(crate) fn execute_encoded(&self, input: Segments) -> Result<Segments, ActionError> {
72        match self.execute(ActionValue::Encoded(input))? {
73            ActionValue::Encoded(output) => Ok(output),
74            ActionValue::Typed(_) => {
75                Err(ActionError::codec("encoded action returned typed output"))
76            }
77        }
78    }
79}
80
81#[derive(Clone, Debug, Default)]
82pub(crate) struct ActionRegistry {
83    entries: BTreeMap<ActionId, RegisteredAction>,
84    schemas: BTreeMap<ActionId, (u64, u64)>,
85}
86
87impl ActionRegistry {
88    pub(crate) fn register<A, F>(&mut self, handler: F) -> Result<(), RuntimeError>
89    where
90        A: Action,
91        F: Fn(A) -> Result<A::Output, ActionError> + Send + Sync + 'static,
92    {
93        let id = ActionId::new(A::ID).map_err(|_| RuntimeError::InvalidActionId)?;
94        if A::SCHEMA_ID == 0 || A::Output::SCHEMA_ID == 0 {
95            return Err(RuntimeError::InvalidSchema);
96        }
97        if self.entries.contains_key(&id) {
98            return Err(RuntimeError::DuplicateAction(id));
99        }
100        let handler = Arc::new(handler);
101        let remote = Arc::clone(&handler);
102        let local = Arc::clone(&handler);
103        self.entries.insert(
104            id,
105            RegisteredAction {
106                handler: Arc::new(move |segments| {
107                    let input = A::decode(segments)?;
108                    remote(input)?.encode()
109                }),
110                local_handler: Some(Arc::new(move |value| {
111                    let input = value
112                        .downcast::<A>()
113                        .map_err(|_| ActionError::codec("typed action input mismatch"))?;
114                    Ok(Box::new(local(*input)?) as Box<dyn Any + Send>)
115                })),
116            },
117        );
118        self.schemas
119            .insert(id, (A::SCHEMA_ID, A::Output::SCHEMA_ID));
120        Ok(())
121    }
122
123    pub(crate) fn insert_erased(
124        &mut self,
125        id: ActionId,
126        input_schema: u64,
127        output_schema: u64,
128        handler: RegisteredAction,
129    ) -> Result<(), RuntimeError> {
130        if input_schema == 0 || output_schema == 0 {
131            return Err(RuntimeError::InvalidSchema);
132        }
133        if self.entries.contains_key(&id) {
134            return Err(RuntimeError::DuplicateAction(id));
135        }
136        self.entries.insert(id, handler);
137        self.schemas.insert(id, (input_schema, output_schema));
138        Ok(())
139    }
140
141    pub(crate) fn get(&self, id: ActionId) -> Option<&RegisteredAction> {
142        self.entries.get(&id)
143    }
144
145    pub(crate) fn fingerprint(&self) -> [u8; 32] {
146        const SEEDS: [u64; 4] = [
147            0xcbf2_9ce4_8422_2325,
148            0x8422_2325_cbf2_9ce4,
149            0x9e37_79b9_7f4a_7c15,
150            0xd6e8_feb8_6659_fd93,
151        ];
152        let mut hashes = SEEDS;
153        for (id, (input, output)) in &self.schemas {
154            for byte in id
155                .get()
156                .to_le_bytes()
157                .into_iter()
158                .chain(input.to_le_bytes())
159                .chain(output.to_le_bytes())
160            {
161                for (index, hash) in hashes.iter_mut().enumerate() {
162                    *hash ^= u64::from(byte).wrapping_add(index as u64);
163                    *hash = hash.wrapping_mul(0x100_0000_01b3);
164                }
165            }
166        }
167        let mut result = [0_u8; 32];
168        for (index, hash) in hashes.into_iter().enumerate() {
169            result[index * 8..index * 8 + 8].copy_from_slice(&hash.to_le_bytes());
170        }
171        result
172    }
173}
174
175impl WireValue for () {
176    const SCHEMA_ID: u64 = 1;
177
178    fn encode(self) -> Result<Segments, ActionError> {
179        Ok(Vec::new())
180    }
181
182    fn decode(segments: Segments) -> Result<Self, ActionError> {
183        if segments.is_empty() {
184            Ok(())
185        } else {
186            Err(ActionError::codec("unit value has payload"))
187        }
188    }
189}
190
191impl WireValue for u64 {
192    const SCHEMA_ID: u64 = 2;
193
194    fn encode(self) -> Result<Segments, ActionError> {
195        Ok(vec![self.to_le_bytes().to_vec()])
196    }
197
198    fn decode(segments: Segments) -> Result<Self, ActionError> {
199        let bytes: [u8; 8] = one_segment(segments, "u64")?
200            .try_into()
201            .map_err(|_| ActionError::codec("u64 payload must contain 8 bytes"))?;
202        Ok(Self::from_le_bytes(bytes))
203    }
204}
205
206impl WireValue for Vec<u8> {
207    const SCHEMA_ID: u64 = 3;
208
209    fn encode(self) -> Result<Segments, ActionError> {
210        Ok(vec![self])
211    }
212
213    fn decode(segments: Segments) -> Result<Self, ActionError> {
214        one_segment(segments, "byte vector")
215    }
216}
217
218impl WireValue for String {
219    const SCHEMA_ID: u64 = 4;
220
221    fn encode(self) -> Result<Segments, ActionError> {
222        Ok(vec![self.into_bytes()])
223    }
224
225    fn decode(segments: Segments) -> Result<Self, ActionError> {
226        String::from_utf8(one_segment(segments, "string")?)
227            .map_err(|_| ActionError::codec("string payload is not UTF-8"))
228    }
229}
230
231fn one_segment(mut segments: Segments, name: &'static str) -> Result<Vec<u8>, ActionError> {
232    if segments.len() != 1 {
233        return Err(ActionError::codec(format!(
234            "{name} requires exactly one segment"
235        )));
236    }
237    Ok(segments.pop().unwrap())
238}
239
240#[cfg(test)]
241mod tests;