swactor/crates/datastore/tests/stream_integration.rs

//! End-to-end test: store a blob on node A, download via QUIC stream on node B.

use std::sync::Arc;
use std::time::{Duration, Instant};

use distribution::iroh_driver::{IrohDriver, IrohDriverConfig};
use distribution::node::DistributedNodeConfig;
use iroh::RelayMode;

use swactor::config::RuntimeConfig;
use swactor::runtime::Runtime;
use swactor_datastore::bridge::{DatastoreGroup, DatastoreGroupConfig};
use swactor_datastore::messages::{DatastoreNodeMsg, DatastoreResponse};
use swactor_std::RuntimeNaming;

fn make_driver_with_streams() -> IrohDriver {
    IrohDriver::new(IrohDriverConfig {
        secret_key: None,
        relay_mode: RelayMode::Disabled,
        node: DistributedNodeConfig::default(),
        peer_auth: None,
        additional_alpns: vec![swactor_streams::ALPN.to_vec()],
    })
    .expect("create iroh driver")
}

fn make_runtime() -> Arc<Runtime> {
    let rt = Runtime::new(RuntimeConfig {
        num_threads: 2,
        max_actors: 256,
        channel_buffer_size: 2000,
        ..Default::default()
    })
    .with_extension(Arc::new(swactor_std::StdExtension::new()));

    let handle = rt.run().expect("start runtime");
    handle.runtime
}

fn spawn_stream_manager(
    runtime: &Arc<Runtime>,
    driver: &IrohDriver,
) -> swactor::actor::ActorAddress {
    let mgr = swactor_streams::StreamManager::new(
        driver.endpoint().clone(),
        driver.tokio_handle(),
        Arc::clone(runtime),
    );
    let addr = runtime.spawn(mgr).expect("spawn StreamManager");
    runtime
        .register_name(swactor_streams::STREAM_MANAGER_NAME, addr)
        .expect("register StreamManager");
    addr
}

fn spawn_datastore(
    runtime: &Arc<Runtime>,
    driver: &IrohDriver,
    mgr_addr: swactor::actor::ActorAddress,
) -> DatastoreGroup {
    let node_id = driver.node_id();
    let group = DatastoreGroup::spawn(
        Arc::clone(runtime),
        DatastoreGroupConfig {
            node_id,
            node_id_hex: format!("{:?}", node_id),
            chunk_size: 256,
            storage_path: None, // in-memory
            auth: None,
            gc_interval: u64::MAX,
            disseminate_interval: u64::MAX,
        },
    )
    .expect("spawn datastore");
    group.configure_streams(mgr_addr, driver.tokio_handle());
    group
}

/// Poll for a response from the inbox, routing stream connections between
/// the two nodes. Actor message processing is handled by worker threads.
fn pump_until_response(
    rt_a: &Arc<Runtime>,
    rt_b: &Arc<Runtime>,
    driver_a: &mut IrohDriver,
    driver_b: &mut IrohDriver,
    mgr_a: swactor::actor::ActorAddress,
    mgr_b: swactor::actor::ActorAddress,
    inbox: &swactor::runtime::Inbox<DatastoreResponse>,
    timeout: Duration,
) -> Option<DatastoreResponse> {
    let deadline = Instant::now() + timeout;
    let tokio_handle = driver_a.tokio_handle();

    loop {
        // SWIM protocol ticks
        driver_a.recv();
        driver_a.tick();
        driver_b.recv();
        driver_b.tick();

        // Route incoming stream connections on node A
        for (node_id, conn) in driver_a.drain_other_connections() {
            let rt = Arc::clone(rt_a);
            let mgr = mgr_a;
            let node_bytes = node_id.0;
            tokio_handle.spawn(async move {
                let _ = swactor_streams::accept::handle_incoming(node_bytes, conn, &rt, mgr).await;
            });
        }

        // Route incoming stream connections on node B
        for (node_id, conn) in driver_b.drain_other_connections() {
            let rt = Arc::clone(rt_b);
            let mgr = mgr_b;
            let node_bytes = node_id.0;
            tokio_handle.spawn(async move {
                let _ = swactor_streams::accept::handle_incoming(node_bytes, conn, &rt, mgr).await;
            });
        }

        // Check for completion
        if let Some(resp) = inbox.try_recv() {
            return Some(resp);
        }

        if Instant::now() >= deadline {
            return None;
        }

        std::thread::sleep(Duration::from_millis(10));
    }
}

