rushdb-query 3.0.4

//! Per-query working-set cap (config #3): a query that would materialize more
//! than the configured budget returns `resource exhausted` instead of
//! ballooning RSS. The default (0) is off, so existing behavior is unchanged.

use std::sync::Arc;

use rushdb_lmpg::plan::FieldChange;
use rushdb_lmpg::Lmpg;
use rushdb_query::document::parse_document;
use rushdb_query::{find_records_in_view, QueryError};
use rushdb_storage::resources::StorageResources;
use rushdb_storage::Store;
use rushdb_types::ids::TenantId;
use serde_json::json;

fn store_with_cap(cap: u64) -> Arc<Store> {
    Arc::new(
        Store::new_with_resources(StorageResources {
            query_memory_max_bytes: cap,
            ..StorageResources::default()
        })
        .unwrap(),
    )
}

fn seed(lmpg: &Lmpg) {
    for i in 0..300 {
        lmpg.create_record("Big", None, vec![("n".into(), FieldChange::Set(json!(i)))])
            .unwrap();
    }
}

#[test]
fn query_memory_cap_rejects_an_oversized_working_set() {
    let store = store_with_cap(1024);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    seed(&lmpg);
    let parsed = parse_document(&json!({ "labels": ["Big"] })).unwrap();
    let view = Arc::new(store.read_view().unwrap());
    let err = find_records_in_view(&lmpg, &parsed.search, &|| false, view).unwrap_err();
    assert!(
        matches!(err, QueryError::ResourceLimit(_)),
        "expected a resource cap, got {err:?}"
    );
}

#[test]
fn cancellation_interrupts_candidate_discovery() {
    // The unscoped candidate scan must poll cancellation between batches, not
    // only per-candidate during predicate evaluation.
    let store = store_with_cap(0);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    for i in 0..5000u64 {
        lmpg.create_record("Big", None, vec![("n".into(), FieldChange::Set(json!(i)))])
            .unwrap();
    }
    let view = Arc::new(store.read_view().unwrap());
    let parsed = parse_document(&json!({})).unwrap();
    let polls = std::sync::atomic::AtomicUsize::new(0);
    let cancelled = || {
        // First poll (the pre-dispatch check) passes; the next (inside the
        // candidate scan) cancels.
        polls.fetch_add(1, std::sync::atomic::Ordering::Relaxed) + 1 >= 2
    };
    let err = find_records_in_view(&lmpg, &parsed.search, &cancelled, view).unwrap_err();
    assert!(matches!(err, QueryError::Cancelled), "{err:?}");
}

#[test]
fn query_memory_cap_is_off_by_default() {
    let store = store_with_cap(0);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    seed(&lmpg);
    let parsed = parse_document(&json!({ "labels": ["Big"] })).unwrap();
    let view = Arc::new(store.read_view().unwrap());
    let page = find_records_in_view(&lmpg, &parsed.search, &|| false, view).unwrap();
    assert_eq!(page.total_matched, 300);
}

#[test]
fn record_collect_respects_the_configured_working_set_cap() {
    use rushdb_query::{CollectSpec, Direction};
    let store = store_with_cap(64 * 1024);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    let root = lmpg.create_record("Root", None, vec![]).unwrap();
    for _ in 0..4 {
        let child = lmpg
            .create_record(
                "Child",
                None,
                vec![(
                    "payload".into(),
                    FieldChange::Set(json!("x".repeat(32 * 1024))),
                )],
            )
            .unwrap();
        lmpg.create_relationship("HAS", root, child, vec![])
            .unwrap();
    }
    let mut parsed = parse_document(&json!({ "labels": ["Root"] })).unwrap();
    parsed.search.collects = vec![CollectSpec {
        steps: vec![],
        alias: "children".into(),
        rel_type: Some("HAS".into()),
        direction: Direction::Outgoing,
        target_labels: vec![],
        fields: vec![],
        nested: vec![],
        filter: None,
        order_by: vec![],
        skip: 0,
        unique: true,
        limit: None,
    }];
    let view = Arc::new(store.read_view().unwrap());
    let err = find_records_in_view(&lmpg, &parsed.search, &|| false, view).unwrap_err();
    assert!(
        matches!(err, QueryError::ResourceLimit(_) | QueryError::Storage(_)),
        "collect expansion must be charged against the working-set cap: {err:?}"
    );
}

