#![allow(clippy::module_name_repetitions)]
use super::{panic_message, SourceMessage};
use crate::error::{Error, Result};
use crate::pipeline::{wait_if_paused, ErrorPolicy, SampleChunk};
use crate::source::Source;
use crossbeam_channel::bounded;
use crossbeam_channel::{Receiver, RecvTimeoutError, Sender};
use std::panic::{catch_unwind, AssertUnwindSafe};
use std::sync::atomic::{AtomicBool, AtomicU64, Ordering};
use std::sync::Arc;
use std::thread::JoinHandle;
use std::time::Duration;
pub(crate) fn recv_source_message(
source_rx: &Receiver<SourceMessage>,
source_timeout: Option<Duration>,
shutdown_flag: &std::sync::atomic::AtomicBool,
) -> Result<SourceMessage> {
let start = std::time::Instant::now();
loop {
if shutdown_flag.load(std::sync::atomic::Ordering::Relaxed) {
return Err(Error::SourceFailed {
source_name: "pipeline".to_string(),
reason: "pipeline shutdown cleanly requested".to_string(),
});
}
let remaining = source_timeout.map_or(Duration::from_secs(1), |limit| {
limit
.saturating_sub(start.elapsed())
.max(Duration::from_millis(1))
});
match source_rx.recv_timeout(remaining) {
Ok(msg) => return Ok(msg),
Err(RecvTimeoutError::Timeout) => {
if let Some(limit) = source_timeout {
if start.elapsed() >= limit {
return Err(Error::SourceFailed {
source_name: "pipeline".to_string(),
reason: format!("source read timed out after {limit:?}"),
});
}
}
}
Err(RecvTimeoutError::Disconnected) => {
return Err(Error::SourceFailed {
source_name: "pipeline".to_string(),
reason: "source reader disconnected unexpectedly".to_string(),
});
}
}
}
}
pub(crate) fn run_source_epoch(source: Arc<dyn Source>, source_tx: Sender<SourceMessage>) {
let source_name = source.name().to_string();
let fatal = match catch_unwind(AssertUnwindSafe(|| source.open())) {
Ok(Ok(iterator)) => run_source_iterator(iterator, &source_name, &source_tx),
Ok(Err(error)) => Some(error),
Err(payload) => Some(Error::SourceFailed {
source_name: source_name.clone(),
reason: format!("open panicked: {}", panic_message(payload)),
}),
};
if let Some(error) = fatal {
let _ = source_tx.send(SourceMessage::Fatal(error));
} else {
let _ = source_tx.send(SourceMessage::EndOfEpoch);
}
}
pub(crate) fn run_source_iterator(
mut iterator: Box<dyn crate::source::SourceIterator>,
source_name: &str,
source_tx: &Sender<SourceMessage>,
) -> Option<Error> {
loop {
match catch_unwind(AssertUnwindSafe(|| iterator.next_sample())) {
Ok(Some(Ok(sample))) => {
if source_tx.send(SourceMessage::Sample(sample)).is_err() {
return None;
}
}
Ok(Some(Err(error))) => {
if source_tx.send(SourceMessage::ItemError(error)).is_err() {
return None;
}
}
Ok(None) => return None,
Err(payload) => {
return Some(Error::SourceFailed {
source_name: source_name.to_string(),
reason: format!("next_sample panicked: {}", panic_message(payload)),
});
}
}
}
}
#[allow(clippy::too_many_arguments, clippy::too_many_lines)]
pub(crate) fn spawn_source_thread(
source: Arc<dyn Source>,
raw_tx: Sender<SampleChunk>,
chunk_size: usize,
epochs: usize,
source_timeout: Option<std::time::Duration>,
error_policy: ErrorPolicy,
test_start_sequence: u64,
output_paused: Arc<AtomicBool>,
shutdown: Arc<AtomicBool>,
errors_skipped: Arc<AtomicU64>,
fatal_error: Arc<std::sync::Mutex<Option<crate::error::Error>>>,
) -> Result<JoinHandle<()>> {
std::thread::Builder::new()
.name("tenshift-source".into())
.spawn(move || {
let epoch_count = if epochs == 0 { usize::MAX } else { epochs };
let mut sequence = test_start_sequence;
let mut last_error_log = std::time::Instant::now();
let mut suppressed_errors = 0_u64;
for _ in 0..epoch_count {
if shutdown.load(Ordering::Relaxed) {
break;
}
let (source_tx, source_rx) = bounded::<SourceMessage>(chunk_size.max(1));
let source_impl = Arc::clone(&source);
if std::thread::Builder::new()
.name("tenshift-source-epoch".into())
.spawn(move || run_source_epoch(source_impl, source_tx))
.is_err()
{
tracing::error!("source init failed: could not spawn source epoch thread");
shutdown.store(true, Ordering::Relaxed);
return;
}
let mut chunk_buf = Vec::with_capacity(chunk_size);
while !shutdown.load(Ordering::Relaxed) {
if wait_if_paused(output_paused.as_ref(), shutdown.as_ref()) {
break;
}
let message = match recv_source_message(&source_rx, source_timeout, &shutdown) {
Ok(message) => message,
Err(error) => {
tracing::error!("source failed: {error}");
if let Ok(mut lock) = fatal_error.lock() {
if lock.is_none() {
*lock = Some(error);
}
}
shutdown.store(true, Ordering::Relaxed);
return;
}
};
match message {
SourceMessage::Sample(sample) => {
chunk_buf.push(sample);
if chunk_buf.len() >= chunk_size {
let samples = std::mem::replace(
&mut chunk_buf,
Vec::with_capacity(chunk_size),
);
if wait_if_paused(output_paused.as_ref(), shutdown.as_ref()) {
break;
}
if raw_tx.send(SampleChunk { sequence, samples }).is_err() {
return;
}
if sequence == u64::MAX {
tracing::error!(
"sequence number overflow: stopping source after {} chunks",
sequence
);
shutdown.store(true, Ordering::Relaxed);
return;
}
sequence += 1;
}
}
SourceMessage::ItemError(error) => match error_policy {
ErrorPolicy::Skip => {
errors_skipped.fetch_add(1, Ordering::Relaxed);
if last_error_log.elapsed() >= std::time::Duration::from_secs(1) {
if suppressed_errors > 0 {
tracing::warn!("skipping bad sample from source: {error} (and {} more suppressed)", suppressed_errors);
} else {
tracing::warn!("skipping bad sample from source: {error}");
}
suppressed_errors = 0;
last_error_log = std::time::Instant::now();
} else {
suppressed_errors += 1;
}
}
ErrorPolicy::Fail => {
tracing::error!("source error: {error}");
if let Ok(mut lock) = fatal_error.lock() {
if lock.is_none() {
*lock = Some(error);
}
}
shutdown.store(true, Ordering::Relaxed);
return;
}
},
SourceMessage::Fatal(error) => {
tracing::error!("source failed: {error}");
if let Ok(mut lock) = fatal_error.lock() {
if lock.is_none() {
*lock = Some(error);
}
}
shutdown.store(true, Ordering::Relaxed);
return;
}
SourceMessage::EndOfEpoch => break,
}
}
if !chunk_buf.is_empty() && !shutdown.load(Ordering::Relaxed) {
if wait_if_paused(output_paused.as_ref(), shutdown.as_ref()) {
break;
}
if raw_tx
.send(SampleChunk {
sequence,
samples: chunk_buf,
})
.is_err()
{
return;
}
if sequence == u64::MAX {
tracing::error!(
"sequence number overflow: stopping source after {} chunks",
sequence
);
shutdown.store(true, Ordering::Relaxed);
return;
}
sequence += 1;
}
}
})
.map_err(Into::into)
}