flwrs_plugin/sink/
runner.rs

1use crate::plugin::core::ConnectionConfig;
2use crate::plugin::error::Error;
3use crate::plugin::logger::PluginLogger;
4use crate::plugin::msg_client::MSG_CLIENT;
5use crate::schema::common::log_level::Enum as LogLevel;
6use crate::schema::sink::sink_message::Payload;
7use crate::schema::sink::{
8    runtime_sink_message::Payload as RuntimeSinkMessagePayload, RuntimeSinkMessage, SinkMessage,
9};
10use crate::sink::plugin::Sink;
11use prost::Message;
12
13pub struct SinkRunnerConfig {
14    pub plugin_id: String,
15    pub log_level: LogLevel,
16    pub hub_connection: ConnectionConfig,
17}
18
19pub struct SinkRunner<T>
20where
21    T: Sink,
22{
23    plugin: T,
24    plugin_id: String,
25    log_level: LogLevel,
26}
27
28impl<T> SinkRunner<T>
29where
30    T: Sink,
31{
32    #[allow(dead_code)]
33    pub async fn initialize(plugin: T, config: SinkRunnerConfig) -> Result<Self, Error> {
34        MSG_CLIENT
35            .write()
36            .await
37            .connect(
38                format!(
39                    "{}:{}",
40                    config.hub_connection.host, config.hub_connection.port
41                )
42                .as_str(),
43            )
44            .await?;
45        PluginLogger::initialize(config.plugin_id.as_str(), config.log_level.into())?;
46        Ok(Self::new(config.plugin_id, plugin, config.log_level))
47    }
48
49    pub(crate) fn new(id: String, plugin: T, log_level: LogLevel) -> Self {
50        Self {
51            plugin,
52            plugin_id: id,
53            log_level,
54        }
55    }
56
57    #[allow(dead_code)]
58    pub async fn run(&mut self) -> Result<(), Error> {
59        // send hello to runtime
60        let payload = match self
61            .plugin
62            .initialize(self.plugin_id.clone(), self.log_level)
63        {
64            Ok(result) => result,
65            Err(err) => {
66                log::error!("Error initializing plugin: {}", err);
67                return Err(Error::InitError(err));
68            }
69        };
70        let hello_msg = SinkMessage {
71            payload: Some(Payload::Initialize(payload.into())),
72        };
73
74        match MSG_CLIENT
75            .read()
76            .await
77            .send(hello_msg.encode_to_vec().as_slice())
78            .await
79        {
80            Ok(_) => {}
81            Err(err) => {
82                log::error!("Error sending hello message: {}", err);
83                return Err(Error::IOError(err));
84            }
85        };
86
87        loop {
88            let bytes = match MSG_CLIENT.read().await.receive().await {
89                Ok(bytes) => match bytes {
90                    None => {
91                        continue;
92                    }
93                    Some(bytes) => bytes,
94                },
95                Err(err) => {
96                    log::error!("Error receiving message: {}", err);
97                    continue;
98                }
99            };
100            let msg = match RuntimeSinkMessage::decode(bytes) {
101                Ok(message) => message,
102                Err(err) => {
103                    log::error!("Error parsing message: {}", err);
104                    continue;
105                }
106            };
107            let pyld = match msg.payload {
108                Some(payload) => payload,
109                None => {
110                    log::error!("Message payload is missing");
111                    continue;
112                }
113            };
114
115            match pyld {
116                RuntimeSinkMessagePayload::Initialize(_) => {
117                    // noop for sink
118                    continue;
119                }
120                RuntimeSinkMessagePayload::Event(payload) => {
121                    log::debug!("Received event: {:?}", payload.plugin_id.clone());
122                    let result = self.plugin.consume_event(payload);
123                    if let Err(err) = result {
124                        log::error!("Error processing event: {}", err);
125                        continue;
126                    }
127                    continue;
128                }
129                RuntimeSinkMessagePayload::Shutdown(_) => {
130                    log::debug!("Received shutdown message");
131                    let result = self.plugin.shutdown();
132                    if let Err(err) = result {
133                        log::error!("Error shutting down: {}", err);
134                        break Err(Error::ShutdownError(err));
135                    }
136                    log::info!("Plugin shutdown: {}", self.plugin_id);
137                    break Ok(());
138                }
139            }
140        }
141    }
142}
flwrs_plugin/sink/runner.rs

flwrs_plugin/sink/
runner.rs