#[test]
fn records_explicit_zero_budget_does_not_mean_unlimited() {
    // An exhausted request budget arrives as `Some(0)`. It must be a hard zero
    // cap, not silently disable the limit and admit the whole result set.
    let store = store_with_cap(0);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    seed(&lmpg);
    let parsed = parse_document(&json!({ "labels": ["Big"] })).unwrap();
    let view = Arc::new(store.read_view().unwrap());
    let err = rushdb_query::find_records_in_view_with_budget(
        &lmpg,
        &parsed.search,
        &|| false,
        view,
        Some(0),
    )
    .unwrap_err();
    assert!(
        matches!(err, QueryError::ResourceLimit(_)),
        "explicit zero budget must reject, got {err:?}"
    );
}

#[test]
fn collect_expansion_is_cumulative_across_page_records() {
    // `$collect` must be charged cumulatively: two roots whose collected
    // children together exceed the budget must fail even though each root
    // alone would fit the full remaining budget.
    use rushdb_query::{CollectSpec, Direction};
    let store = store_with_cap(96 * 1024);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    for _ in 0..2 {
        let root = lmpg.create_record("Root", None, vec![]).unwrap();
        let child = lmpg
            .create_record(
                "Child",
                None,
                vec![(
                    "payload".into(),
                    FieldChange::Set(json!("x".repeat(64 * 1024))),
                )],
            )
            .unwrap();
        lmpg.create_relationship("HAS", root, child, vec![])
            .unwrap();
    }
    let mut parsed = parse_document(&json!({ "labels": ["Root"] })).unwrap();
    parsed.search.collects = vec![CollectSpec {
        steps: vec![],
        alias: "children".into(),
        rel_type: Some("HAS".into()),
        direction: Direction::Outgoing,
        target_labels: vec![],
        fields: vec![],
        nested: vec![],
        filter: None,
        order_by: vec![],
        skip: 0,
        unique: true,
        limit: None,
    }];
    let view = Arc::new(store.read_view().unwrap());
    let err = find_records_in_view(&lmpg, &parsed.search, &|| false, view).unwrap_err();
    assert!(
        matches!(err, QueryError::ResourceLimit(_) | QueryError::Storage(_)),
        "collect expansion must be cumulative across records: {err:?}"
    );
}

#[test]
fn aggregate_applies_the_configured_limit() {
    let store = store_with_cap(128);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    let result = rushdb_query::aggregate_with_cancel_in_view(
        &lmpg,
        &rushdb_query::AggregateRequest::default(),
        &|| false,
        Arc::new(store.read_view().unwrap()),
    );
    assert!(matches!(result, Err(QueryError::ResourceLimit(_))));
}

/// Seeds `count` trips with a few scalar columns in one commit.
fn seed_trips(lmpg: &Lmpg, count: u32) {
    let rows: Vec<serde_json::Value> = (0..count)
        .map(|i| {
            json!({
                "vendor": if i % 3 == 0 { "A" } else { "B" },
                "hour": i % 24,
                "fare": f64::from(i % 50) + 0.5,
                "distance": f64::from(i % 7),
                "note": format!("trip {i}"),
            })
        })
        .collect();
    let options = rushdb_lmpg::import::ImportOptions::default();
    let plan = rushdb_lmpg::import::serialize_bfs(&json!(rows), "Trip", &options);
    rushdb_lmpg::import::import_bfs(lmpg, &plan, &options).unwrap();
}

