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 resolve_schema(
179        &mut self,
180        config: &ConnectorConfig,
181        explicit: Option<SchemaRef>,
182    ) -> Result<crate::schema::resolution::SchemaBinding, ConnectorError> {
183        crate::schema::resolution::fixed_binding(config, explicit, &Self::generator_schema())
184    }
185
186    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
187        let (config, position, _) = request.into_parts();
188        let binding = self.resolve_schema(&config, config.arrow_schema()).await?;
189        self.schema = Arc::new(binding.logical);
190        let source_name = config
191            .get("laminar.source.name")
192            .filter(|name| !name.is_empty())
193            .unwrap_or("generator");
194        self.input_channel = generator_input_channel(source_name)?;
195        if let Some(rps) = config.get_parsed::<u64>("rows.per.second")? {
196            if rps == 0 {
197                return Err(ConnectorError::ConfigurationError(
198                    "rows.per.second must be > 0".into(),
199                ));
200            }
201            self.rows_per_second = rps;
202        }
203        if let Some(n) = config.get_parsed::<usize>("batch.max.size")? {
204            self.batch_max = n.max(1);
205        }
206        self.max_rows = config.get_parsed::<u64>("max.rows")?;
207
208        self.next_seq = match position {
209            SourcePosition::Initial => 0,
210            SourcePosition::Initialized { .. } => {
211                return Err(ConnectorError::ConfigurationError(
212                    "generator has no sealed topology startup contract".into(),
213                ));
214            }
215            SourcePosition::Resume {
216                attempt,
217                checkpoint,
218            } => {
219                validate_generator_checkpoint(&checkpoint)?;
220                if checkpoint.input_channels() != Some(std::slice::from_ref(&self.input_channel)) {
221                    return Err(ConnectorError::ConfigurationError(format!(
222                        "generator checkpoint {attempt:?} has the wrong input-channel inventory"
223                    )));
224                }
225                let seq = checkpoint.get_offset("seq").ok_or_else(|| {
226                    ConnectorError::ConfigurationError(format!(
227                        "generator checkpoint {attempt:?} is missing required 'seq' offset"
228                    ))
229                })?;
230                seq.parse().map_err(|e| {
231                    ConnectorError::ConfigurationError(format!(
232                        "bad generator offset '{seq}' in checkpoint {attempt:?}: {e}"
233                    ))
234                })?
235            }
236        };
237        // Position must be installed before the rate anchor is created. Otherwise a
238        // resumed source can accrue tokens against the pre-resume sequence and burst.
239        self.anchor = None;
240        Ok(())
241    }
242
243    // Cast lints: see build_batch.
244    #[allow(
245        clippy::cast_possible_truncation,
246        clippy::cast_sign_loss,
247        clippy::cast_precision_loss
248    )]
249    async fn poll_batch(
250        &mut self,
251        max_records: usize,
252    ) -> Result<Option<SourceBatch>, ConnectorError> {
253        let (anchored_at, anchor_seq) = *self
254            .anchor
255            .get_or_insert_with(|| (Instant::now(), self.next_seq));
256        let mut allowed = anchor_seq.saturating_add(
257            (anchored_at.elapsed().as_secs_f64() * self.rows_per_second as f64) as u64,
258        );
259        if let Some(max) = self.max_rows {
260            allowed = allowed.min(max);
261        }
262        let pending = allowed.saturating_sub(self.next_seq);
263        let n = (pending as usize).min(max_records).min(self.batch_max);
264        if n == 0 {
265            return Ok(None);
266        }
267        let batch = self.build_batch(self.next_seq, n)?;
268        let row_positions = self.build_row_positions(self.next_seq, n)?;
269        self.next_seq =
270            self.next_seq
271                .checked_add(u64::try_from(n).map_err(|_| {
272                    ConnectorError::Internal("generator batch size exceeds u64".into())
273                })?)
274                .ok_or_else(|| ConnectorError::Internal("generator sequence overflow".into()))?;
275        Ok(Some(SourceBatch::positioned(batch, row_positions)?))
276    }
277
278    fn schema(&self) -> SchemaRef {
279        Arc::clone(&self.schema)
280    }
281
282    fn checkpoint(&self) -> SourceCheckpoint {
283        let mut cp = SourceCheckpoint::new();
284        cp.set_offset("seq", self.next_seq.to_string());
285        cp.set_metadata("connector", GENERATOR_CHECKPOINT_CONNECTOR);
286        cp.set_metadata(CHECKPOINT_VERSION_METADATA, GENERATOR_CHECKPOINT_VERSION);
287        cp.set_input_channels(vec![self.input_channel.clone()])
288            .expect("the cached generator input-channel identity is valid");
289        cp
290    }
291
292    async fn close(&mut self) -> Result<(), ConnectorError> {
293        Ok(())
294    }
295}
296
297/// Registers the generator source so `CREATE SOURCE ... FROM GENERATOR (...)` resolves.
298///
299/// # Errors
300///
301/// Returns the registry error when the name is already registered or the registry is frozen.
302pub fn register_generator_source(registry: &ConnectorRegistry) -> Result<(), ConnectorError> {
303    let info = ConnectorInfo {
304        schema_capabilities: crate::schema::resolution::SchemaCapabilities::built_in(),
305        name: "generator".to_string(),
306        display_name: "Synthetic Data Generator".to_string(),
307        version: env!("CARGO_PKG_VERSION").to_string(),
308        is_source: true,
309        is_sink: false,
310        config_keys: vec![
311            ConfigKeySpec::optional("rows.per.second", "Emission rate", "1000"),
312            ConfigKeySpec::optional("batch.max.size", "Max rows per batch", "1024"),
313            ConfigKeySpec::optional(
314                "max.rows",
315                "Stop after this many rows (unbounded if unset)",
316                "",
317            ),
318        ],
319    };
320    registry.register_source(
321        "generator",
322        info,
323        Arc::new(|_registry: Option<&Arc<prometheus::Registry>>| {
324            Ok(Box::new(GeneratorSource::default()))
325        }),
326    )
327}
328
329#[cfg(test)]
330mod tests;
laminar_connectors/generator/mod.rs

laminar_connectors/generator/
mod.rs