#[test]
fn small_blob_transfers_between_two_nodes_via_stream() {
    let mut driver_a = make_driver_with_streams();
    let mut driver_b = make_driver_with_streams();

    let rt_a = make_runtime();
    let rt_b = make_runtime();

    let mgr_a = spawn_stream_manager(&rt_a, &driver_a);
    let mgr_b = spawn_stream_manager(&rt_b, &driver_b);

    let _ds_a = spawn_datastore(&rt_a, &driver_a, mgr_a);
    let ds_b = spawn_datastore(&rt_b, &driver_b, mgr_b);
    let _ = &ds_b; // keep alive

    // Have both nodes discover each other via SWIM
    let addr_a = driver_a.endpoint_addr();
    let addr_b = driver_b.endpoint_addr();
    driver_a.join(&[addr_b]);
    driver_b.join(&[addr_a]);

    // Pump until SWIM membership converges
    let swim_deadline = Instant::now() + Duration::from_secs(10);
    loop {
        driver_a.recv();
        driver_a.tick();
        driver_b.recv();
        driver_b.tick();

        let snap_a = driver_a.snapshot();
        let snap_b = driver_b.snapshot();
        if snap_a.alive_count >= 1 && snap_b.alive_count >= 1 {
            break;
        }
        if Instant::now() >= swim_deadline {
            panic!("SWIM convergence timed out");
        }
        std::thread::sleep(Duration::from_millis(50));
    }

    // Give worker threads time to process spawned actors
    std::thread::sleep(Duration::from_millis(100));

    // Store a small blob on Node A
    let test_data = b"Hello from node A! This is a stream integration test.";
    let put_inbox = rt_a
        .new_inbox::<DatastoreResponse>()
        .expect("create inbox");
    let ds_a_addr = rt_a.where_is("Datastore").expect("Datastore registered on A");
    let _ = rt_a.send_to(
        ds_a_addr,
        DatastoreNodeMsg::Put {
            data: test_data.to_vec(),
            name: Some("test-blob".into()),
            tags: Default::default(),
            reply_to: *put_inbox.addr(),
        },
    );

    // Wait for PutOk (worker threads process messages)
    let content_hash = loop {
        if let Some(resp) = put_inbox.try_recv() {
            match resp {
                DatastoreResponse::PutOk { content_hash } => break content_hash,
                other => panic!("expected PutOk, got: {other:?}"),
            }
        }
        std::thread::sleep(Duration::from_millis(10));
    };

    // PutOk comes from MetadataActor; BlobStore writes are fire-and-forget.
    // Give BlobStore time to finish writing chunks + manifest.
    std::thread::sleep(Duration::from_millis(200));

    // Download the blob on Node B via stream
    let download_inbox = rt_b
        .new_inbox::<DatastoreResponse>()
        .expect("create inbox");
    let ds_b_addr = rt_b.where_is("Datastore").expect("Datastore registered on B");
    let _ = rt_b.send_to(
        ds_b_addr,
        DatastoreNodeMsg::DownloadViaStream {
            content_hash,
            source_node: driver_a.node_id().0,
            reply_to: *download_inbox.addr(),
        },
    );

    // Pump loop until we get a response
    let resp = pump_until_response(
        &rt_a,
        &rt_b,
        &mut driver_a,
        &mut driver_b,
        mgr_a,
        mgr_b,
        &download_inbox,
        Duration::from_secs(15),
    );

    match resp {
        Some(DatastoreResponse::PutOk { content_hash: h }) => {
            assert_eq!(h, content_hash, "downloaded blob hash should match");
        }
        other => panic!("expected PutOk from download, got: {other:?}"),
    }

    // Verify: read the blob back from Node B's datastore
    let verify_inbox = rt_b
        .new_inbox::<DatastoreResponse>()
        .expect("create inbox");
    let _ = rt_b.send_to(
        ds_b_addr,
        DatastoreNodeMsg::Get {
            content_hash,
            reply_to: *verify_inbox.addr(),
        },
    );

    let verify_deadline = Instant::now() + Duration::from_secs(5);
    loop {
        if let Some(resp) = verify_inbox.try_recv() {
            match resp {
                DatastoreResponse::GetOk { entry, manifest } => {
                    assert_eq!(entry.content_hash, content_hash);
                    assert_eq!(manifest.total_size, test_data.len() as u64);
                    break;
                }
                other => panic!("expected GetOk, got: {other:?}"),
            }
        }
        if Instant::now() >= verify_deadline {
            panic!("verify timed out — blob not found on Node B");
        }
        std::thread::sleep(Duration::from_millis(10));
    }

    driver_a.shutdown();
    driver_b.shutdown();
    rt_a.shutdown();
    rt_b.shutdown();
}