#[test]
fn grouped_averages_over_many_records_fit_query_memory_max() {
    // The aggregate working set is the match ordinals, the row being
    // aggregated and the buckets, under `query-memory-max`. Before, the
    // candidate vector's *capacity* was charged twice over (16 B then 8 B
    // per slot) and tripled again for an `$avg`, so 30,000 matches needed
    // over 768 KiB and every row was rebuilt in full (the `note` column
    // included). 30,000 ordinals are 240,000 bytes retained, and the walk
    // charged twice that.
    let store = store_with_cap(600_000);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    seed_trips(&lmpg, 30_000);
    let request = rushdb_query::AggregateRequest {
        labels: vec!["Trip".into()],
        group_by: vec!["vendor".into()],
        avg_terms: vec![("dist".into(), "distance".into())],
        maxs: vec![("max_fare".into(), "fare".into())],
        counts: vec![("n".into(), "*".into())],
        ..Default::default()
    };
    let view = Arc::new(store.read_view().unwrap());
    let (buckets, total) =
        rushdb_query::aggregate_with_cancel_in_view(&lmpg, &request, &|| false, view).unwrap();
    assert_eq!(total, 30_000);
    assert_eq!(buckets.len(), 2);
    let a = buckets
        .iter()
        .find(|b| b.keys["vendor"] == json!("A"))
        .unwrap();
    assert_eq!(a.count_roots, 10_000);
    assert_eq!(a.maxs["max_fare"], json!(49.5));
    let expected_distance = (0..30_000u32)
        .filter(|i| i % 3 == 0)
        .map(|i| f64::from(i % 7))
        .sum::<f64>()
        / 10_000.0;
    assert_eq!(a.sums["dist"], json!(expected_distance));

    // The cap still binds the retained ordinals: 30,000 do not fit 200 KiB.
    let store = store_with_cap(200_000);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    seed_trips(&lmpg, 30_000);
    let view = Arc::new(store.read_view().unwrap());
    let error =
        rushdb_query::aggregate_with_cancel_in_view(&lmpg, &request, &|| false, view).unwrap_err();
    assert!(
        matches!(&error, QueryError::ResourceLimit(message) if message.contains("query-memory-max")),
        "{error:?}"
    );
}

#[test]
fn aggregate_limits_retained_groups_under_query_memory_max() {
    // Buckets are what an aggregate retains: 100 groups with 1 KiB keys do
    // not fit a 64 KiB working set, whether or not `skip` discards them.
    let store = store_with_cap(64 * 1024);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    for i in 0..100 {
        lmpg.create_record(
            "Groups",
            None,
            vec![(
                "group".into(),
                FieldChange::Set(json!(format!("{i}{}", "x".repeat(1024)))),
            )],
        )
        .unwrap();
    }
    let request = rushdb_query::AggregateRequest {
        labels: vec!["Groups".into()],
        group_by: vec!["group".into()],
        skip: 100,
        ..Default::default()
    };
    let view = Arc::new(store.read_view().unwrap());
    let result = rushdb_query::aggregate_with_cancel_in_view(&lmpg, &request, &|| false, view);
    assert!(
        matches!(result, Err(QueryError::ResourceLimit(_))),
        "{result:?}"
    );
}

#[test]
fn a_filtered_grouped_aggregate_reads_only_the_fields_it_names() {
    // The residual filter and the metrics read their properties by point
    // reads; a 64 KiB payload column is never decoded, so the working set
    // stays near the ordinals.
    let store = store_with_cap(96 * 1024);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    for i in 0..40 {
        lmpg.create_record(
            "Wide",
            None,
            vec![
                ("kind".into(), FieldChange::Set(json!(i % 2))),
                ("n".into(), FieldChange::Set(json!(i))),
                (
                    "payload".into(),
                    FieldChange::Set(json!("x".repeat(64 * 1024))),
                ),
            ],
        )
        .unwrap();
    }
    let parsed =
        parse_document(&json!({"labels": ["Wide"], "where": {"n": {"$gte": 20}}})).unwrap();
    let request = rushdb_query::AggregateRequest {
        labels: parsed.search.labels,
        where_: parsed.search.where_,
        group_by: vec!["kind".into()],
        sums: vec![("sum".into(), "n".into())],
        ..Default::default()
    };
    let view = Arc::new(store.read_view().unwrap());
    let (buckets, total) =
        rushdb_query::aggregate_with_cancel_in_view(&lmpg, &request, &|| false, view).unwrap();
    assert_eq!(total, 20);
    let even = buckets.iter().find(|b| b.keys["kind"] == json!(0)).unwrap();
    assert_eq!(
        even.sums["sum"],
        json!(f64::from((20..40).step_by(2).sum::<i32>()))
    );
}

