1use std::marker::PhantomData;
10use std::sync::Arc;
11use std::time::Duration;
12
13use dex_protocol::dex::Value as ProtoValue;
14
15use crate::attribute::AttributeLock;
16use crate::step::{ErasedValue, TypedValue};
17use crate::value_mapper;
18use crate::{Context, HandlerResult, SdkResult, StepMovement, Value};
19
20pub struct Rpc<Input, Output> {
49 name: &'static str,
50 marker: PhantomData<fn(Input) -> Output>,
51}
52
53impl<Input, Output> Rpc<Input, Output> {
54 pub const fn new(name: &'static str) -> Self {
56 Self {
57 name,
58 marker: PhantomData,
59 }
60 }
61
62 pub fn timeout(self, timeout: Duration) -> RpcDefinition<Input, Output> {
64 RpcDefinition::from(self).timeout(timeout)
65 }
66
67 pub fn lock(self, lock: AttributeLock) -> RpcDefinition<Input, Output> {
69 RpcDefinition::from(self).lock(lock)
70 }
71
72 pub fn is_transactional(self) -> RpcDefinition<Input, Output> {
74 RpcDefinition::from(self).is_transactional()
75 }
76
77 pub(crate) fn name(&self) -> &'static str {
78 self.name
79 }
80}
81
82impl<Input, Output> Clone for Rpc<Input, Output> {
83 fn clone(&self) -> Self {
84 *self
85 }
86}
87
88impl<Input, Output> Copy for Rpc<Input, Output> {}
89
90pub struct RpcDefinition<Input, Output> {
95 rpc: Rpc<Input, Output>,
96 timeout: Option<Duration>,
97 locks: Vec<AttributeLock>,
98 is_transactional: bool,
99}
100
101impl<Input, Output> RpcDefinition<Input, Output> {
102 pub fn timeout(mut self, timeout: Duration) -> Self {
104 self.timeout = Some(timeout);
105 self
106 }
107
108 pub fn lock(mut self, lock: AttributeLock) -> Self {
110 self.locks.push(lock);
111 self
112 }
113
114 pub fn is_transactional(mut self) -> Self {
116 self.is_transactional = true;
117 self
118 }
119}
120
121impl<Input, Output> From<Rpc<Input, Output>> for RpcDefinition<Input, Output> {
122 fn from(rpc: Rpc<Input, Output>) -> Self {
123 Self {
124 rpc,
125 timeout: None,
126 locks: Vec::new(),
127 is_transactional: false,
128 }
129 }
130}
131
132pub struct RpcList<FlowType> {
137 definitions: Vec<Box<dyn RpcBinder<FlowType>>>,
138}
139
140impl<FlowType> RpcList<FlowType>
141where
142 FlowType: Send + Sync + 'static,
143{
144 pub fn new() -> Self {
146 Self {
147 definitions: Vec::new(),
148 }
149 }
150
151 pub fn function<Input, Output>(
155 mut self,
156 definition: impl Into<RpcDefinition<Input, Output>>,
157 handler: fn(&FlowType, &mut Context, Input) -> HandlerResult<RpcResult<Output>>,
158 ) -> Self
159 where
160 Input: Value,
161 Output: Value,
162 {
163 self.definitions.push(Box::new(FunctionBinder {
164 definition: definition.into(),
165 handler,
166 }));
167 self
168 }
169
170 pub fn function_without_input<Output>(
172 mut self,
173 definition: impl Into<RpcDefinition<(), Output>>,
174 handler: fn(&FlowType, &mut Context) -> HandlerResult<RpcResult<Output>>,
175 ) -> Self
176 where
177 Output: Value,
178 {
179 self.definitions.push(Box::new(FunctionWithoutInputBinder {
180 definition: definition.into(),
181 handler,
182 }));
183 self
184 }
185
186 pub fn procedure<Input>(
188 mut self,
189 definition: impl Into<RpcDefinition<Input, ()>>,
190 handler: fn(&FlowType, &mut Context, Input) -> HandlerResult<()>,
191 ) -> Self
192 where
193 Input: Value,
194 {
195 self.definitions.push(Box::new(ProcedureBinder {
196 definition: definition.into(),
197 handler,
198 }));
199 self
200 }
201
202 pub fn procedure_without_input(
204 mut self,
205 definition: impl Into<RpcDefinition<(), ()>>,
206 handler: fn(&FlowType, &mut Context) -> HandlerResult<()>,
207 ) -> Self {
208 self.definitions.push(Box::new(ProcedureWithoutInputBinder {
209 definition: definition.into(),
210 handler,
211 }));
212 self
213 }
214
215 pub(crate) fn bind(self, flow: Arc<FlowType>) -> Vec<RegisteredRpc> {
216 self.definitions
217 .into_iter()
218 .map(|definition| definition.bind(Arc::clone(&flow)))
219 .collect()
220 }
221}
222
223impl<FlowType> Default for RpcList<FlowType>
224where
225 FlowType: Send + Sync + 'static,
226{
227 fn default() -> Self {
228 Self::new()
229 }
230}
231
232pub struct RpcResult<Output> {
236 output: Output,
237 next_steps: Vec<StepMovement>,
238 cancel_step_types: Vec<&'static str>,
239}
240
241impl<Output: Value> RpcResult<Output> {
242 pub fn new(output: Output) -> Self {
244 Self {
245 output,
246 next_steps: Vec::new(),
247 cancel_step_types: Vec::new(),
248 }
249 }
250
251 pub fn then(mut self, movement: StepMovement) -> Self {
253 self.next_steps.push(movement);
254 self
255 }
256
257 pub fn cancel_step<SelectedStep>(mut self, step: &SelectedStep) -> Self
262 where
263 SelectedStep: crate::Step,
264 {
265 let step_type = step.step_type();
266 if !self.cancel_step_types.contains(&step_type) {
267 self.cancel_step_types.push(step_type);
268 }
269 self
270 }
271
272 pub fn output(&self) -> &Output {
274 &self.output
275 }
276
277 pub fn next_steps(&self) -> &[StepMovement] {
279 &self.next_steps
280 }
281}
282
283pub(crate) struct ErasedRpcResult {
284 pub(crate) output: Box<dyn ErasedValue>,
285 pub(crate) next_steps: Vec<StepMovement>,
286 pub(crate) cancel_step_types: Vec<&'static str>,
287}
288
289#[derive(Clone)]
290pub(crate) struct RegisteredRpc {
291 pub(crate) name: &'static str,
292 pub(crate) timeout: Option<Duration>,
293 pub(crate) locks: Vec<AttributeLock>,
294 pub(crate) is_transactional: bool,
295 pub(crate) handler: Arc<dyn ErasedRpc>,
296}
297
298pub(crate) trait ErasedRpc: Send + Sync {
299 fn invoke(&self, context: &mut Context, input: &ProtoValue) -> HandlerResult<ErasedRpcResult>;
300}
301
302trait RpcBinder<FlowType>: Send + Sync {
303 fn bind(self: Box<Self>, flow: Arc<FlowType>) -> RegisteredRpc;
304}
305
306struct FunctionBinder<FlowType, Input, Output> {
307 definition: RpcDefinition<Input, Output>,
308 handler: fn(&FlowType, &mut Context, Input) -> HandlerResult<RpcResult<Output>>,
309}
310
311impl<FlowType, Input, Output> RpcBinder<FlowType> for FunctionBinder<FlowType, Input, Output>
312where
313 FlowType: Send + Sync + 'static,
314 Input: Value,
315 Output: Value,
316{
317 fn bind(self: Box<Self>, flow: Arc<FlowType>) -> RegisteredRpc {
318 registered_rpc(
319 self.definition,
320 Arc::new(BoundFunction {
321 flow,
322 handler: self.handler,
323 }),
324 )
325 }
326}
327
328struct BoundFunction<FlowType, Input, Output> {
329 flow: Arc<FlowType>,
330 handler: fn(&FlowType, &mut Context, Input) -> HandlerResult<RpcResult<Output>>,
331}
332
333impl<FlowType, Input, Output> ErasedRpc for BoundFunction<FlowType, Input, Output>
334where
335 FlowType: Send + Sync + 'static,
336 Input: Value,
337 Output: Value,
338{
339 fn invoke(&self, context: &mut Context, input: &ProtoValue) -> HandlerResult<ErasedRpcResult> {
340 let result = (self.handler)(&self.flow, context, value_mapper::decode_handler(input)?)?;
341 Ok(ErasedRpcResult {
342 output: Box::new(TypedValue(result.output)),
343 next_steps: result.next_steps,
344 cancel_step_types: result.cancel_step_types,
345 })
346 }
347}
348
349struct FunctionWithoutInputBinder<FlowType, Output> {
350 definition: RpcDefinition<(), Output>,
351 handler: fn(&FlowType, &mut Context) -> HandlerResult<RpcResult<Output>>,
352}
353
354impl<FlowType, Output> RpcBinder<FlowType> for FunctionWithoutInputBinder<FlowType, Output>
355where
356 FlowType: Send + Sync + 'static,
357 Output: Value,
358{
359 fn bind(self: Box<Self>, flow: Arc<FlowType>) -> RegisteredRpc {
360 registered_rpc(
361 self.definition,
362 Arc::new(BoundFunctionWithoutInput {
363 flow,
364 handler: self.handler,
365 }),
366 )
367 }
368}
369
370struct BoundFunctionWithoutInput<FlowType, Output> {
371 flow: Arc<FlowType>,
372 handler: fn(&FlowType, &mut Context) -> HandlerResult<RpcResult<Output>>,
373}
374
375impl<FlowType, Output> ErasedRpc for BoundFunctionWithoutInput<FlowType, Output>
376where
377 FlowType: Send + Sync + 'static,
378 Output: Value,
379{
380 fn invoke(&self, context: &mut Context, _input: &ProtoValue) -> HandlerResult<ErasedRpcResult> {
381 let result = (self.handler)(&self.flow, context)?;
382 Ok(ErasedRpcResult {
383 output: Box::new(TypedValue(result.output)),
384 next_steps: result.next_steps,
385 cancel_step_types: result.cancel_step_types,
386 })
387 }
388}
389
390struct ProcedureBinder<FlowType, Input> {
391 definition: RpcDefinition<Input, ()>,
392 handler: fn(&FlowType, &mut Context, Input) -> HandlerResult<()>,
393}
394
395impl<FlowType, Input> RpcBinder<FlowType> for ProcedureBinder<FlowType, Input>
396where
397 FlowType: Send + Sync + 'static,
398 Input: Value,
399{
400 fn bind(self: Box<Self>, flow: Arc<FlowType>) -> RegisteredRpc {
401 registered_rpc(
402 self.definition,
403 Arc::new(BoundProcedure {
404 flow,
405 handler: self.handler,
406 }),
407 )
408 }
409}
410
411struct BoundProcedure<FlowType, Input> {
412 flow: Arc<FlowType>,
413 handler: fn(&FlowType, &mut Context, Input) -> HandlerResult<()>,
414}
415
416impl<FlowType, Input> ErasedRpc for BoundProcedure<FlowType, Input>
417where
418 FlowType: Send + Sync + 'static,
419 Input: Value,
420{
421 fn invoke(&self, context: &mut Context, input: &ProtoValue) -> HandlerResult<ErasedRpcResult> {
422 (self.handler)(&self.flow, context, value_mapper::decode_handler(input)?)?;
423 Ok(empty_rpc_result())
424 }
425}
426
427struct ProcedureWithoutInputBinder<FlowType> {
428 definition: RpcDefinition<(), ()>,
429 handler: fn(&FlowType, &mut Context) -> HandlerResult<()>,
430}
431
432impl<FlowType> RpcBinder<FlowType> for ProcedureWithoutInputBinder<FlowType>
433where
434 FlowType: Send + Sync + 'static,
435{
436 fn bind(self: Box<Self>, flow: Arc<FlowType>) -> RegisteredRpc {
437 registered_rpc(
438 self.definition,
439 Arc::new(BoundProcedureWithoutInput {
440 flow,
441 handler: self.handler,
442 }),
443 )
444 }
445}
446
447struct BoundProcedureWithoutInput<FlowType> {
448 flow: Arc<FlowType>,
449 handler: fn(&FlowType, &mut Context) -> HandlerResult<()>,
450}
451
452impl<FlowType> ErasedRpc for BoundProcedureWithoutInput<FlowType>
453where
454 FlowType: Send + Sync + 'static,
455{
456 fn invoke(&self, context: &mut Context, _input: &ProtoValue) -> HandlerResult<ErasedRpcResult> {
457 (self.handler)(&self.flow, context)?;
458 Ok(empty_rpc_result())
459 }
460}
461
462fn registered_rpc<Input, Output>(
463 definition: RpcDefinition<Input, Output>,
464 handler: Arc<dyn ErasedRpc>,
465) -> RegisteredRpc {
466 RegisteredRpc {
467 name: definition.rpc.name,
468 timeout: definition.timeout,
469 locks: definition.locks,
470 is_transactional: definition.is_transactional,
471 handler,
472 }
473}
474
475fn empty_rpc_result() -> ErasedRpcResult {
476 ErasedRpcResult {
477 output: Box::new(TypedValue(())),
478 next_steps: Vec::new(),
479 cancel_step_types: Vec::new(),
480 }
481}
482
483pub(crate) fn encode_rpc_output(output: &dyn ErasedValue) -> SdkResult<ProtoValue> {
484 output.encode()
485}