#[test]
fn multi_chunk_blob_transfers_between_two_nodes_via_stream() {
    let mut driver_a = make_driver_with_streams();
    let mut driver_b = make_driver_with_streams();

    let rt_a = make_runtime();
    let rt_b = make_runtime();

    let mgr_a = spawn_stream_manager(&rt_a, &driver_a);
    let mgr_b = spawn_stream_manager(&rt_b, &driver_b);

    let _ds_a = spawn_datastore(&rt_a, &driver_a, mgr_a);
    let ds_b = spawn_datastore(&rt_b, &driver_b, mgr_b);
    let _ = &ds_b;

    // Discover each other via SWIM
    let addr_a = driver_a.endpoint_addr();
    let addr_b = driver_b.endpoint_addr();
    driver_a.join(&[addr_b]);
    driver_b.join(&[addr_a]);

    let swim_deadline = Instant::now() + Duration::from_secs(10);
    loop {
        driver_a.recv();
        driver_a.tick();
        driver_b.recv();
        driver_b.tick();

        let snap_a = driver_a.snapshot();
        let snap_b = driver_b.snapshot();
        if snap_a.alive_count >= 1 && snap_b.alive_count >= 1 {
            break;
        }
        if Instant::now() >= swim_deadline {
            panic!("SWIM convergence timed out");
        }
        std::thread::sleep(Duration::from_millis(50));
    }

    // Give worker threads time to process spawned actors
    std::thread::sleep(Duration::from_millis(100));

    // 4 chunks at 256 bytes each = 1024 bytes
    let test_data: Vec<u8> = (0..1024).map(|i| (i % 251) as u8).collect();
    let put_inbox = rt_a
        .new_inbox::<DatastoreResponse>()
        .expect("create inbox");
    let ds_a_addr = rt_a.where_is("Datastore").expect("Datastore registered on A");
    let _ = rt_a.send_to(
        ds_a_addr,
        DatastoreNodeMsg::Put {
            data: test_data.clone(),
            name: Some("multi-chunk".into()),
            tags: Default::default(),
            reply_to: *put_inbox.addr(),
        },
    );

    let content_hash = loop {
        if let Some(resp) = put_inbox.try_recv() {
            match resp {
                DatastoreResponse::PutOk { content_hash } => break content_hash,
                other => panic!("expected PutOk, got: {other:?}"),
            }
        }
        std::thread::sleep(Duration::from_millis(10));
    };

    // PutOk comes from MetadataActor; BlobStore writes are fire-and-forget.
    std::thread::sleep(Duration::from_millis(200));

    // Download on Node B
    let download_inbox = rt_b
        .new_inbox::<DatastoreResponse>()
        .expect("create inbox");
    let ds_b_addr = rt_b.where_is("Datastore").expect("Datastore registered on B");
    let _ = rt_b.send_to(
        ds_b_addr,
        DatastoreNodeMsg::DownloadViaStream {
            content_hash,
            source_node: driver_a.node_id().0,
            reply_to: *download_inbox.addr(),
        },
    );

    let resp = pump_until_response(
        &rt_a,
        &rt_b,
        &mut driver_a,
        &mut driver_b,
        mgr_a,
        mgr_b,
        &download_inbox,
        Duration::from_secs(15),
    );

    match resp {
        Some(DatastoreResponse::PutOk { content_hash: h }) => {
            assert_eq!(h, content_hash);
        }
        other => panic!("expected PutOk from download, got: {other:?}"),
    }

    // Verify the data on Node B by reading each chunk
    let verify_inbox = rt_b
        .new_inbox::<DatastoreResponse>()
        .expect("create inbox");
    let _ = rt_b.send_to(
        ds_b_addr,
        DatastoreNodeMsg::Get {
            content_hash,
            reply_to: *verify_inbox.addr(),
        },
    );

    let verify_deadline = Instant::now() + Duration::from_secs(5);
    loop {
        if let Some(resp) = verify_inbox.try_recv() {
            match resp {
                DatastoreResponse::GetOk { entry, manifest } => {
                    assert_eq!(entry.content_hash, content_hash);
                    assert_eq!(manifest.total_size, test_data.len() as u64);
                    assert!(manifest.chunks.len() > 1, "should be multi-chunk");
                    break;
                }
                other => panic!("expected GetOk, got: {other:?}"),
            }
        }
        if Instant::now() >= verify_deadline {
            panic!("verify timed out — blob not found on Node B");
        }
        std::thread::sleep(Duration::from_millis(10));
    }

    driver_a.shutdown();
    driver_b.shutdown();
    rt_a.shutdown();
    rt_b.shutdown();
}