#[test]
fn aggregate_collect_budget_covers_work_discarded_by_child_pagination() {
    use rushdb_query::{AggregateRequest, CollectSpec, Direction};
    let store = store_with_cap(0);
    let lmpg = Lmpg::new(store.clone(), TenantId::single());
    let root = rushdb_lmpg::import::import_document(&lmpg, "Root", &json!({"n": 0})).unwrap();
    let child = rushdb_lmpg::import::import_document(
        &lmpg,
        "Child",
        &json!({"payload": "x".repeat(128 * 1024)}),
    )
    .unwrap();
    lmpg.create_relationship("HAS", root, child, vec![])
        .unwrap();
    let request = AggregateRequest {
        labels: vec!["Root".into()],
        collects: vec![CollectSpec {
            steps: vec![],
            alias: "children".into(),
            rel_type: Some("HAS".into()),
            direction: Direction::Outgoing,
            target_labels: vec![],
            fields: vec![],
            nested: vec![],
            filter: None,
            order_by: vec![],
            skip: 1,
            unique: true,
            limit: None,
        }],
        ..Default::default()
    };
    let view = Arc::new(store.read_view().unwrap());
    let (buckets, _) = rushdb_query::aggregate_in_view(&lmpg, &request, view.clone()).unwrap();
    let result = rushdb_query::aggregate_to_rows_in_view_with_budget(
        &lmpg,
        &view,
        &buckets,
        &request,
        Some(64 * 1024),
    );
    assert!(
        matches!(result, Err(QueryError::ResourceLimit(_))),
        "{result:?}"
    );
    let rows = rushdb_query::aggregate_to_rows_in_view_with_budget(
        &lmpg,
        &view,
        &buckets,
        &request,
        Some(8 * 1024 * 1024),
    )
    .unwrap();
    assert_eq!(rows[0]["children"], json!([]));
}

#[test]
fn the_response_budget_caps_the_page_not_the_working_set() {
    // READ-1: the request's response byte budget (the HTTP body limit) used
    // to cap the candidate and match ordinals too, so a label over ~2M
    // records refused even a 1-row page. The working set is capped by
    // `query-memory-max` alone, and it holds only the matches a page keeps:
    // candidates stream past and the total comes from the label counter.
    let read = |cap: u64, response_budget: usize, limit: u64| {
        let store = store_with_cap(cap);
        let lmpg = Lmpg::new(store.clone(), TenantId::single());
        seed(&lmpg);
        let parsed = parse_document(&json!({ "labels": ["Big"], "limit": limit })).unwrap();
        let view = Arc::new(store.read_view().unwrap());
        rushdb_query::find_records_in_view_with_budget(
            &lmpg,
            &parsed.search,
            &|| false,
            view,
            Some(response_budget),
        )
    };
    // 300 candidates, a 1-row page: fits a 1 KiB working set and a 1 KiB
    // response budget, with the exact total.
    let page = read(1024, 1024, 1).unwrap();
    assert_eq!(page.total_matched, 300);
    assert_eq!(page.records.len(), 1);
    // `query-memory-max` still binds what a page keeps: 300 ordinals are
    // 2,400 bytes.
    assert!(matches!(
        read(1024, 1024 * 1024, 1_000),
        Err(QueryError::ResourceLimit(_))
    ));
}