laminar_connectors/generator/
mod.rs

1//! Synthetic data generator source — no external infrastructure.
2//!
3//! Emits a deterministic sequence at a configured rate: row `i` is
4//! always `(seq = i, ts_ms = i * 1000 / rows_per_second, value = "v{i}")`,
5//! so output is a pure function of the offset. That makes the source
6//! fully replayable (exactly-once capable) and lets harnesses verify
7//! sink completeness by recomputing the expected rows — its primary
8//! consumer is the cluster soak test, but it works anywhere a
9//! self-driving source is needed (demos, benchmarks).
10
11use std::sync::Arc;
12use std::time::Instant;
13
14use arrow_array::builder::BinaryBuilder;
15use arrow_array::{Int64Array, RecordBatch, StringArray, UInt32Array};
16use arrow_schema::{DataType, Field, Schema, SchemaRef};
17use async_trait::async_trait;
18
19use crate::checkpoint::SourceCheckpoint;
20use crate::config::{ConfigKeySpec, ConnectorConfig, ConnectorInfo};
21use crate::connector::{
22    SourceBatch, SourceConnector, SourceConsistency, SourceContract, SourceInputMode,
23    SourcePosition, SourceRowPositionCapability, SourceRowPositions, SourceStart, SourceTopology,
24};
25use crate::error::ConnectorError;
26use crate::registry::ConnectorRegistry;
27
28const GENERATOR_CHECKPOINT_CONNECTOR: &str = "generator";
29const CHECKPOINT_VERSION_METADATA: &str = "checkpoint.version";
30const GENERATOR_CHECKPOINT_VERSION: &str = "2";
31
32fn generator_input_channel(source_name: &str) -> Result<Vec<u8>, ConnectorError> {
33    let source_len = u32::try_from(source_name.len()).map_err(|_| {
34        ConnectorError::ConfigurationError(
35            "generator source identity exceeds the input-channel encoding limit".into(),
36        )
37    })?;
38    let mut channel = Vec::with_capacity(source_name.len() + 8);
39    channel.extend_from_slice(&source_len.to_be_bytes());
40    channel.extend_from_slice(source_name.as_bytes());
41    channel.extend_from_slice(&0_u32.to_be_bytes());
42    Ok(channel)
43}
44
45fn validate_generator_checkpoint(checkpoint: &SourceCheckpoint) -> Result<(), ConnectorError> {
46    match checkpoint.get_metadata("connector") {
47        Some(GENERATOR_CHECKPOINT_CONNECTOR) => {}
48        Some(connector) => {
49            return Err(ConnectorError::ConfigurationError(format!(
50                "generator checkpoint belongs to connector '{connector}'"
51            )));
52        }
53        None => {
54            return Err(ConnectorError::ConfigurationError(
55                "generator checkpoint is missing connector identity".into(),
56            ));
57        }
58    }
59    if checkpoint.get_metadata(CHECKPOINT_VERSION_METADATA) != Some(GENERATOR_CHECKPOINT_VERSION) {
60        return Err(ConnectorError::ConfigurationError(format!(
61            "generator checkpoint requires {CHECKPOINT_VERSION_METADATA}={GENERATOR_CHECKPOINT_VERSION}"
62        )));
63    }
64    Ok(())
65}
66
67/// Deterministic rate-limited source. See module docs.
68pub struct GeneratorSource {
69    schema: SchemaRef,
70    rows_per_second: u64,
71    batch_max: usize,
72    max_rows: Option<u64>,
73    input_channel: Vec<u8>,
74    /// Next sequence number to emit (== rows emitted so far across
75    /// restarts; restored from the checkpoint).
76    next_seq: u64,
77    /// Rate-limit anchor: emission is allowed up to
78    /// `anchor_seq + elapsed_since(anchor) * rows_per_second`.
79    /// Created only after `start` installs the initial or recovered sequence,
80    /// so a restart cannot burst or stall against a stale anchor.
81    anchor: Option<(Instant, u64)>,
82}
83
84impl GeneratorSource {
85    fn generator_schema() -> SchemaRef {
86        Arc::new(Schema::new(vec![
87            Field::new("seq", DataType::Int64, false),
88            Field::new("ts_ms", DataType::Int64, false),
89            Field::new("value", DataType::Utf8, false),
90        ]))
91    }
92
93    /// Build rows `[start, start + n)` — a pure function of `start`,
94    /// which is what makes the source replayable.
95    // Cast lints: seq is a monotonic generator counter and rates are
96    // operator-supplied config — both far below the 2^52/2^63 edges.
97    #[allow(
98        clippy::cast_possible_wrap,
99        clippy::cast_sign_loss,
100        clippy::cast_possible_truncation,
101        clippy::cast_precision_loss
102    )]
103    fn build_batch(&self, start: u64, n: usize) -> Result<RecordBatch, ConnectorError> {
104        let seqs: Vec<i64> = (0..n as u64).map(|i| (start + i) as i64).collect();
105        let ts: Vec<i64> = seqs
106            .iter()
107            .map(|&s| {
108                (s as u64)
109                    .saturating_mul(1000)
110                    .wrapping_div(self.rows_per_second) as i64
111            })
112            .collect();
113        let values: Vec<String> = seqs.iter().map(|s| format!("v{s}")).collect();
114        RecordBatch::try_new(
115            Arc::clone(&self.schema),
116            vec![
117                Arc::new(Int64Array::from(seqs)),
118                Arc::new(Int64Array::from(ts)),
119                Arc::new(StringArray::from(values)),
120            ],
121        )
122        .map_err(|e| ConnectorError::ReadError(e.to_string()))
123    }
124
125    fn build_row_positions(
126        &self,
127        start: u64,
128        n: usize,
129    ) -> Result<SourceRowPositions, ConnectorError> {
130        let mut partitions =
131            BinaryBuilder::with_capacity(n, self.input_channel.len().saturating_mul(n));
132        let mut order_keys = BinaryBuilder::with_capacity(n, 8_usize.saturating_mul(n));
133        for row in 0..n {
134            let sequence = start
135                .checked_add(u64::try_from(row).map_err(|_| {
136                    ConnectorError::Internal("generator batch row index exceeds u64".into())
137                })?)
138                .ok_or_else(|| ConnectorError::Internal("generator sequence overflow".into()))?;
139            partitions.append_value(&self.input_channel);
140            order_keys.append_value(sequence.to_be_bytes());
141        }
142        SourceRowPositions::try_new(
143            partitions.finish(),
144            order_keys.finish(),
145            UInt32Array::from(vec![0; n]),
146        )
147    }
148}
149
150impl Default for GeneratorSource {
151    fn default() -> Self {
152        let input_channel = generator_input_channel("generator")
153            .expect("the built-in generator source identity is valid");
154        Self {
155            schema: Self::generator_schema(),
156            rows_per_second: 1000,
157            batch_max: 1024,
158            max_rows: None,
159            input_channel,
160            next_seq: 0,
161            anchor: None,
162        }
163    }
164}
165
166#[async_trait]
167impl SourceConnector for GeneratorSource {
168    fn contract(&self, _config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
169        Ok(SourceContract::new(
170            SourceConsistency::Replayable,
171            SourceTopology::Singleton,
172            SourceInputMode::AppendOnly,
173        )
174        .with_row_positions(SourceRowPositionCapability::OrderedDeterministic)
175        .with_exact_delivery_certification())
176    }
177
178    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
179        let (config, position, _) = request.into_parts();
180        let source_name = config
181            .get("laminar.source.name")
182            .filter(|name| !name.is_empty())
183            .unwrap_or("generator");
184        self.input_channel = generator_input_channel(source_name)?;
185        if let Some(rps) = config.get_parsed::<u64>("rows.per.second")? {
186            if rps == 0 {
187                return Err(ConnectorError::ConfigurationError(
188                    "rows.per.second must be > 0".into(),
189                ));
190            }
191            self.rows_per_second = rps;
192        }
193        if let Some(n) = config.get_parsed::<usize>("batch.max.size")? {
194            self.batch_max = n.max(1);
195        }
196        self.max_rows = config.get_parsed::<u64>("max.rows")?;
197
198        self.next_seq = match position {
199            SourcePosition::Initial => 0,
200            SourcePosition::Resume {
201                attempt,
202                checkpoint,
203            } => {
204                validate_generator_checkpoint(&checkpoint)?;
205                if checkpoint.input_channels() != Some(std::slice::from_ref(&self.input_channel)) {
206                    return Err(ConnectorError::ConfigurationError(format!(
207                        "generator checkpoint {attempt:?} has the wrong input-channel inventory"
208                    )));
209                }
210                let seq = checkpoint.get_offset("seq").ok_or_else(|| {
211                    ConnectorError::ConfigurationError(format!(
212                        "generator checkpoint {attempt:?} is missing required 'seq' offset"
213                    ))
214                })?;
215                seq.parse().map_err(|e| {
216                    ConnectorError::ConfigurationError(format!(
217                        "bad generator offset '{seq}' in checkpoint {attempt:?}: {e}"
218                    ))
219                })?
220            }
221        };
222        // Position must be installed before the rate anchor is created. Otherwise a
223        // resumed source can accrue tokens against the pre-resume sequence and burst.
224        self.anchor = None;
225        Ok(())
226    }
227
228    // Cast lints: see build_batch.
229    #[allow(
230        clippy::cast_possible_truncation,
231        clippy::cast_sign_loss,
232        clippy::cast_precision_loss
233    )]
234    async fn poll_batch(
235        &mut self,
236        max_records: usize,
237    ) -> Result<Option<SourceBatch>, ConnectorError> {
238        let (anchored_at, anchor_seq) = *self
239            .anchor
240            .get_or_insert_with(|| (Instant::now(), self.next_seq));
241        let mut allowed = anchor_seq.saturating_add(
242            (anchored_at.elapsed().as_secs_f64() * self.rows_per_second as f64) as u64,
243        );
244        if let Some(max) = self.max_rows {
245            allowed = allowed.min(max);
246        }
247        let pending = allowed.saturating_sub(self.next_seq);
248        let n = (pending as usize).min(max_records).min(self.batch_max);
249        if n == 0 {
250            return Ok(None);
251        }
252        let batch = self.build_batch(self.next_seq, n)?;
253        let row_positions = self.build_row_positions(self.next_seq, n)?;
254        self.next_seq =
255            self.next_seq
256                .checked_add(u64::try_from(n).map_err(|_| {
257                    ConnectorError::Internal("generator batch size exceeds u64".into())
258                })?)
259                .ok_or_else(|| ConnectorError::Internal("generator sequence overflow".into()))?;
260        Ok(Some(SourceBatch::positioned(batch, row_positions)?))
261    }
262
263    fn schema(&self) -> SchemaRef {
264        Arc::clone(&self.schema)
265    }
266
267    fn checkpoint(&self) -> SourceCheckpoint {
268        let mut cp = SourceCheckpoint::new();
269        cp.set_offset("seq", self.next_seq.to_string());
270        cp.set_metadata("connector", GENERATOR_CHECKPOINT_CONNECTOR);
271        cp.set_metadata(CHECKPOINT_VERSION_METADATA, GENERATOR_CHECKPOINT_VERSION);
272        cp.set_input_channels(vec![self.input_channel.clone()])
273            .expect("the cached generator input-channel identity is valid");
274        cp
275    }
276
277    async fn close(&mut self) -> Result<(), ConnectorError> {
278        Ok(())
279    }
280}
281
282/// Registers the generator source so `CREATE SOURCE ... FROM GENERATOR (...)` resolves.
283///
284/// # Errors
285///
286/// Returns the registry error when the name is already registered or the registry is frozen.
287pub fn register_generator_source(registry: &ConnectorRegistry) -> Result<(), ConnectorError> {
288    let info = ConnectorInfo {
289        name: "generator".to_string(),
290        display_name: "Synthetic Data Generator".to_string(),
291        version: env!("CARGO_PKG_VERSION").to_string(),
292        is_source: true,
293        is_sink: false,
294        config_keys: vec![
295            ConfigKeySpec::optional("rows.per.second", "Emission rate", "1000"),
296            ConfigKeySpec::optional("batch.max.size", "Max rows per batch", "1024"),
297            ConfigKeySpec::optional(
298                "max.rows",
299                "Stop after this many rows (unbounded if unset)",
300                "",
301            ),
302        ],
303    };
304    registry.register_source(
305        "generator",
306        info,
307        Arc::new(|_registry: Option<&Arc<prometheus::Registry>>| {
308            Ok(Box::new(GeneratorSource::default()))
309        }),
310    )
311}
312
313#[cfg(test)]
314mod tests;
laminar_connectors/generator/mod.rs

laminar_connectors/generator/
mod.rs
