laminar_connectors/websocket/
mod.rs

1//! WebSocket client source plus client/server sinks.
2//!
3//! WebSocket is non-replayable. Sources and sinks are best-effort, and failures
4//! can produce gaps or duplicates.
5
6mod backpressure;
7mod connection;
8mod fanout;
9mod metrics;
10mod parser;
11mod protocol;
12mod serializer;
13mod sink;
14mod sink_client;
15mod sink_config;
16mod sink_metrics;
17mod source;
18mod source_config;
19
20use std::sync::{Arc, OnceLock, Weak};
21
22use metrics::WebSocketSourceMetrics;
23use sink::WebSocketSinkServer;
24use sink_client::WebSocketSinkClient;
25use sink_config::WebSocketSinkConfig;
26use sink_metrics::WebSocketSinkMetrics;
27use source::WebSocketSource;
28use source_config::WebSocketSourceConfig;
29
30use crate::config::{ConfigKeySpec, ConnectorInfo};
31use crate::registry::ConnectorRegistry;
32
33struct RegisteredMetricFamily<T> {
34    registry: Weak<prometheus::Registry>,
35    metrics: T,
36}
37
38struct MetricFamilyCache<T> {
39    family: OnceLock<RegisteredMetricFamily<T>>,
40    initialization: parking_lot::Mutex<()>,
41}
42
43impl<T> Default for MetricFamilyCache<T> {
44    fn default() -> Self {
45        Self {
46            family: OnceLock::new(),
47            initialization: parking_lot::Mutex::new(()),
48        }
49    }
50}
51
52impl<T: Clone> MetricFamilyCache<T> {
53    fn get_or_try_init(
54        &self,
55        registry: &Arc<prometheus::Registry>,
56        initialize: impl FnOnce() -> Result<T, crate::error::ConnectorError>,
57    ) -> Result<T, crate::error::ConnectorError> {
58        if let Some(family) = self.family.get() {
59            return family.for_registry(registry);
60        }
61        let _guard = self.initialization.lock();
62        if let Some(family) = self.family.get() {
63            return family.for_registry(registry);
64        }
65        let family = RegisteredMetricFamily {
66            registry: Arc::downgrade(registry),
67            metrics: initialize()?,
68        };
69        self.family
70            .set(family)
71            .unwrap_or_else(|_| unreachable!("metric initialization is serialized"));
72        Ok(self
73            .family
74            .get()
75            .expect("metric family was initialized")
76            .metrics
77            .clone())
78    }
79}
80
81impl<T: Clone> RegisteredMetricFamily<T> {
82    fn for_registry(
83        &self,
84        registry: &Arc<prometheus::Registry>,
85    ) -> Result<T, crate::error::ConnectorError> {
86        if !Weak::ptr_eq(&self.registry, &Arc::downgrade(registry)) {
87            return Err(crate::error::ConnectorError::ConfigurationError(
88                "WebSocket connector registry is already bound to a different Prometheus registry"
89                    .into(),
90            ));
91        }
92        Ok(self.metrics.clone())
93    }
94}
95
96/// Registers the WebSocket source connector with the given registry.
97///
98/// After registration, the runtime can instantiate `WebSocketSource` by
99/// name when processing `CREATE SOURCE ... FROM WEBSOCKET (...)`.
100///
101/// # Errors
102///
103/// Returns an error when the source name is already registered or the
104/// connector registry is frozen.
105pub fn register_websocket_source(
106    registry: &ConnectorRegistry,
107) -> Result<(), crate::error::ConnectorError> {
108    let info = ConnectorInfo {
109        name: "websocket".to_string(),
110        display_name: "WebSocket Source".to_string(),
111        version: env!("CARGO_PKG_VERSION").to_string(),
112        is_source: true,
113        is_sink: false,
114        config_keys: websocket_source_config_keys(),
115    };
116
117    let metric_cache = Arc::new(MetricFamilyCache::<WebSocketSourceMetrics>::default());
118    registry.register_source(
119        "websocket",
120        info,
121        Arc::new(move |registry: Option<&Arc<prometheus::Registry>>| {
122            let metrics = if let Some(registry) = registry {
123                metric_cache
124                    .get_or_try_init(registry, || WebSocketSourceMetrics::register(registry))?
125            } else {
126                WebSocketSourceMetrics::local()
127            };
128            Ok(Box::new(WebSocketSource::new(
129                Arc::new(arrow_schema::Schema::empty()),
130                WebSocketSourceConfig::default(),
131                metrics,
132            )))
133        }),
134    )
135}
136
137/// Registers the WebSocket sink connector with the given registry.
138///
139/// The sink factory selects server or client mode from the validated config
140/// before either implementation performs network I/O.
141///
142/// # Errors
143///
144/// Returns an error when the sink name is already registered or the connector
145/// registry is frozen.
146pub fn register_websocket_sink(
147    registry: &ConnectorRegistry,
148) -> Result<(), crate::error::ConnectorError> {
149    let info = ConnectorInfo {
150        name: "websocket".to_string(),
151        display_name: "WebSocket Sink".to_string(),
152        version: env!("CARGO_PKG_VERSION").to_string(),
153        is_source: false,
154        is_sink: true,
155        config_keys: websocket_sink_config_keys(),
156    };
157
158    let metric_cache = Arc::new(MetricFamilyCache::<WebSocketSinkMetrics>::default());
159    registry.register_sink(
160        "websocket",
161        info,
162        Arc::new(
163            move |config, registry: Option<&Arc<prometheus::Registry>>| {
164                let sink_config = WebSocketSinkConfig::from_config(config)?;
165                let decoded_schema = config.arrow_schema();
166                if config.get("_arrow_schema").is_some() && decoded_schema.is_none() {
167                    return Err(crate::error::ConnectorError::ConfigurationError(
168                        "invalid WebSocket sink _arrow_schema encoding".into(),
169                    ));
170                }
171                let schema = decoded_schema.ok_or_else(|| {
172                    crate::error::ConnectorError::ConfigurationError(
173                        "WebSocket sink requires a declared Arrow schema".into(),
174                    )
175                })?;
176                let metrics = if let Some(registry) = registry {
177                    metric_cache
178                        .get_or_try_init(registry, || WebSocketSinkMetrics::register(registry))?
179                } else {
180                    WebSocketSinkMetrics::local()
181                };
182                let is_server = matches!(&sink_config, WebSocketSinkConfig::Server { .. });
183                let sink: Box<dyn crate::connector::SinkConnector> = if is_server {
184                    Box::new(WebSocketSinkServer::new(schema, sink_config, metrics))
185                } else {
186                    Box::new(WebSocketSinkClient::new(schema, sink_config, metrics))
187                };
188                Ok(sink)
189            },
190        ),
191    )
192}
193
194fn websocket_source_config_keys() -> Vec<ConfigKeySpec> {
195    vec![
196        ConfigKeySpec::required("url", "WebSocket URL to connect to (ws:// or wss://)"),
197        ConfigKeySpec::optional("format", "Message format (json/csv/binary)", "json"),
198        ConfigKeySpec::optional(
199            "subscribe.message",
200            "Text subscription message to send after handshake",
201            "",
202        ),
203        ConfigKeySpec::optional("reconnect.enabled", "Enable automatic reconnection", "true"),
204        ConfigKeySpec::optional(
205            "reconnect.initial.delay.ms",
206            "Initial reconnect delay in ms",
207            "100",
208        ),
209        ConfigKeySpec::optional(
210            "reconnect.max.delay.ms",
211            "Maximum reconnect delay in ms",
212            "30000",
213        ),
214        ConfigKeySpec::optional(
215            "reconnect.max.retries",
216            "Maximum reconnect attempts; empty means unlimited",
217            "",
218        ),
219        ConfigKeySpec::optional(
220            "on.backpressure",
221            "Backpressure strategy (block/drop_newest)",
222            "block",
223        ),
224        ConfigKeySpec::optional(
225            "max.message.size",
226            "Max WebSocket message size in bytes",
227            "67108864",
228        ),
229    ]
230}
231
232fn websocket_sink_config_keys() -> Vec<ConfigKeySpec> {
233    vec![
234        ConfigKeySpec::optional(
235            "bind.address",
236            "Socket address required in server mode (e.g., 0.0.0.0:8080)",
237            "",
238        ),
239        ConfigKeySpec::optional("mode", "Operating mode (server/client)", "server"),
240        ConfigKeySpec::optional(
241            "max.connections",
242            "Max concurrent client connections",
243            "10000",
244        ),
245        ConfigKeySpec::optional("ping.interval.ms", "Ping interval in ms", "30000"),
246        ConfigKeySpec::optional("ping.timeout.ms", "Pong timeout in ms", "10000"),
247        ConfigKeySpec::optional("url", "WebSocket URL required in client mode", ""),
248    ]
249}
250
251#[cfg(test)]
252mod tests;
laminar_connectors/websocket/mod.rs

laminar_connectors/websocket/
mod.rs
