laminar_connectors/
generator.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::{Int64Array, RecordBatch, StringArray};
15use arrow_schema::{DataType, Field, Schema, SchemaRef};
16use async_trait::async_trait;
17
18use crate::checkpoint::SourceCheckpoint;
19use crate::config::{ConfigKeySpec, ConnectorConfig, ConnectorInfo};
20use crate::connector::{
21    SourceBatch, SourceConnector, SourceConsistency, SourceContract, SourcePosition, SourceStart,
22    SourceTopology,
23};
24use crate::error::ConnectorError;
25use crate::registry::ConnectorRegistry;
26
27/// Deterministic rate-limited source. See module docs.
28pub struct GeneratorSource {
29    schema: SchemaRef,
30    rows_per_second: u64,
31    batch_max: usize,
32    max_rows: Option<u64>,
33    /// Next sequence number to emit (== rows emitted so far across
34    /// restarts; restored from the checkpoint).
35    next_seq: u64,
36    /// Rate-limit anchor: emission is allowed up to
37    /// `anchor_seq + elapsed_since(anchor) * rows_per_second`.
38    /// Created only after `start` installs the initial or recovered sequence,
39    /// so a restart cannot burst or stall against a stale anchor.
40    anchor: Option<(Instant, u64)>,
41}
42
43impl GeneratorSource {
44    fn generator_schema() -> SchemaRef {
45        Arc::new(Schema::new(vec![
46            Field::new("seq", DataType::Int64, false),
47            Field::new("ts_ms", DataType::Int64, false),
48            Field::new("value", DataType::Utf8, false),
49        ]))
50    }
51
52    /// Build rows `[start, start + n)` — a pure function of `start`,
53    /// which is what makes the source replayable.
54    // Cast lints: seq is a monotonic generator counter and rates are
55    // operator-supplied config — both far below the 2^52/2^63 edges.
56    #[allow(
57        clippy::cast_possible_wrap,
58        clippy::cast_sign_loss,
59        clippy::cast_possible_truncation,
60        clippy::cast_precision_loss
61    )]
62    fn build_batch(&self, start: u64, n: usize) -> Result<RecordBatch, ConnectorError> {
63        let seqs: Vec<i64> = (0..n as u64).map(|i| (start + i) as i64).collect();
64        let ts: Vec<i64> = seqs
65            .iter()
66            .map(|&s| {
67                (s as u64)
68                    .saturating_mul(1000)
69                    .wrapping_div(self.rows_per_second) as i64
70            })
71            .collect();
72        let values: Vec<String> = seqs.iter().map(|s| format!("v{s}")).collect();
73        RecordBatch::try_new(
74            Arc::clone(&self.schema),
75            vec![
76                Arc::new(Int64Array::from(seqs)),
77                Arc::new(Int64Array::from(ts)),
78                Arc::new(StringArray::from(values)),
79            ],
80        )
81        .map_err(|e| ConnectorError::ReadError(e.to_string()))
82    }
83}
84
85impl Default for GeneratorSource {
86    fn default() -> Self {
87        Self {
88            schema: Self::generator_schema(),
89            rows_per_second: 1000,
90            batch_max: 1024,
91            max_rows: None,
92            next_seq: 0,
93            anchor: None,
94        }
95    }
96}
97
98#[async_trait]
99impl SourceConnector for GeneratorSource {
100    fn contract(&self, _config: &ConnectorConfig) -> Result<SourceContract, ConnectorError> {
101        Ok(
102            SourceContract::new(SourceConsistency::Replayable, SourceTopology::Singleton)
103                .with_exact_delivery_certification(),
104        )
105    }
106
107    async fn start(&mut self, request: SourceStart) -> Result<(), ConnectorError> {
108        let (config, position, _) = request.into_parts();
109        if let Some(rps) = config.get_parsed::<u64>("rows.per.second")? {
110            if rps == 0 {
111                return Err(ConnectorError::ConfigurationError(
112                    "rows.per.second must be > 0".into(),
113                ));
114            }
115            self.rows_per_second = rps;
116        }
117        if let Some(n) = config.get_parsed::<usize>("batch.max.size")? {
118            self.batch_max = n.max(1);
119        }
120        self.max_rows = config.get_parsed::<u64>("max.rows")?;
121
122        self.next_seq = match position {
123            SourcePosition::Initial => 0,
124            SourcePosition::Resume {
125                attempt,
126                checkpoint,
127            } => {
128                let seq = checkpoint.get_offset("seq").ok_or_else(|| {
129                    ConnectorError::ConfigurationError(format!(
130                        "generator checkpoint {attempt:?} is missing required 'seq' offset"
131                    ))
132                })?;
133                seq.parse().map_err(|e| {
134                    ConnectorError::ConfigurationError(format!(
135                        "bad generator offset '{seq}' in checkpoint {attempt:?}: {e}"
136                    ))
137                })?
138            }
139        };
140        // Position must be installed before the rate anchor is created. Otherwise a
141        // resumed source can accrue tokens against the pre-resume sequence and burst.
142        self.anchor = None;
143        Ok(())
144    }
145
146    // Cast lints: see build_batch.
147    #[allow(
148        clippy::cast_possible_truncation,
149        clippy::cast_sign_loss,
150        clippy::cast_precision_loss
151    )]
152    async fn poll_batch(
153        &mut self,
154        max_records: usize,
155    ) -> Result<Option<SourceBatch>, ConnectorError> {
156        let (anchored_at, anchor_seq) = *self
157            .anchor
158            .get_or_insert_with(|| (Instant::now(), self.next_seq));
159        let mut allowed = anchor_seq.saturating_add(
160            (anchored_at.elapsed().as_secs_f64() * self.rows_per_second as f64) as u64,
161        );
162        if let Some(max) = self.max_rows {
163            allowed = allowed.min(max);
164        }
165        let pending = allowed.saturating_sub(self.next_seq);
166        let n = (pending as usize).min(max_records).min(self.batch_max);
167        if n == 0 {
168            return Ok(None);
169        }
170        let batch = self.build_batch(self.next_seq, n)?;
171        self.next_seq += n as u64;
172        Ok(Some(SourceBatch::new(batch)))
173    }
174
175    fn schema(&self) -> SchemaRef {
176        Arc::clone(&self.schema)
177    }
178
179    fn checkpoint(&self) -> SourceCheckpoint {
180        let mut cp = SourceCheckpoint::new();
181        cp.set_offset("seq", self.next_seq.to_string());
182        cp
183    }
184
185    async fn close(&mut self) -> Result<(), ConnectorError> {
186        Ok(())
187    }
188}
189
190/// Registers the generator source so
191/// `CREATE SOURCE ... WITH (connector = 'generator')` resolves.
192///
193/// # Errors
194///
195/// Returns the registry error when the name is already registered or the registry is frozen.
196pub fn register_generator_source(registry: &ConnectorRegistry) -> Result<(), ConnectorError> {
197    let info = ConnectorInfo {
198        name: "generator".to_string(),
199        display_name: "Synthetic Data Generator".to_string(),
200        version: env!("CARGO_PKG_VERSION").to_string(),
201        is_source: true,
202        is_sink: false,
203        config_keys: vec![
204            ConfigKeySpec::optional("rows.per.second", "Emission rate", "1000"),
205            ConfigKeySpec::optional("batch.max.size", "Max rows per batch", "1024"),
206            ConfigKeySpec::optional(
207                "max.rows",
208                "Stop after this many rows (unbounded if unset)",
209                "",
210            ),
211        ],
212    };
213    registry.register_source(
214        "generator",
215        info,
216        Arc::new(|_registry: Option<&Arc<prometheus::Registry>>| {
217            Ok(Box::new(GeneratorSource::default()))
218        }),
219    )
220}
221
222#[cfg(test)]
223mod tests {
224    use super::*;
225    use crate::connector::{DeliveryGuarantee, SourcePosition, SourceStart};
226    use laminar_core::state::CheckpointAttempt;
227
228    fn start_request(config: ConnectorConfig, position: SourcePosition) -> SourceStart {
229        SourceStart::new(config, position, DeliveryGuarantee::AtLeastOnce).unwrap()
230    }
231
232    #[test]
233    fn source_contract_is_replayable_singleton() {
234        let source = GeneratorSource::default();
235        let contract = source
236            .contract(&ConnectorConfig::new("generator"))
237            .expect("static generator contract");
238        assert_eq!(contract.consistency, SourceConsistency::Replayable);
239        assert_eq!(contract.topology, SourceTopology::Singleton);
240        assert!(contract.is_exact_delivery_certified());
241    }
242
243    #[tokio::test]
244    async fn deterministic_and_replayable_across_resume() {
245        let mut config = ConnectorConfig::new("generator");
246        config.set("rows.per.second", "1000000"); // effectively unthrottled
247        let mut a = GeneratorSource::default();
248        a.start(start_request(config.clone(), SourcePosition::Initial))
249            .await
250            .unwrap();
251        // First poll anchors the rate limit (no tokens accrue before
252        // polling starts); rows become available after it.
253        let _ = a.poll_batch(8).await.unwrap();
254        std::thread::sleep(std::time::Duration::from_millis(10));
255        let first = a.poll_batch(8).await.unwrap().expect("rows");
256        assert_eq!(first.num_rows(), 8);
257
258        // Restore a fresh instance from a's checkpoint at seq=8 and
259        // verify the next rows equal what `a` produces next.
260        let cp = a.checkpoint();
261        let mut b = GeneratorSource::default();
262        b.start(start_request(
263            config,
264            SourcePosition::Resume {
265                attempt: CheckpointAttempt::new(1, 1),
266                checkpoint: cp,
267            },
268        ))
269        .await
270        .unwrap();
271        let _ = b.poll_batch(4).await.unwrap();
272        std::thread::sleep(std::time::Duration::from_millis(10));
273        let from_a = a.poll_batch(4).await.unwrap().expect("rows").records;
274        let from_b = b.poll_batch(4).await.unwrap().expect("rows").records;
275        assert_eq!(from_a, from_b, "replay from offset must be byte-identical");
276    }
277
278    #[tokio::test]
279    async fn rate_limit_and_max_rows_bound_emission() {
280        let mut config = ConnectorConfig::new("generator");
281        config.set("rows.per.second", "1000");
282        config.set("max.rows", "3");
283        let mut g = GeneratorSource::default();
284        g.start(start_request(config, SourcePosition::Initial))
285            .await
286            .unwrap();
287        let _ = g.poll_batch(100).await.unwrap(); // anchor
288        std::thread::sleep(std::time::Duration::from_millis(50)); // ~50 tokens
289        let batch = g.poll_batch(100).await.unwrap().expect("rows");
290        assert_eq!(batch.num_rows(), 3, "max.rows caps emission");
291        assert!(g.poll_batch(100).await.unwrap().is_none(), "exhausted");
292    }
293
294    #[tokio::test]
295    async fn malformed_resume_fails_before_rate_anchor() {
296        let mut checkpoint = SourceCheckpoint::new();
297        checkpoint.set_offset("seq", "not-a-sequence");
298        let mut source = GeneratorSource::default();
299        let error = source
300            .start(start_request(
301                ConnectorConfig::new("generator"),
302                SourcePosition::Resume {
303                    attempt: CheckpointAttempt::new(1, 1),
304                    checkpoint,
305                },
306            ))
307            .await
308            .expect_err("malformed durable cursor must fail closed");
309        assert!(error.to_string().contains("bad generator offset"));
310        assert!(source.anchor.is_none());
311    }
312
313    #[tokio::test]
314    async fn resume_at_finite_end_remains_exhausted() {
315        let mut config = ConnectorConfig::new("generator");
316        config.set("rows.per.second", "1000000");
317        config.set("max.rows", "8");
318        let mut checkpoint = SourceCheckpoint::new();
319        checkpoint.set_offset("seq", "8");
320        let mut source = GeneratorSource::default();
321        source
322            .start(start_request(
323                config,
324                SourcePosition::Resume {
325                    attempt: CheckpointAttempt::new(7, 7),
326                    checkpoint,
327                },
328            ))
329            .await
330            .unwrap();
331
332        assert!(source.poll_batch(8).await.unwrap().is_none());
333        std::thread::sleep(std::time::Duration::from_millis(2));
334        assert!(source.poll_batch(8).await.unwrap().is_none());
335        assert_eq!(source.checkpoint().get_offset("seq"), Some("8"));
336    }
337}
laminar_connectors/generator.rs

laminar_connectors/
generator.rs
