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

flwrs_plugin/transform/
runner.rs
