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        schema_capabilities: crate::schema::resolution::SchemaCapabilities::declared(false),
110        name: "websocket".to_string(),
111        display_name: "WebSocket Source".to_string(),
112        version: env!("CARGO_PKG_VERSION").to_string(),
113        is_source: true,
114        is_sink: false,
115        config_keys: websocket_source_config_keys(),
116    };
117
118    let metric_cache = Arc::new(MetricFamilyCache::<WebSocketSourceMetrics>::default());
119    registry.register_source(
120        "websocket",
121        info,
122        Arc::new(move |registry: Option<&Arc<prometheus::Registry>>| {
123            let metrics = if let Some(registry) = registry {
124                metric_cache
125                    .get_or_try_init(registry, || WebSocketSourceMetrics::register(registry))?
126            } else {
127                WebSocketSourceMetrics::local()
128            };
129            Ok(Box::new(WebSocketSource::new(
130                Arc::new(arrow_schema::Schema::empty()),
131                WebSocketSourceConfig::default(),
132                metrics,
133            )))
134        }),
135    )
136}
137
138/// Registers the WebSocket sink connector with the given registry.
139///
140/// The sink factory selects server or client mode from the validated config
141/// before either implementation performs network I/O.
142///
143/// # Errors
144///
145/// Returns an error when the sink name is already registered or the connector
146/// registry is frozen.
147pub fn register_websocket_sink(
148    registry: &ConnectorRegistry,
149) -> Result<(), crate::error::ConnectorError> {
150    let info = ConnectorInfo {
151        schema_capabilities: crate::schema::resolution::SchemaCapabilities::declared(true),
152        name: "websocket".to_string(),
153        display_name: "WebSocket Sink".to_string(),
154        version: env!("CARGO_PKG_VERSION").to_string(),
155        is_source: false,
156        is_sink: true,
157        config_keys: websocket_sink_config_keys(),
158    };
159
160    let metric_cache = Arc::new(MetricFamilyCache::<WebSocketSinkMetrics>::default());
161    registry.register_sink(
162        "websocket",
163        info,
164        Arc::new(
165            move |config, registry: Option<&Arc<prometheus::Registry>>| {
166                let sink_config = WebSocketSinkConfig::from_config(config)?;
167                let decoded_schema = config.arrow_schema();
168                if config.get("_arrow_schema").is_some() && decoded_schema.is_none() {
169                    return Err(crate::error::ConnectorError::ConfigurationError(
170                        "invalid WebSocket sink _arrow_schema encoding".into(),
171                    ));
172                }
173                let schema = decoded_schema.ok_or_else(|| {
174                    crate::error::ConnectorError::ConfigurationError(
175                        "WebSocket sink requires a declared Arrow schema".into(),
176                    )
177                })?;
178                let metrics = if let Some(registry) = registry {
179                    metric_cache
180                        .get_or_try_init(registry, || WebSocketSinkMetrics::register(registry))?
181                } else {
182                    WebSocketSinkMetrics::local()
183                };
184                let is_server = matches!(&sink_config, WebSocketSinkConfig::Server { .. });
185                let sink: Box<dyn crate::connector::SinkConnector> = if is_server {
186                    Box::new(WebSocketSinkServer::new(schema, sink_config, metrics))
187                } else {
188                    Box::new(WebSocketSinkClient::new(schema, sink_config, metrics))
189                };
190                Ok(sink)
191            },
192        ),
193    )
194}
195
196fn websocket_source_config_keys() -> Vec<ConfigKeySpec> {
197    vec![
198        ConfigKeySpec::required("url", "WebSocket URL to connect to (ws:// or wss://)"),
199        ConfigKeySpec::optional("format", "Message format (json/csv/binary)", "json"),
200        ConfigKeySpec::optional(
201            "subscribe.message",
202            "Text subscription message to send after handshake",
203            "",
204        ),
205        ConfigKeySpec::optional("reconnect.enabled", "Enable automatic reconnection", "true"),
206        ConfigKeySpec::optional(
207            "reconnect.initial.delay.ms",
208            "Initial reconnect delay in ms",
209            "100",
210        ),
211        ConfigKeySpec::optional(
212            "reconnect.max.delay.ms",
213            "Maximum reconnect delay in ms",
214            "30000",
215        ),
216        ConfigKeySpec::optional(
217            "reconnect.max.retries",
218            "Maximum reconnect attempts; empty means unlimited",
219            "",
220        ),
221        ConfigKeySpec::optional(
222            "on.backpressure",
223            "Backpressure strategy (block/drop_newest)",
224            "block",
225        ),
226        ConfigKeySpec::optional(
227            "max.message.size",
228            "Max WebSocket message size in bytes",
229            "67108864",
230        ),
231    ]
232}
233
234fn websocket_sink_config_keys() -> Vec<ConfigKeySpec> {
235    vec![
236        ConfigKeySpec::optional(
237            "bind.address",
238            "Socket address required in server mode (e.g., 0.0.0.0:8080)",
239            "",
240        ),
241        ConfigKeySpec::optional("mode", "Operating mode (server/client)", "server"),
242        ConfigKeySpec::optional(
243            "max.connections",
244            "Max concurrent client connections",
245            "10000",
246        ),
247        ConfigKeySpec::optional("ping.interval.ms", "Ping interval in ms", "30000"),
248        ConfigKeySpec::optional("ping.timeout.ms", "Pong timeout in ms", "10000"),
249        ConfigKeySpec::optional("url", "WebSocket URL required in client mode", ""),
250    ]
251}
252
253#[cfg(test)]
254mod tests;
laminar_connectors/websocket/mod.rs

laminar_connectors/websocket/
mod.rs
