laminar_connectors/otel/
source.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, 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 start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
213        let (config, position, _) = request.into_parts();
214        if let SourcePosition::Resume { attempt, .. } = position {
215            return Err(ConnectorError::ConfigurationError(format!(
216                "OTLP is an ephemeral source and cannot resume checkpoint attempt {attempt:?}"
217            )));
218        }
219        if !matches!(self.state, ConnectorState::Created | ConnectorState::Closed)
220            || self.server_task.is_some()
221        {
222            return Err(ConnectorError::InvalidState {
223                expected: "Created or fully closed".into(),
224                actual: format!("{}", self.state),
225            });
226        }
227
228        let candidate_config = OtelSourceConfig::from_config(&config)?;
229
230        let candidate_schema = match candidate_config.signals {
231            OtelSignal::Traces => traces_schema(),
232            OtelSignal::Metrics => metrics_schema(),
233            OtelSignal::Logs => logs_schema(),
234        };
235
236        let (batch_tx, batch_rx) =
237            mpsc::bounded_async::<RecordBatch>(candidate_config.channel_capacity);
238
239        let addr = candidate_config.socket_addr();
240
241        // Bind the TCP listener here so port conflicts fail start(),
242        // not silently inside the background task.
243        let listener = TcpListener::bind(&addr)
244            .await
245            .map_err(|e| ConnectorError::ConnectionFailed(format!("failed to bind {addr}: {e}")))?;
246        let incoming = TcpIncoming::from(listener).with_nodelay(Some(true));
247
248        let (shutdown_tx, shutdown_rx) = watch::channel(false);
249        let service_guard = self.task_owner.track().ok_or_else(|| {
250            ConnectorError::Internal("OTel source task generation is already retired".into())
251        })?;
252
253        let receiver = OtelReceiver::new(
254            batch_tx,
255            Arc::clone(&candidate_schema),
256            Arc::clone(&self.data_ready),
257            Arc::clone(&self.records_received),
258            Arc::clone(&self.requests_received),
259            candidate_config.batch_size,
260            service_guard,
261        );
262
263        // Spawn signal-specific gRPC server on the already-bound listener
264        let server_task = match candidate_config.signals {
265            OtelSignal::Traces => spawn_grpc_server(
266                &self.task_owner,
267                TraceServiceServer::new(receiver),
268                incoming,
269                shutdown_rx,
270                Arc::clone(&self.data_ready),
271            ),
272            OtelSignal::Metrics => spawn_grpc_server(
273                &self.task_owner,
274                MetricsServiceServer::new(receiver),
275                incoming,
276                shutdown_rx,
277                Arc::clone(&self.data_ready),
278            ),
279            OtelSignal::Logs => spawn_grpc_server(
280                &self.task_owner,
281                LogsServiceServer::new(receiver),
282                incoming,
283                shutdown_rx,
284                Arc::clone(&self.data_ready),
285            ),
286        }?;
287
288        self.config = candidate_config;
289        self.schema = candidate_schema;
290        self.batch_rx = Some(batch_rx);
291        self.shutdown_tx = Some(shutdown_tx);
292        self.server_task = Some(server_task);
293        self.server_failure = None;
294        self.state = ConnectorState::Running;
295
296        tracing::info!(
297            %addr,
298            signals = ?self.config.signals,
299            batch_size = self.config.batch_size,
300            "OTel source connector started"
301        );
302
303        Ok(())
304    }
305
306    async fn poll_batch(
307        &mut self,
308        max_records: usize,
309    ) -> Result<Option<SourceBatch>, ConnectorError> {
310        if let Some(error) = &self.server_failure {
311            return Err(Self::terminal_server_error(error));
312        }
313        let rx = self.batch_rx.as_ref().ok_or(ConnectorError::InvalidState {
314            expected: "Running".into(),
315            actual: format!("{}", self.state),
316        })?;
317
318        let mut total_rows = 0usize;
319        let mut batches: Vec<RecordBatch> = Vec::new();
320        let mut disconnected = false;
321
322        loop {
323            match rx.try_recv() {
324                Ok(batch) => {
325                    total_rows += batch.num_rows();
326                    batches.push(batch);
327                    if total_rows >= max_records {
328                        break;
329                    }
330                }
331                Err(TryRecvError::Empty) => break,
332                Err(TryRecvError::Disconnected) => {
333                    disconnected = true;
334                    break;
335                }
336            }
337        }
338
339        self.observe_server_exit().await;
340
341        if batches.is_empty() {
342            return if let Some(error) = &self.server_failure {
343                Err(Self::terminal_server_error(error))
344            } else if disconnected {
345                self.state = ConnectorState::Closed;
346                Err(ConnectorError::Closed)
347            } else {
348                Ok(None)
349            };
350        }
351
352        self.checkpoint_seq += 1;
353
354        if batches.len() == 1 {
355            return Ok(Some(SourceBatch::new(batches.into_iter().next().unwrap())));
356        }
357
358        let schema = batches[0].schema();
359        let combined =
360            arrow_select::concat::concat_batches(&schema, batches.iter()).map_err(|e| {
361                ConnectorError::ReadError(format!("failed to concatenate OTel batches: {e}"))
362            })?;
363
364        Ok(Some(SourceBatch::new(combined)))
365    }
366
367    async fn discover_schema(
368        &mut self,
369        properties: &std::collections::HashMap<String, String>,
370    ) -> Result<(), ConnectorError> {
371        let Some(sig) = properties
372            .get("signals")
373            .or_else(|| properties.get("signal"))
374        else {
375            return Ok(());
376        };
377        let signal = OtelSignal::parse(sig).map_err(|e| {
378            ConnectorError::ConfigurationError(format!("invalid OTel signal '{sig}': {e}"))
379        })?;
380        self.schema = match signal {
381            OtelSignal::Traces => traces_schema(),
382            OtelSignal::Metrics => metrics_schema(),
383            OtelSignal::Logs => logs_schema(),
384        };
385        Ok(())
386    }
387
388    fn schema(&self) -> SchemaRef {
389        Arc::clone(&self.schema)
390    }
391
392    fn checkpoint(&self) -> SourceCheckpoint {
393        let mut cp = SourceCheckpoint::new();
394        cp.set_offset("batch_sequence", self.checkpoint_seq.to_string());
395        cp.set_offset(
396            "records_received",
397            self.records_received.load(Ordering::Relaxed).to_string(),
398        );
399        cp.set_offset(
400            "requests_received",
401            self.requests_received.load(Ordering::Relaxed).to_string(),
402        );
403        cp.set_metadata("connector", "otel");
404        cp.set_metadata("signals", format!("{:?}", self.config.signals));
405        cp
406    }
407
408    async fn close(&mut self) -> Result<(), ConnectorError> {
409        tracing::info!("OTel source connector shutting down");
410
411        self.request_shutdown();
412
413        let mut completed = false;
414        let mut close_error = self
415            .server_failure
416            .as_ref()
417            .map(|error| Self::terminal_server_error(error));
418        if let Some(server) = self.server_task.as_mut() {
419            let deadline = tokio::time::Instant::now() + SERVER_CLOSE_TIMEOUT;
420            match server.wait_until(deadline).await {
421                ServerWait::Completed(Ok(())) => completed = true,
422                ServerWait::Completed(Err(error)) => {
423                    completed = true;
424                    tracing::warn!(%error, "OTel gRPC server task failed while closing");
425                    close_error = Some(error);
426                }
427                ServerWait::TimedOut => {
428                    server.abort();
429                    completed = true;
430                    close_error = Some(ConnectorError::Internal(
431                        "OTel gRPC server exceeded its close deadline; connector generation retired"
432                            .into(),
433                    ));
434                    tracing::warn!("OTel gRPC server exceeded its close deadline and was aborted");
435                }
436            }
437        }
438        if completed {
439            self.server_task.take();
440        }
441
442        if let Some(error) = close_error {
443            self.state = ConnectorState::Failed;
444            Err(error)
445        } else {
446            self.state = ConnectorState::Closed;
447            Ok(())
448        }
449    }
450
451    fn data_ready_notify(&self) -> Option<Arc<Notify>> {
452        Some(Arc::clone(&self.data_ready))
453    }
454
455    fn contract(&self, _config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
456        Ok(SourceContract::new(
457            SourceConsistency::Ephemeral,
458            SourceTopology::NodeLocalIngress,
459        ))
460    }
461}
462
463/// Spawn a tonic gRPC server on a pre-bound listener with graceful shutdown.
464fn spawn_grpc_server<S>(
465    owner: &ConnectorTaskOwner,
466    svc: S,
467    incoming: TcpIncoming,
468    mut shutdown_rx: watch::Receiver<bool>,
469    data_ready: Arc<Notify>,
470) -> Result<TrackedServerTask, ConnectorError>
471where
472    S: tonic::codegen::Service<
473            tonic::codegen::http::Request<tonic::body::Body>,
474            Response = tonic::codegen::http::Response<tonic::body::Body>,
475            Error = std::convert::Infallible,
476        > + tonic::server::NamedService
477        + Clone
478        + Send
479        + Sync
480        + 'static,
481    S::Future: Send + 'static,
482{
483    TrackedServerTask::spawn(owner, data_ready, async move {
484        tonic::transport::Server::builder()
485            .add_service(svc)
486            .serve_with_incoming_shutdown(incoming, async move {
487                let _ = shutdown_rx.wait_for(|&v| v).await;
488            })
489            .await
490            .map_err(|error| {
491                tracing::error!(%error, "OTel gRPC server exited with error");
492                ConnectorError::ConnectionFailed(format!(
493                    "OTel gRPC server exited with error: {error}"
494                ))
495            })
496    })
497}
498
499impl std::fmt::Debug for OtelSource {
500    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
501        f.debug_struct("OtelSource")
502            .field("state", &self.state)
503            .field("config", &self.config)
504            .field(
505                "records_received",
506                &self.records_received.load(Ordering::Relaxed),
507            )
508            .field(
509                "requests_received",
510                &self.requests_received.load(Ordering::Relaxed),
511            )
512            .field("checkpoint_seq", &self.checkpoint_seq)
513            .field(
514                "server_running",
515                &self
516                    .server_task
517                    .as_ref()
518                    .and_then(|task| task.handle.as_ref().map(|handle| !handle.is_finished())),
519            )
520            .field("has_shutdown_tx", &self.shutdown_tx.is_some())
521            .finish_non_exhaustive()
522    }
523}
524
525#[cfg(test)]
526mod tests {
527    use super::*;
528    use crate::connector::DeliveryGuarantee;
529
530    fn start_request() -> SourceStart {
531        let mut config = ConnectorConfig::new("otel");
532        config.set("bind.address", "127.0.0.1");
533        config.set("port", "0");
534        SourceStart::new(
535            config,
536            SourcePosition::Initial,
537            DeliveryGuarantee::BestEffort,
538        )
539        .unwrap()
540    }
541
542    #[tokio::test]
543    async fn clean_close_can_restart_the_same_tracked_generation() {
544        let mut source = OtelSource::new(traces_schema(), None);
545        let terminal = source.terminal_task_tracker().unwrap();
546
547        source.start(start_request()).await.unwrap();
548        assert!(source.server_task.is_some());
549        source.close().await.unwrap();
550        assert!(source.server_task.is_none());
551        assert!(!terminal.is_terminated(), "the generation owner is live");
552
553        source.start(start_request()).await.unwrap();
554        assert!(source.server_task.is_some());
555        drop(source);
556
557        tokio::time::timeout(Duration::from_secs(2), terminal.wait_terminated())
558            .await
559            .expect("restarted OTel server generation did not terminate");
560    }
561
562    #[tokio::test]
563    async fn aborted_close_waiter_aborts_the_owned_server_task() {
564        let mut source = OtelSource::new(traces_schema(), None);
565        let terminal = source.terminal_task_tracker().unwrap();
566        let (started_tx, started_rx) = tokio::sync::oneshot::channel();
567        let (release_tx, release_rx) = tokio::sync::oneshot::channel();
568        source.server_task = Some(
569            TrackedServerTask::spawn(
570                &source.task_owner,
571                Arc::clone(&source.data_ready),
572                async move {
573                    let _ = started_tx.send(());
574                    let _ = release_rx.await;
575                    Ok(())
576                },
577            )
578            .unwrap(),
579        );
580        let (shutdown_tx, shutdown_rx) = watch::channel(false);
581        source.shutdown_tx = Some(shutdown_tx);
582        source.state = ConnectorState::Running;
583        started_rx.await.expect("test server task started");
584
585        let close = tokio::spawn(async move { source.close().await });
586        tokio::task::yield_now().await;
587        assert!(!close.is_finished(), "close must be joining the server");
588        close.abort();
589        assert!(
590            close
591                .await
592                .expect_err("close waiter cancelled")
593                .is_cancelled(),
594            "close waiter must be cancelled"
595        );
596
597        assert!(*shutdown_rx.borrow(), "close must publish shutdown");
598        tokio::time::timeout(Duration::from_secs(2), terminal.wait_terminated())
599            .await
600            .expect("cancelled close left the server generation live");
601        assert!(
602            release_tx.send(()).is_err(),
603            "the aborted server must drop its test receiver"
604        );
605    }
606
607    #[tokio::test(start_paused = true)]
608    async fn close_timeout_aborts_the_server_task() {
609        let mut source = OtelSource::new(traces_schema(), None);
610        let terminal = source.terminal_task_tracker().unwrap();
611        source.server_task = Some(
612            TrackedServerTask::spawn(
613                &source.task_owner,
614                Arc::clone(&source.data_ready),
615                std::future::pending::<Result<(), ConnectorError>>(),
616            )
617            .unwrap(),
618        );
619        let (shutdown_tx, shutdown_rx) = watch::channel(false);
620        source.shutdown_tx = Some(shutdown_tx);
621        source.state = ConnectorState::Running;
622
623        let error = source.close().await.unwrap_err();
624
625        assert!(*shutdown_rx.borrow(), "close must publish shutdown");
626        assert!(error.to_string().contains("close deadline"), "{error}");
627        assert_eq!(source.state, ConnectorState::Failed);
628        assert!(
629            source.server_task.is_none(),
630            "the aborted task must not block restart"
631        );
632        assert!(matches!(
633            source.start(start_request()).await,
634            Err(ConnectorError::InvalidState { .. })
635        ));
636        drop(source);
637        tokio::time::timeout(Duration::from_secs(1), terminal.wait_terminated())
638            .await
639            .expect("aborted OTel server generation did not terminate");
640    }
641
642    #[tokio::test]
643    async fn late_server_completion_is_a_close_timeout() {
644        let owner = ConnectorTaskOwner::new().0;
645        let mut task = TrackedServerTask::spawn(&owner, Arc::new(Notify::new()), async {
646            std::thread::sleep(Duration::from_millis(25));
647            Ok(())
648        })
649        .unwrap();
650
651        let result = task
652            .wait_until(tokio::time::Instant::now() + Duration::from_millis(5))
653            .await;
654
655        assert!(matches!(result, ServerWait::TimedOut));
656        assert!(task.handle.is_none());
657    }
658
659    #[tokio::test]
660    async fn unexpected_server_exit_fails_live_polling() {
661        let mut source = OtelSource::new(traces_schema(), None);
662        let (_batch_tx, batch_rx) = mpsc::bounded_async::<RecordBatch>(1);
663        source.batch_rx = Some(batch_rx);
664        source.server_task = Some(
665            TrackedServerTask::spawn(&source.task_owner, Arc::clone(&source.data_ready), async {
666                Err(ConnectorError::ConnectionFailed(
667                    "accept loop failed".into(),
668                ))
669            })
670            .unwrap(),
671        );
672        source.state = ConnectorState::Running;
673        tokio::time::timeout(Duration::from_secs(1), async {
674            while !source
675                .server_task
676                .as_ref()
677                .and_then(|task| task.handle.as_ref())
678                .is_some_and(tokio::task::JoinHandle::is_finished)
679            {
680                tokio::task::yield_now().await;
681            }
682        })
683        .await
684        .expect("test server failure did not become observable");
685        tokio::time::timeout(Duration::from_secs(1), source.data_ready.notified())
686            .await
687            .expect("terminal server exit did not wake source polling");
688
689        let error = source.poll_batch(1).await.unwrap_err();
690
691        assert!(error.to_string().contains("accept loop failed"), "{error}");
692        assert!(
693            !error.is_transient(),
694            "terminal generation errors must not retry"
695        );
696        assert_eq!(source.state, ConnectorState::Failed);
697        assert!(source.server_task.is_none());
698    }
699}
laminar_connectors/otel/source.rs

laminar_connectors/otel/
source.rs
