cloudpub_client/plugins/
plugin_trait.rs

1use crate::config::{ClientConfig, ClientOpts};
2use crate::shell::SubProcess;
3use anyhow::{bail, Context, Result};
4use async_trait::async_trait;
5use cloudpub_common::fair_channel::FairSender;
6use cloudpub_common::protocol::message::Message;
7use cloudpub_common::protocol::{Break, ErrorInfo, ErrorKind, ServerEndpoint};
8use cloudpub_common::utils::is_tcp_port_available;
9use parking_lot::RwLock;
10use std::sync::Arc;
11use tokio::sync::mpsc;
12use tokio::task::JoinHandle;
13use tokio::time::{self, Duration, Instant};
14use tracing::{debug, error};
15
16/// ClientEndpoint intentionally has no PartialEq (see common/build.rs);
17/// its deterministic prost encoding stands in for config equality.
18fn encode_client_config(endpoint: &ServerEndpoint) -> Vec<u8> {
19    use cloudpub_common::prost::Message as ProstMessage;
20    endpoint
21        .client
22        .as_ref()
23        .map(|c| c.encode_to_vec())
24        .unwrap_or_default()
25}
26
27#[async_trait]
28pub trait Plugin: Send + Sync {
29    /// Name of the plugin
30    fn name(&self) -> &'static str;
31
32    /// Setup the plugin environment
33    async fn setup(
34        &self,
35        config: &Arc<RwLock<ClientConfig>>,
36        opts: &ClientOpts,
37        command_rx: &mut mpsc::Receiver<Message>,
38        result_tx: &mpsc::Sender<Message>,
39    ) -> Result<()>;
40
41    /// Publish a service using this plugin
42    async fn publish(
43        &self,
44        endpoint: &ServerEndpoint,
45        config: &Arc<RwLock<ClientConfig>>,
46        opts: &ClientOpts,
47        result_tx: &mpsc::Sender<Message>,
48    ) -> Result<SubProcess>;
49}
50
51pub struct PluginHandle {
52    guid: String,
53    /// prost-encoded ClientEndpoint this handle was spawned for, used to
54    /// recognize duplicate EndpointStart deliveries for the same config
55    client_config: Vec<u8>,
56    server: Arc<RwLock<Option<SubProcess>>>,
57    cancel_tx: mpsc::Sender<Message>,
58    task: JoinHandle<()>,
59}
60
61impl PluginHandle {
62    pub fn spawn(
63        plugin: Arc<dyn Plugin>,
64        mut endpoint: ServerEndpoint,
65        config: Arc<RwLock<ClientConfig>>,
66        opts: ClientOpts,
67        to_server_tx: FairSender<Message>,
68    ) -> Self {
69        let guid = endpoint.guid.clone();
70        let client_config = encode_client_config(&endpoint);
71        let guid_for_task = guid.clone();
72        let server = Arc::new(RwLock::new(None));
73        let server_task = server.clone();
74
75        // Per-process cancel and event channels
76        let (cancel_tx, mut cancel_rx) = mpsc::channel::<Message>(1);
77        let (proc_event_tx, mut proc_event_rx) = mpsc::channel::<Message>(1024);
78
79        let task = tokio::spawn(async move {
80            // Forward per-process events to the main client event channel
81            tokio::spawn({
82                let mut endpoint = endpoint.clone();
83                let to_server_tx = to_server_tx.clone();
84                let guid_for_task = guid_for_task.clone();
85                async move {
86                    use tokio::time::{Duration, Instant};
87
88                    let mut last_progress_time = Instant::now() - Duration::from_secs(1); // Allow first progress immediately
89
90                    while let Some(mut msg) = proc_event_rx.recv().await {
91                        let should_send = match &mut msg {
92                            Message::Progress(progress_info) => {
93                                progress_info.guid = guid_for_task.clone();
94
95                                let now = Instant::now();
96                                // Always send 0% and 100% progress messages, or throttle to 1 per second
97                                if progress_info.current == 0
98                                    || progress_info.current >= progress_info.total
99                                    || now.duration_since(last_progress_time)
100                                        >= Duration::from_secs(1)
101                                {
102                                    last_progress_time = now;
103                                    true
104                                } else {
105                                    false
106                                }
107                            }
108                            _ => true, // Always send non-progress messages
109                        };
110
111                        if should_send {
112                            to_server_tx.send(msg.clone()).await.ok();
113                        }
114                    }
115
116                    endpoint.status = Some("offline".to_string());
117                    to_server_tx
118                        .send(Message::EndpointStatus(endpoint.clone()))
119                        .await
120                        .ok();
121                }
122            });
123
124            // Initial lifecycle status: waiting
125            endpoint.status = Some("waiting".to_string());
126            let _ = to_server_tx
127                .send(Message::EndpointStatus(endpoint.clone()))
128                .await;
129
130            let res: Result<()> = async {
131                // Long running setup: cancellable via cancel_rx
132                plugin
133                    .setup(&config, &opts, &mut cancel_rx, &proc_event_tx)
134                    .await
135                    .context("Failed to setup plugin")?;
136
137                let server_process = plugin
138                    .publish(&endpoint, &config, &opts, &proc_event_tx)
139                    .await
140                    .context("Failed to publish plugin service")?;
141
142                // Wait until the port is bound
143                let now = Instant::now();
144                while is_tcp_port_available("127.0.0.1", server_process.port)
145                    .await
146                    .context("Check port availability")?
147                {
148                    if server_process.result.read().is_err() {
149                        return Ok(()); // Error already reported by subprocess
150                    }
151
152                    if now.elapsed() > Duration::from_secs(60) {
153                        bail!("{}", crate::t!("error-start-server"));
154                    }
155                    debug!(
156                        "Waiting for server to start on port {}",
157                        server_process.port
158                    );
159                    time::sleep(Duration::from_secs(1)).await;
160                }
161
162                *server_task.write() = Some(server_process);
163                Ok(())
164            }
165            .await;
166
167            match res {
168                Ok(()) => {
169                    endpoint.status = Some("online".into());
170                    let _ = to_server_tx.send(Message::EndpointStatus(endpoint)).await;
171                }
172                Err(e) => {
173                    error!("Error handling endpoint {}: {:#}", &guid_for_task, e);
174                    let _ = to_server_tx
175                        .send(Message::Error(ErrorInfo {
176                            kind: ErrorKind::PublishFailed.into(),
177                            message: e.to_string(),
178                            guid: guid_for_task,
179                        }))
180                        .await;
181                }
182            }
183        });
184
185        Self {
186            guid,
187            client_config,
188            server,
189            cancel_tx,
190            task,
191        }
192    }
193
194    /// True when this handle already serves exactly this endpoint
195    /// configuration: a duplicate EndpointStart (the server re-sends them
196    /// after every EndpointStartAll and on repeated registrations) must
197    /// not bounce a running publication.
198    pub fn same_config(&self, endpoint: &ServerEndpoint) -> bool {
199        self.client_config == encode_client_config(endpoint)
200    }
201
202    /// The publication is being set up or its subprocess is alive. A
203    /// finished task with no healthy subprocess means setup failed or the
204    /// server died - a fresh spawn is warranted then.
205    pub fn is_healthy(&self) -> bool {
206        if !self.task.is_finished() {
207            return true;
208        }
209        match self.server.read().as_ref() {
210            Some(process) => process.result.read().is_ok(),
211            None => false,
212        }
213    }
214
215    pub fn port(&self) -> Option<u16> {
216        self.server.read().as_ref().map(|s| s.port)
217    }
218
219    pub fn guid(&self) -> &str {
220        &self.guid
221    }
222
223    pub fn send_break(&self) {
224        let _ = self.cancel_tx.try_send(Message::Break(Break {
225            guid: self.guid.clone(),
226        }));
227    }
228}
229
230impl Drop for PluginHandle {
231    fn drop(&mut self) {
232        // Abort main task and drop cancel_tx to notify setup to stop
233        self.task.abort();
234        // cancel_tx dropped here; receiver sees end-of-stream and should terminate
235        let _ = self.cancel_tx.try_send(Message::Break(Break {
236            guid: self.guid.clone(),
237        }));
238    }
239}
cloudpub_client/plugins/plugin_trait.rs

cloudpub_client/plugins/
plugin_trait.rs
