laminar_connectors/otel/source/
mod.rs

1//! `SourceConnector` implementation for the OTel OTLP/gRPC receiver.
2//!
3//! Spawns a tonic gRPC server that accepts OTLP export RPCs for traces,
4//! metrics, or logs (one signal type per source), converts to Arrow
5//! `RecordBatch`, and delivers them via `poll_batch()`.
6
7use std::sync::atomic::{AtomicU64, Ordering};
8use std::sync::Arc;
9use std::time::Duration;
10
11use arrow_array::RecordBatch;
12use arrow_schema::SchemaRef;
13use async_trait::async_trait;
14use crossfire::{mpsc, AsyncRx, TryRecvError};
15use tokio::net::TcpListener;
16use tokio::sync::{watch, Notify};
17use tokio::task::JoinHandle;
18use tonic::transport::server::TcpIncoming;
19
20use opentelemetry_proto::tonic::collector::logs::v1::logs_service_server::LogsServiceServer;
21use opentelemetry_proto::tonic::collector::metrics::v1::metrics_service_server::MetricsServiceServer;
22use opentelemetry_proto::tonic::collector::trace::v1::trace_service_server::TraceServiceServer;
23
24use crate::checkpoint::SourceCheckpoint;
25use crate::config::{ConnectorConfig, ConnectorState};
26use crate::connector::{
27    ConnectorTaskOwner, ConnectorTaskTracker, SourceBatch, SourceConnector, SourceConsistency,
28    SourceContract, SourceInputMode, SourcePosition, SourceStart, SourceTopology,
29};
30use crate::error::ConnectorError;
31
32use super::config::{OtelSignal, OtelSourceConfig};
33use super::schema::{logs_schema, metrics_schema, traces_schema};
34use super::server::OtelReceiver;
35
36const SERVER_CLOSE_TIMEOUT: Duration = Duration::from_secs(5);
37
38struct TrackedServerTask {
39    handle: Option<JoinHandle<Result<(), ConnectorError>>>,
40}
41
42struct TaskExitNotify(Arc<Notify>);
43
44impl Drop for TaskExitNotify {
45    fn drop(&mut self) {
46        self.0.notify_one();
47    }
48}
49
50enum ServerWait {
51    Completed(Result<(), ConnectorError>),
52    TimedOut,
53}
54
55impl TrackedServerTask {
56    fn spawn(
57        owner: &ConnectorTaskOwner,
58        exit_notify: Arc<Notify>,
59        future: impl std::future::Future<Output = Result<(), ConnectorError>> + Send + 'static,
60    ) -> Result<Self, ConnectorError> {
61        let task_guard = owner.track().ok_or_else(|| {
62            ConnectorError::Internal("OTel source task generation is already retired".into())
63        })?;
64        let handle = tokio::spawn(async move {
65            let _task_guard = task_guard;
66            let _exit_notify = TaskExitNotify(exit_notify);
67            future.await
68        });
69        Ok(Self {
70            handle: Some(handle),
71        })
72    }
73
74    async fn wait_until(&mut self, deadline: tokio::time::Instant) -> ServerWait {
75        let Some(handle) = self.handle.as_mut() else {
76            return ServerWait::Completed(Ok(()));
77        };
78        match tokio::time::timeout_at(deadline, handle).await {
79            Ok(_) if tokio::time::Instant::now() >= deadline => {
80                self.handle.take();
81                ServerWait::TimedOut
82            }
83            Ok(Ok(result)) => {
84                self.handle.take();
85                ServerWait::Completed(result)
86            }
87            Ok(Err(error)) => {
88                self.handle.take();
89                ServerWait::Completed(Err(ConnectorError::Internal(format!(
90                    "OTel gRPC server task failed: {error}"
91                ))))
92            }
93            Err(_) => ServerWait::TimedOut,
94        }
95    }
96
97    async fn take_finished(&mut self) -> Option<Result<(), ConnectorError>> {
98        let handle = self.handle.as_ref()?;
99        if !handle.is_finished() {
100            return None;
101        }
102        let result = self.handle.take()?.await;
103        Some(result.unwrap_or_else(|error| {
104            Err(ConnectorError::Internal(format!(
105                "OTel gRPC server task failed: {error}"
106            )))
107        }))
108    }
109
110    fn abort(&mut self) {
111        if let Some(handle) = self.handle.take() {
112            handle.abort();
113        }
114    }
115}
116
117/// OTel OTLP/gRPC source connector.
118///
119/// Binds a gRPC server on the configured port and receives telemetry
120/// data from `OpenTelemetry` exporters and collectors. Each source
121/// handles exactly one signal type (traces, metrics, or logs).
122pub struct OtelSource {
123    config: OtelSourceConfig,
124    schema: SchemaRef,
125    state: ConnectorState,
126    batch_rx: Option<AsyncRx<mpsc::Array<RecordBatch>>>,
127    data_ready: Arc<Notify>,
128    server_task: Option<TrackedServerTask>,
129    shutdown_tx: Option<watch::Sender<bool>>,
130    /// Monotonic counter of records received (spans, data points, or log records).
131    records_received: Arc<AtomicU64>,
132    requests_received: Arc<AtomicU64>,
133    checkpoint_seq: u64,
134    server_failure: Option<String>,
135    task_owner: ConnectorTaskOwner,
136    task_tracker: ConnectorTaskTracker,
137}
138
139impl OtelSource {
140    /// Create a new OTel source with the given default schema.
141    #[must_use]
142    pub fn new(schema: SchemaRef, _registry: Option<&prometheus::Registry>) -> Self {
143        let (task_owner, task_tracker) = ConnectorTaskOwner::new();
144        Self {
145            config: OtelSourceConfig::default(),
146            schema,
147            state: ConnectorState::Created,
148            batch_rx: None,
149            data_ready: Arc::new(Notify::new()),
150            server_task: None,
151            shutdown_tx: None,
152            records_received: Arc::new(AtomicU64::new(0)),
153            requests_received: Arc::new(AtomicU64::new(0)),
154            checkpoint_seq: 0,
155            server_failure: None,
156            task_owner,
157            task_tracker,
158        }
159    }
160
161    fn request_shutdown(&mut self) {
162        if let Some(shutdown) = self.shutdown_tx.take() {
163            shutdown.send_replace(true);
164        }
165        self.batch_rx.take();
166    }
167
168    async fn observe_server_exit(&mut self) {
169        if self.server_failure.is_some() || self.state != ConnectorState::Running {
170            return;
171        }
172        let Some(server) = self.server_task.as_mut() else {
173            return;
174        };
175        let Some(result) = server.take_finished().await else {
176            return;
177        };
178        self.server_task.take();
179        self.server_failure = Some(match result {
180            Ok(()) => "OTel gRPC server stopped unexpectedly".into(),
181            Err(ConnectorError::ConnectionFailed(message) | ConnectorError::Internal(message)) => {
182                message
183            }
184            Err(error) => error.to_string(),
185        });
186        self.state = ConnectorState::Failed;
187    }
188
189    fn terminal_server_error(message: &str) -> ConnectorError {
190        ConnectorError::InvalidState {
191            expected: "live OTLP gRPC server".into(),
192            actual: format!("server generation terminated: {message}"),
193        }
194    }
195}
196
197impl Drop for OtelSource {
198    fn drop(&mut self) {
199        self.request_shutdown();
200        if let Some(server) = self.server_task.as_mut() {
201            server.abort();
202        }
203    }
204}
205
206#[async_trait]
207impl SourceConnector for OtelSource {
208    fn terminal_task_tracker(&self) -> Option<ConnectorTaskTracker> {
209        Some(self.task_tracker.clone())
210    }
211
212    async fn resolve_schema(
213        &mut self,
214        config: &ConnectorConfig,
215        explicit: Option<SchemaRef>,
216    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
217        self.discover_schema(config.properties()).await?;
218        crate::schema::resolution::fixed_binding(config, explicit, &self.schema())
219    }
220
221    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
222        let (config, position, _) = request.into_parts();
223        if matches!(&position, SourcePosition::Initialized { .. }) {
224            return Err(ConnectorError::ConfigurationError(
225                "OTLP has no sealed topology startup contract".into(),
226            ));
227        }
228        if let SourcePosition::Resume { attempt, .. } = position {
229            return Err(ConnectorError::ConfigurationError(format!(
230                "OTLP is an ephemeral source and cannot resume checkpoint attempt {attempt:?}"
231            )));
232        }
233        if !matches!(self.state, ConnectorState::Created | ConnectorState::Closed)
234            || self.server_task.is_some()
235        {
236            return Err(ConnectorError::InvalidState {
237                expected: "Created or fully closed".into(),
238                actual: format!("{}", self.state),
239            });
240        }
241
242        let candidate_config = OtelSourceConfig::from_config(&config)?;
243
244        let protocol = match candidate_config.signals {
245            OtelSignal::Traces => traces_schema(),
246            OtelSignal::Metrics => metrics_schema(),
247            OtelSignal::Logs => logs_schema(),
248        };
249
250        let binding =
251            crate::schema::resolution::fixed_binding(&config, config.arrow_schema(), &protocol)?;
252        let candidate_schema = Arc::new(binding.logical);
253        let (batch_tx, batch_rx) =
254            mpsc::bounded_async::<RecordBatch>(candidate_config.channel_capacity);
255
256        let addr = candidate_config.socket_addr();
257
258        // Bind the TCP listener here so port conflicts fail start(),
259        // not silently inside the background task.
260        let listener = TcpListener::bind(&addr)
261            .await
262            .map_err(|e| ConnectorError::ConnectionFailed(format!("failed to bind {addr}: {e}")))?;
263        let incoming = TcpIncoming::from(listener).with_nodelay(Some(true));
264
265        let (shutdown_tx, shutdown_rx) = watch::channel(false);
266        let service_guard = self.task_owner.track().ok_or_else(|| {
267            ConnectorError::Internal("OTel source task generation is already retired".into())
268        })?;
269
270        let receiver = OtelReceiver::new(
271            batch_tx,
272            Arc::clone(&candidate_schema),
273            Arc::clone(&self.data_ready),
274            Arc::clone(&self.records_received),
275            Arc::clone(&self.requests_received),
276            candidate_config.batch_size,
277            service_guard,
278        );
279
280        // Spawn signal-specific gRPC server on the already-bound listener
281        let server_task = match candidate_config.signals {
282            OtelSignal::Traces => spawn_grpc_server(
283                &self.task_owner,
284                TraceServiceServer::new(receiver),
285                incoming,
286                shutdown_rx,
287                Arc::clone(&self.data_ready),
288            ),
289            OtelSignal::Metrics => spawn_grpc_server(
290                &self.task_owner,
291                MetricsServiceServer::new(receiver),
292                incoming,
293                shutdown_rx,
294                Arc::clone(&self.data_ready),
295            ),
296            OtelSignal::Logs => spawn_grpc_server(
297                &self.task_owner,
298                LogsServiceServer::new(receiver),
299                incoming,
300                shutdown_rx,
301                Arc::clone(&self.data_ready),
302            ),
303        }?;
304
305        self.config = candidate_config;
306        self.schema = candidate_schema;
307        self.batch_rx = Some(batch_rx);
308        self.shutdown_tx = Some(shutdown_tx);
309        self.server_task = Some(server_task);
310        self.server_failure = None;
311        self.state = ConnectorState::Running;
312
313        tracing::info!(
314            %addr,
315            signals = ?self.config.signals,
316            batch_size = self.config.batch_size,
317            "OTel source connector started"
318        );
319
320        Ok(())
321    }
322
323    async fn poll_batch(
324        &mut self,
325        max_records: usize,
326    ) -> Result<Option<SourceBatch>, ConnectorError> {
327        if let Some(error) = &self.server_failure {
328            return Err(Self::terminal_server_error(error));
329        }
330        let rx = self.batch_rx.as_ref().ok_or(ConnectorError::InvalidState {
331            expected: "Running".into(),
332            actual: format!("{}", self.state),
333        })?;
334
335        let mut total_rows = 0usize;
336        let mut batches: Vec<RecordBatch> = Vec::new();
337        let mut disconnected = false;
338
339        loop {
340            match rx.try_recv() {
341                Ok(batch) => {
342                    total_rows += batch.num_rows();
343                    batches.push(batch);
344                    if total_rows >= max_records {
345                        break;
346                    }
347                }
348                Err(TryRecvError::Empty) => break,
349                Err(TryRecvError::Disconnected) => {
350                    disconnected = true;
351                    break;
352                }
353            }
354        }
355
356        self.observe_server_exit().await;
357
358        if batches.is_empty() {
359            return if let Some(error) = &self.server_failure {
360                Err(Self::terminal_server_error(error))
361            } else if disconnected {
362                self.state = ConnectorState::Closed;
363                Err(ConnectorError::Closed)
364            } else {
365                Ok(None)
366            };
367        }
368
369        self.checkpoint_seq += 1;
370
371        if batches.len() == 1 {
372            return Ok(Some(SourceBatch::new(
373                crate::schema::resolution::project_batch(&batches.remove(0), &self.schema)?,
374            )));
375        }
376
377        let schema = batches[0].schema();
378        let combined =
379            arrow_select::concat::concat_batches(&schema, batches.iter()).map_err(|e| {
380                ConnectorError::ReadError(format!("failed to concatenate OTel batches: {e}"))
381            })?;
382
383        Ok(Some(SourceBatch::new(
384            crate::schema::resolution::project_batch(&combined, &self.schema)?,
385        )))
386    }
387
388    async fn discover_schema(
389        &mut self,
390        properties: &std::collections::HashMap<String, String>,
391    ) -> Result<(), ConnectorError> {
392        let Some(sig) = properties
393            .get("signals")
394            .or_else(|| properties.get("signal"))
395        else {
396            return Ok(());
397        };
398        let signal = OtelSignal::parse(sig).map_err(|e| {
399            ConnectorError::ConfigurationError(format!("invalid OTel signal '{sig}': {e}"))
400        })?;
401        self.schema = match signal {
402            OtelSignal::Traces => traces_schema(),
403            OtelSignal::Metrics => metrics_schema(),
404            OtelSignal::Logs => logs_schema(),
405        };
406        Ok(())
407    }
408
409    fn schema(&self) -> SchemaRef {
410        Arc::clone(&self.schema)
411    }
412
413    fn checkpoint(&self) -> SourceCheckpoint {
414        let mut cp = SourceCheckpoint::new();
415        cp.set_offset("batch_sequence", self.checkpoint_seq.to_string());
416        cp.set_offset(
417            "records_received",
418            self.records_received.load(Ordering::Relaxed).to_string(),
419        );
420        cp.set_offset(
421            "requests_received",
422            self.requests_received.load(Ordering::Relaxed).to_string(),
423        );
424        cp.set_metadata("connector", "otel");
425        cp.set_metadata("signals", format!("{:?}", self.config.signals));
426        cp
427    }
428
429    async fn close(&mut self) -> Result<(), ConnectorError> {
430        tracing::info!("OTel source connector shutting down");
431
432        self.request_shutdown();
433
434        let mut completed = false;
435        let mut close_error = self
436            .server_failure
437            .as_ref()
438            .map(|error| Self::terminal_server_error(error));
439        if let Some(server) = self.server_task.as_mut() {
440            let deadline = tokio::time::Instant::now() + SERVER_CLOSE_TIMEOUT;
441            match server.wait_until(deadline).await {
442                ServerWait::Completed(Ok(())) => completed = true,
443                ServerWait::Completed(Err(error)) => {
444                    completed = true;
445                    tracing::warn!(%error, "OTel gRPC server task failed while closing");
446                    close_error = Some(error);
447                }
448                ServerWait::TimedOut => {
449                    server.abort();
450                    completed = true;
451                    close_error = Some(ConnectorError::Internal(
452                        "OTel gRPC server exceeded its close deadline; connector generation retired"
453                            .into(),
454                    ));
455                    tracing::warn!("OTel gRPC server exceeded its close deadline and was aborted");
456                }
457            }
458        }
459        if completed {
460            self.server_task.take();
461        }
462
463        if let Some(error) = close_error {
464            self.state = ConnectorState::Failed;
465            Err(error)
466        } else {
467            self.state = ConnectorState::Closed;
468            Ok(())
469        }
470    }
471
472    fn data_ready_notify(&self) -> Option<Arc<Notify>> {
473        Some(Arc::clone(&self.data_ready))
474    }
475
476    fn contract(&self, _config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
477        Ok(SourceContract::new(
478            SourceConsistency::Ephemeral,
479            SourceTopology::NodeLocalIngress,
480            SourceInputMode::AppendOnly,
481        ))
482    }
483}
484
485/// Spawn a tonic gRPC server on a pre-bound listener with graceful shutdown.
486fn spawn_grpc_server<S>(
487    owner: &ConnectorTaskOwner,
488    svc: S,
489    incoming: TcpIncoming,
490    mut shutdown_rx: watch::Receiver<bool>,
491    data_ready: Arc<Notify>,
492) -> Result<TrackedServerTask, ConnectorError>
493where
494    S: tonic::codegen::Service<
495            tonic::codegen::http::Request<tonic::body::Body>,
496            Response = tonic::codegen::http::Response<tonic::body::Body>,
497            Error = std::convert::Infallible,
498        > + tonic::server::NamedService
499        + Clone
500        + Send
501        + Sync
502        + 'static,
503    S::Future: Send + 'static,
504{
505    TrackedServerTask::spawn(owner, data_ready, async move {
506        tonic::transport::Server::builder()
507            .add_service(svc)
508            .serve_with_incoming_shutdown(incoming, async move {
509                let _ = shutdown_rx.wait_for(|&v| v).await;
510            })
511            .await
512            .map_err(|error| {
513                tracing::error!(%error, "OTel gRPC server exited with error");
514                ConnectorError::ConnectionFailed(format!(
515                    "OTel gRPC server exited with error: {error}"
516                ))
517            })
518    })
519}
520
521impl std::fmt::Debug for OtelSource {
522    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
523        f.debug_struct("OtelSource")
524            .field("state", &self.state)
525            .field("config", &self.config)
526            .field(
527                "records_received",
528                &self.records_received.load(Ordering::Relaxed),
529            )
530            .field(
531                "requests_received",
532                &self.requests_received.load(Ordering::Relaxed),
533            )
534            .field("checkpoint_seq", &self.checkpoint_seq)
535            .field(
536                "server_running",
537                &self
538                    .server_task
539                    .as_ref()
540                    .and_then(|task| task.handle.as_ref().map(|handle| !handle.is_finished())),
541            )
542            .field("has_shutdown_tx", &self.shutdown_tx.is_some())
543            .finish_non_exhaustive()
544    }
545}
546
547#[cfg(test)]
548mod tests;
laminar_connectors/otel/source/mod.rs

laminar_connectors/otel/source/
mod.rs
