sort-governor 0.1.3

Process-wide governor for bounded asynchronous sorting: plans in-memory versus external sorts and rations file descriptors and spill memory
Documentation
//! Spill metadata bounds, ordinal exhaustion, and stable merge semantics.

use async_fs_io::DirectoryReader;
use futures_util::StreamExt;

use super::SortSession;
use crate::plan::SortPlan;

#[tokio::test]
async fn spill_frontier_stays_bounded_and_preserves_first_equal_row() {
    for rows in [255, 257] {
        for max_fan_in in [2, 3, 7] {
            for dedup in [false, true] {
                assert_bounded_sort(rows, max_fan_in, dedup).await;
            }
        }
    }
}

async fn assert_bounded_sort(rows: u32, max_fan_in: u32, dedup: bool) {
    let mut session = SortSession::with_temp_dir(
        SortPlan::External {
            run_buffer_bytes: 1,
            max_fan_in,
        },
        &std::env::temp_dir(),
        dedup,
    );
    for sequence in 0..rows {
        session
            .push_with_size(sequence % 7, sequence, 1)
            .await
            .expect("push");
        assert!(
            session.spills.len() <= 64,
            "spill metadata grows with input"
        );
        let Some(mut files) = DirectoryReader::open_if_exists(session.scratch_dir())
            .await
            .expect("open scratch")
        else {
            continue;
        };
        let mut count = 0;
        while files.next().await.expect("next file").is_some() {
            count += 1;
            assert!(count <= 64, "spill files grow with input before finish");
        }
    }
    let mut output = session.finish().await.expect("finish");
    for key in 0..7u32 {
        for sequence in (key..rows).step_by(7) {
            assert_eq!(output.next().await.expect("row").expect("value"), sequence);
            if dedup {
                break;
            }
        }
    }
    assert!(output.next().await.is_none());
}

#[tokio::test]
async fn exhausted_spill_ordinal_returns_error_before_creating_a_file() {
    let mut session = SortSession::with_temp_dir(
        SortPlan::External {
            run_buffer_bytes: 1,
            max_fan_in: 2,
        },
        &std::env::temp_dir(),
        false,
    );
    session.next_run = !0;
    session
        .push_with_size(0u32, 0u32, 1)
        .await
        .expect("first row");
    assert!(
        session.push_with_size(1, 1, 1).await.is_err(),
        "ordinal must not wrap"
    );
    assert!(
        DirectoryReader::open_if_exists(session.scratch_dir())
            .await
            .expect("inspect scratch")
            .is_none()
    );
}

#[tokio::test]
async fn failed_spill_cannot_be_retried_or_finished() {
    let root = async_fs_io::TempDir::create(std::env::temp_dir())
        .await
        .expect("temp dir");
    let dir = root.path().join("blocked");
    drop(
        async_fs_io::AsyncFile::create(&dir)
            .await
            .expect("block scratch"),
    );
    let mut session = SortSession::new(
        SortPlan::External {
            run_buffer_bytes: 1,
            max_fan_in: 2,
        },
        dir.clone(),
        false,
    );
    session
        .push_with_size(0u32, 0u32, 1)
        .await
        .expect("first row");
    assert!(session.push_with_size(1, 1, 1).await.is_err());
    async_fs_io::remove_if_exists(&dir)
        .await
        .expect("remove blocker");
    assert!(
        session.push_with_size(2, 2, 1).await.is_err(),
        "failed session must stay failed"
    );
    assert!(
        session.finish().await.is_err(),
        "failed session cannot yield partial results"
    );
    root.remove().await.expect("cleanup");
}