use std::sync::Arc;
use std::sync::atomic::{AtomicU64, Ordering};
use crate::node::{PortType, Value};
#[derive(Clone, Debug)]
pub struct SourceItem {
pub ordinal: u64,
pub fields: Vec<(String, Value)>,
}
impl SourceItem {
pub fn ordinal(ordinal: u64) -> Self {
Self { ordinal, fields: Vec::new() }
}
pub fn with_fields(ordinal: u64, fields: Vec<(String, Value)>) -> Self {
Self { ordinal, fields }
}
pub fn field(&self, name: &str) -> Option<&Value> {
self.fields.iter().find(|(n, _)| n == name).map(|(_, v)| v)
}
}
#[derive(Clone, Debug)]
pub struct SourceSchema {
pub name: String,
pub projections: Vec<(String, PortType)>,
pub extent: Option<u64>,
pub extent_outputs: Option<(String, String)>,
pub extent_limit: Option<u64>,
pub cursor_kind: CursorKind,
pub partition_output: Option<String>,
}
#[derive(Clone, Debug, Default)]
pub enum CursorKind {
#[default]
Range,
ExtendingTimed {
min_ms_output: String,
delta_output: Option<String>,
},
ExtendingPasses {
min_passes_output: String,
delta_output: Option<String>,
},
ExtendingCount {
min_count_output: String,
delta_output: Option<String>,
},
ExtendingElapsedAndPasses {
min_ms_output: String,
min_passes_output: String,
delta_output: Option<String>,
},
ExtendingElapsedOrPasses {
min_ms_output: String,
min_passes_output: String,
delta_output: Option<String>,
},
}
pub trait DataSource: Send {
fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>>;
fn next(&mut self) -> Option<SourceItem> {
let range = self.reserve(1)?;
Some(self.render_item(range.start))
}
fn next_chunk(&mut self, limit: usize) -> Vec<SourceItem> {
let range = match self.reserve(limit) {
Some(r) => r,
None => return Vec::new(),
};
(range.start..range.end)
.map(|ordinal| self.render_item(ordinal))
.collect()
}
fn render_item(&self, ordinal: u64) -> SourceItem;
fn extent(&self) -> Option<u64>;
fn consumed(&self) -> u64;
fn schema(&self) -> &SourceSchema;
}
pub trait DataSourceFactory: Send + Sync {
fn create_reader(&self) -> Box<dyn DataSource>;
fn schema(&self) -> &SourceSchema;
fn global_consumed(&self) -> u64;
fn global_extent(&self) -> Option<u64> {
self.schema().extent
}
fn rewind_for_poll(&self) -> bool {
false
}
}
pub struct RangeSourceFactory {
cursor: Arc<AtomicU64>,
end: u64,
schema: SourceSchema,
}
impl RangeSourceFactory {
pub fn new(start: u64, end: u64) -> Self {
Self {
cursor: Arc::new(AtomicU64::new(start)),
end,
schema: SourceSchema {
name: "_range".into(),
projections: vec![("ordinal".into(), PortType::U64)],
extent: Some(end.saturating_sub(start)),
extent_outputs: None,
extent_limit: None,
cursor_kind: CursorKind::Range,
partition_output: None,
},
}
}
pub fn named(name: &str, start: u64, end: u64) -> Self {
let mut factory = Self::new(start, end);
factory.schema.name = name.to_string();
factory
}
}
impl DataSourceFactory for RangeSourceFactory {
fn create_reader(&self) -> Box<dyn DataSource> {
Box::new(RangeSource {
cursor: self.cursor.clone(),
end: self.end,
consumed: 0,
schema: self.schema.clone(),
})
}
fn schema(&self) -> &SourceSchema {
&self.schema
}
fn global_consumed(&self) -> u64 {
let pos = self.cursor.load(Ordering::Relaxed);
let start = self.end.saturating_sub(self.schema.extent.unwrap_or(0));
pos.saturating_sub(start).min(self.schema.extent.unwrap_or(u64::MAX))
}
fn rewind_for_poll(&self) -> bool {
let start = self.end.saturating_sub(self.schema.extent.unwrap_or(0));
self.cursor.store(start, Ordering::Relaxed);
true
}
}
struct RangeSource {
cursor: Arc<AtomicU64>,
end: u64,
consumed: u64,
schema: SourceSchema,
}
impl DataSource for RangeSource {
fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>> {
let base = self.cursor.fetch_add(stride as u64, Ordering::Relaxed);
if base >= self.end {
return None;
}
let actual_end = (base + stride as u64).min(self.end);
let count = actual_end - base;
self.consumed += count;
Some(base..actual_end)
}
fn render_item(&self, ordinal: u64) -> SourceItem {
SourceItem::ordinal(ordinal)
}
fn extent(&self) -> Option<u64> {
self.schema.extent
}
fn consumed(&self) -> u64 {
self.consumed
}
fn schema(&self) -> &SourceSchema {
&self.schema
}
}
#[derive(Clone, Copy, Debug)]
pub struct ExtensionContext {
pub elapsed_ms: u64,
pub consumed: u64,
pub base: u64,
}
impl ExtensionContext {
pub fn passes(&self) -> u64 {
if self.base == 0 { 0 } else { self.consumed / self.base }
}
}
pub trait ExtensionPolicy: Send + Sync {
fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64>;
}
pub struct ExtendingRangeSourceFactory {
cursor: Arc<AtomicU64>,
end: Arc<AtomicU64>,
start: u64,
base: u64,
started: std::time::Instant,
policy: Arc<dyn ExtensionPolicy>,
schema: SourceSchema,
}
impl ExtendingRangeSourceFactory {
pub fn new(
name: &str,
start: u64,
initial_extent: u64,
policy: Arc<dyn ExtensionPolicy>,
) -> Self {
let end = start.saturating_add(initial_extent);
Self {
cursor: Arc::new(AtomicU64::new(start)),
end: Arc::new(AtomicU64::new(end)),
start,
base: initial_extent,
started: std::time::Instant::now(),
policy,
schema: SourceSchema {
name: name.to_string(),
projections: vec![("ordinal".into(), PortType::U64)],
extent: Some(initial_extent),
extent_outputs: None,
cursor_kind: CursorKind::Range,
extent_limit: None,
partition_output: None,
},
}
}
}
impl DataSourceFactory for ExtendingRangeSourceFactory {
fn create_reader(&self) -> Box<dyn DataSource> {
Box::new(ExtendingRangeSource {
cursor: self.cursor.clone(),
end: self.end.clone(),
policy: self.policy.clone(),
start: self.start,
base: self.base,
started: self.started,
consumed: 0,
schema: self.schema.clone(),
})
}
fn schema(&self) -> &SourceSchema {
&self.schema
}
fn global_consumed(&self) -> u64 {
let pos = self.cursor.load(Ordering::Relaxed);
pos.saturating_sub(self.start)
}
fn global_extent(&self) -> Option<u64> {
Some(self.end.load(Ordering::Acquire)
.saturating_sub(self.start))
}
}
struct ExtendingRangeSource {
cursor: Arc<AtomicU64>,
end: Arc<AtomicU64>,
policy: Arc<dyn ExtensionPolicy>,
start: u64,
base: u64,
started: std::time::Instant,
consumed: u64,
schema: SourceSchema,
}
impl DataSource for ExtendingRangeSource {
fn reserve(&mut self, stride: usize) -> Option<std::ops::Range<u64>> {
loop {
let cur = self.cursor.load(Ordering::Acquire);
let end = self.end.load(Ordering::Acquire);
if cur < end {
let target = (cur.saturating_add(stride as u64)).min(end);
match self.cursor.compare_exchange(
cur, target,
Ordering::AcqRel, Ordering::Acquire,
) {
Ok(_) => {
let count = target - cur;
self.consumed += count;
return Some(cur..target);
}
Err(_) => continue, }
}
let ctx = ExtensionContext {
elapsed_ms: self.started.elapsed().as_millis() as u64,
consumed: end.saturating_sub(self.start),
base: self.base,
};
match self.policy.next_extension(&ctx) {
Some(delta) if delta > 0 => {
let _ = self.end.compare_exchange(
end, end.saturating_add(delta),
Ordering::AcqRel, Ordering::Acquire,
);
continue;
}
_ => return None,
}
}
}
fn render_item(&self, ordinal: u64) -> SourceItem {
SourceItem::ordinal(ordinal)
}
fn extent(&self) -> Option<u64> {
Some(self.end.load(Ordering::Acquire)
.saturating_sub(self.start))
}
fn consumed(&self) -> u64 {
self.consumed
}
fn schema(&self) -> &SourceSchema {
&self.schema
}
}
struct CursorTarget {
reader: Box<dyn DataSource>,
input_index: usize,
#[allow(dead_code)]
source_name: String,
}
pub struct Cursors {
targets: Vec<CursorTarget>,
last_items: Vec<Option<SourceItem>>,
advances: u64,
}
impl Cursors {
pub fn for_fields(
program: &crate::kernel::GkProgram,
field_names: &[&str],
source_factories: &std::collections::HashMap<String, Arc<dyn DataSourceFactory>>,
) -> Self {
let mut combined_provenance: u64 = 0;
for name in field_names {
if let Some((node_idx, _)) = program.resolve_output(name) {
combined_provenance |= program.input_provenance_for(node_idx);
}
}
let input_names = program.input_names();
let mut targets = Vec::new();
let mut seen_sources = std::collections::HashSet::new();
for (idx, input_name) in input_names.iter().enumerate() {
if combined_provenance & (1u64 << idx) == 0 { continue; }
if let Some(source_name) = input_name.strip_suffix("__ordinal") {
if seen_sources.contains(source_name) { continue; }
seen_sources.insert(source_name.to_string());
if let Some(factory) = source_factories.get(source_name) {
targets.push(CursorTarget {
reader: factory.create_reader(),
input_index: idx,
source_name: source_name.to_string(),
});
}
}
}
let target_count = targets.len();
Cursors {
targets,
last_items: vec![None; target_count],
advances: 0,
}
}
pub fn advance(&mut self) -> bool {
for (i, target) in self.targets.iter_mut().enumerate() {
match target.reader.next() {
Some(item) => {
self.last_items[i] = Some(item);
}
None => return false, }
}
self.advances += 1;
true
}
pub fn inject_into_state(&self, state: &mut crate::kernel::GkState) {
for (i, target) in self.targets.iter().enumerate() {
if let Some(ref item) = self.last_items[i] {
state.set_input(target.input_index, crate::node::Value::U64(item.ordinal));
}
}
}
pub fn last_items(&self) -> &[Option<SourceItem>] {
&self.last_items
}
pub fn extent(&self) -> Option<u64> {
self.targets.iter()
.filter_map(|t| t.reader.extent())
.min()
}
pub fn consumed(&self) -> u64 {
self.advances
}
pub fn target_count(&self) -> usize {
self.targets.len()
}
pub fn is_empty(&self) -> bool {
self.targets.is_empty()
}
}
pub struct UntilElapsedPolicy { pub min_ms: u64, pub delta: u64 }
impl ExtensionPolicy for UntilElapsedPolicy {
fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
if ctx.elapsed_ms >= self.min_ms { return None; }
if ctx.elapsed_ms == 0 || ctx.consumed == 0 {
return Some(self.delta.max(1));
}
let remaining_ms = self.min_ms - ctx.elapsed_ms;
let est_remaining = (ctx.consumed as u128)
.saturating_mul(remaining_ms as u128)
/ (ctx.elapsed_ms as u128);
let biased = est_remaining.saturating_mul(95) / 100;
let base = ctx.base.max(1) as u128;
let repeats = biased / base;
if repeats == 0 { return None; }
let delta = repeats.saturating_mul(base);
Some(u64::try_from(delta).unwrap_or(u64::MAX))
}
}
pub struct UntilPassesPolicy { pub min_passes: u64, pub delta: u64 }
impl ExtensionPolicy for UntilPassesPolicy {
fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
if ctx.passes() < self.min_passes { Some(self.delta) } else { None }
}
}
pub struct UntilCountPolicy { pub min_count: u64, pub delta: u64 }
impl ExtensionPolicy for UntilCountPolicy {
fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
if ctx.consumed < self.min_count { Some(self.delta) } else { None }
}
}
pub struct AndPolicy { pub policies: Vec<Arc<dyn ExtensionPolicy>> }
impl ExtensionPolicy for AndPolicy {
fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
let mut min_delta = u64::MAX;
for p in &self.policies {
match p.next_extension(ctx) {
Some(d) => min_delta = min_delta.min(d),
None => return None,
}
}
if min_delta == u64::MAX || min_delta == 0 { None } else { Some(min_delta) }
}
}
pub struct OrPolicy { pub policies: Vec<Arc<dyn ExtensionPolicy>> }
impl ExtensionPolicy for OrPolicy {
fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
let mut max_delta: Option<u64> = None;
for p in &self.policies {
if let Some(d) = p.next_extension(ctx) {
max_delta = Some(max_delta.map(|m| m.max(d)).unwrap_or(d));
}
}
max_delta
}
}
pub struct TimeElapsedPolicy { inner: UntilElapsedPolicy }
impl TimeElapsedPolicy {
pub fn new(base: u64, min_ms: u64) -> Self {
Self { inner: UntilElapsedPolicy { min_ms, delta: base } }
}
}
impl ExtensionPolicy for TimeElapsedPolicy {
fn next_extension(&self, ctx: &ExtensionContext) -> Option<u64> {
self.inner.next_extension(ctx)
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn range_source_yields_ordinals() {
let factory = RangeSourceFactory::new(0, 5);
let mut reader = factory.create_reader();
assert_eq!(reader.extent(), Some(5));
for i in 0..5 {
let item = reader.next().unwrap();
assert_eq!(item.ordinal, i);
assert!(item.fields.is_empty());
}
assert!(reader.next().is_none());
assert_eq!(reader.consumed(), 5);
}
#[test]
fn range_source_chunk() {
let factory = RangeSourceFactory::new(0, 10);
let mut reader = factory.create_reader();
let chunk = reader.next_chunk(3);
assert_eq!(chunk.len(), 3);
assert_eq!(chunk[0].ordinal, 0);
assert_eq!(chunk[2].ordinal, 2);
let chunk = reader.next_chunk(100);
assert_eq!(chunk.len(), 7); assert_eq!(chunk[0].ordinal, 3);
assert_eq!(chunk[6].ordinal, 9);
let chunk = reader.next_chunk(1);
assert!(chunk.is_empty()); }
#[test]
fn range_source_concurrent_readers() {
let factory = RangeSourceFactory::new(0, 100);
let mut r1 = factory.create_reader();
let mut r2 = factory.create_reader();
let a = r1.next().unwrap().ordinal;
let b = r2.next().unwrap().ordinal;
assert_ne!(a, b);
let mut total = 2;
while r1.next().is_some() { total += 1; }
while r2.next().is_some() { total += 1; }
assert_eq!(total, 100);
}
#[test]
fn source_item_field_access() {
let item = SourceItem::with_fields(42, vec![
("name".into(), Value::Str("test".into())),
("score".into(), Value::F64(0.95)),
]);
assert_eq!(item.ordinal, 42);
assert_eq!(item.field("name"), Some(&Value::Str("test".into())));
assert_eq!(item.field("score"), Some(&Value::F64(0.95)));
assert_eq!(item.field("missing"), None);
}
#[test]
fn range_source_named() {
let factory = RangeSourceFactory::named("users", 0, 1000);
assert_eq!(factory.schema().name, "users");
assert_eq!(factory.schema().extent, Some(1000));
}
struct FixedExtensions {
delta: u64,
remaining: std::sync::atomic::AtomicU64,
}
impl FixedExtensions {
fn new(delta: u64, times: u64) -> Self {
Self {
delta,
remaining: std::sync::atomic::AtomicU64::new(times),
}
}
}
impl ExtensionPolicy for FixedExtensions {
fn next_extension(&self, _ctx: &ExtensionContext) -> Option<u64> {
let prev = self.remaining.fetch_sub(1, Ordering::Relaxed);
if prev == 0 || prev > i64::MAX as u64 {
self.remaining.store(0, Ordering::Relaxed);
None
} else {
Some(self.delta)
}
}
}
#[test]
fn extending_source_consumes_initial_extent_then_extends() {
let policy = Arc::new(FixedExtensions::new(5, 1));
let factory = ExtendingRangeSourceFactory::new("ext", 0, 5, policy);
let mut reader = factory.create_reader();
let mut got: Vec<u64> = Vec::new();
while let Some(item) = reader.next() {
got.push(item.ordinal);
if got.len() > 50 { panic!("runaway extension"); }
}
assert_eq!(got, (0..10).collect::<Vec<u64>>());
assert_eq!(reader.consumed(), 10);
}
#[test]
fn extending_source_zero_extensions_behaves_like_fixed_range() {
let policy = Arc::new(FixedExtensions::new(0, 0));
let factory = ExtendingRangeSourceFactory::new("ext", 0, 3, policy);
let mut reader = factory.create_reader();
let mut got: Vec<u64> = Vec::new();
while let Some(item) = reader.next() {
got.push(item.ordinal);
}
assert_eq!(got, vec![0, 1, 2]);
}
#[test]
fn extending_source_global_extent_grows_after_extension() {
let policy = Arc::new(FixedExtensions::new(10, 2));
let factory = ExtendingRangeSourceFactory::new("ext", 0, 5, policy);
assert_eq!(factory.global_extent(), Some(5));
let mut reader = factory.create_reader();
for _ in 0..5 { reader.next().unwrap(); }
let _ = reader.next().unwrap();
assert_eq!(factory.global_extent(), Some(15),
"extent should grow by the extension delta");
}
#[test]
fn extending_source_chunked_reservation_caps_at_current_end() {
let policy = Arc::new(FixedExtensions::new(5, 1));
let factory = ExtendingRangeSourceFactory::new("ext", 0, 3, policy);
let mut reader = factory.create_reader();
let range = reader.reserve(10).expect("first reserve");
assert_eq!(range, 0..3);
let range = reader.reserve(10).expect("second reserve");
assert_eq!(range, 3..8);
}
fn ctx_at(elapsed_ms: u64, consumed: u64, base: u64) -> ExtensionContext {
ExtensionContext { elapsed_ms, consumed, base }
}
#[test]
fn until_elapsed_policy_bootstrap_falls_back_to_delta_when_no_rate_signal() {
let policy = UntilElapsedPolicy { min_ms: 50, delta: 7 };
assert_eq!(policy.next_extension(&ctx_at(0, 0, 0)), Some(7));
assert_eq!(policy.next_extension(&ctx_at(49, 0, 0)), Some(7));
assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(7));
}
#[test]
fn until_elapsed_policy_stops_at_or_past_min_ms() {
let policy = UntilElapsedPolicy { min_ms: 50, delta: 7 };
assert_eq!(policy.next_extension(&ctx_at(50, 100, 100)), None);
assert_eq!(policy.next_extension(&ctx_at(1000, 100, 100)), None);
}
#[test]
fn until_elapsed_policy_projects_remaining_cycles_with_under_bias() {
let policy = UntilElapsedPolicy { min_ms: 1000, delta: 100 };
assert_eq!(
policy.next_extension(&ctx_at(200, 100, 100)),
Some(300),
"expected 3-pass batch (380 under-biased, floored to base multiples)",
);
}
#[test]
fn until_elapsed_policy_returns_none_when_remaining_is_under_one_pass() {
let policy = UntilElapsedPolicy { min_ms: 1000, delta: 100 };
assert_eq!(policy.next_extension(&ctx_at(990, 10000, 100)), None);
}
#[test]
fn until_elapsed_policy_converges_geometrically() {
let policy = UntilElapsedPolicy { min_ms: 1000, delta: 10 };
let mut elapsed = 10u64;
let mut consumed = 10u64;
let mut iters = 0;
while let Some(delta) = policy.next_extension(&ctx_at(elapsed, consumed, 10)) {
iters += 1;
assert!(iters < 10, "geometric series should converge fast");
consumed += delta;
elapsed += delta;
}
assert!(elapsed <= 1000, "must under-shoot, got elapsed={elapsed}");
assert!(elapsed >= 940, "must come within ~6% of target, got {elapsed}");
}
#[test]
fn until_passes_policy_counts_in_base_multiples() {
let policy = UntilPassesPolicy { min_passes: 3, delta: 100 };
assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(100));
assert_eq!(policy.next_extension(&ctx_at(0, 200, 100)), Some(100));
assert_eq!(policy.next_extension(&ctx_at(0, 300, 100)), None);
assert_eq!(policy.next_extension(&ctx_at(0, 400, 100)), None);
}
#[test]
fn until_count_policy_uses_raw_consumed() {
let policy = UntilCountPolicy { min_count: 250, delta: 50 };
assert_eq!(policy.next_extension(&ctx_at(0, 0, 100)), Some(50));
assert_eq!(policy.next_extension(&ctx_at(0, 249, 100)), Some(50));
assert_eq!(policy.next_extension(&ctx_at(0, 250, 100)), None);
}
#[test]
fn and_policy_stops_when_any_child_stops() {
let time = Arc::new(UntilElapsedPolicy { min_ms: 5000, delta: 10 });
let passes = Arc::new(UntilPassesPolicy { min_passes: 3, delta: 20 });
let and = AndPolicy { policies: vec![time, passes] };
assert_eq!(and.next_extension(&ctx_at(0, 0, 100)), Some(10));
assert_eq!(and.next_extension(&ctx_at(5000, 0, 100)), None);
assert_eq!(and.next_extension(&ctx_at(0, 300, 100)), None);
}
#[test]
fn or_policy_continues_if_any_child_continues() {
let time = Arc::new(UntilElapsedPolicy { min_ms: 5000, delta: 10 });
let passes = Arc::new(UntilPassesPolicy { min_passes: 3, delta: 20 });
let or = OrPolicy { policies: vec![time, passes] };
assert_eq!(or.next_extension(&ctx_at(0, 0, 100)), Some(20));
assert_eq!(or.next_extension(&ctx_at(5000, 0, 100)), Some(20));
assert_eq!(or.next_extension(&ctx_at(0, 300, 100)), Some(10));
assert_eq!(or.next_extension(&ctx_at(5000, 300, 100)), None);
}
#[test]
fn time_elapsed_policy_compat_alias_still_works() {
let p = TimeElapsedPolicy::new(7, 1);
std::thread::sleep(std::time::Duration::from_millis(20));
assert_eq!(p.next_extension(&ctx_at(20, 0, 0)), None);
}
}