pie 0.1.1

Programmable Inference Engine (PIE)
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142

// Minimalistic driver implementation

use crate::backend;
use crate::backend::Backend;
use crate::service::Service;
use crate::utils::IdPool;
use dashmap::DashMap;
use prost::Message;

use std::sync::{Arc, Mutex};
use tokio::sync::{mpsc, oneshot};

pub const PROTOCOL: &str = "ping";

mod pb_bindings {
    include!(concat!(env!("OUT_DIR"), "/ping.rs"));
}

#[derive(Debug)]
pub enum Command {
    Ping {
        message: String,
        handler: oneshot::Sender<String>,
    },
}

#[derive(Debug)]
pub enum Event {
    Pong(oneshot::Sender<String>),
}
#[derive(Debug)]
pub struct Ping<B> {
    backend: B,
    protocol_id: u8,
    cmd_id_pool: IdPool<u32>,
    cmd_queue: Arc<Mutex<Vec<Command>>>,
    event_table: Arc<DashMap<u32, Event>>,
    event_loop_handle: tokio::task::JoinHandle<()>,
}

impl<B> Ping<B>
where
    B: Backend,
{
    pub async fn new(backend: B) -> Self {
        let protocol_id = backend
            .protocol_index(PROTOCOL)
            .expect("Failed to get protocol index");

        let (tx, rx) = mpsc::channel(1000);
        backend.register_listener(protocol_id, tx).await;

        let event_table = Arc::new(DashMap::new());
        let event_loop_handle = tokio::spawn(Self::event_loop(rx, event_table.clone()));

        Self {
            backend,
            protocol_id,

            cmd_id_pool: IdPool::new(u32::MAX),
            cmd_queue: Arc::new(Mutex::new(Vec::new())),
            event_table,
            event_loop_handle,
        }
    }

    async fn event_loop(
        mut rx: mpsc::Receiver<Vec<u8>>,
        event_table: Arc<DashMap<u32, Event>>,
    ) {
        while let Some(pong) = rx.recv().await {
            let pong = pb_bindings::Pong::decode(pong.as_slice()).unwrap();

            let correlation_id = pong.correlation_id;
            if let Some((_, event)) = event_table.remove(&correlation_id) {
                match event {
                    Event::Pong(sender) => {
                        // should be ok, not unwrap.
                        // the inst may have dropped even before the pong is received
                        sender.send(pong.message).ok();
                    }
                }
            }
        }
    }
}

//#[async_trait]
impl<B> Service for Ping<B>
where
    B: Backend,
{
    type Command = Command;

    async fn handle(&mut self, cmd: Self::Command) {
        match cmd {
            Command::Ping { message, handler } => {
                let correlation_id = self.cmd_id_pool.acquire().unwrap();

                let ping = pb_bindings::Ping {
                    correlation_id,
                    message,
                };
                self.event_table
                    .insert(correlation_id, Event::Pong(handler));
                self.backend
                    .send(self.protocol_id, ping.encode_to_vec())
                    .unwrap();
            }
        }
    }
}

#[derive(Clone)]
pub struct Simulator {
    pub protocols: Vec<String>,
}

impl Simulator {
    pub fn new() -> Self {
        Self {
            protocols: vec![PROTOCOL.to_string()],
        }
    }
}

impl backend::Simulate for Simulator {
    fn protocols(&self) -> &[String] {
        self.protocols.as_slice()
    }

    fn simulate(&mut self, command: Vec<u8>) -> Option<Vec<u8>> {
        let cmd = pb_bindings::Ping::decode(command.as_slice()).unwrap();
        Some(
            pb_bindings::Pong {
                correlation_id: cmd.correlation_id,
                message: format!("Pong: {}", cmd.message),
            }
            .encode_to_vec(),
        )
    }
}