frxd/tests/concurrency.rs

mod common;

use std::collections::BTreeSet;
use std::fs;
use std::path::Path;
use std::time::{Duration, Instant};

use common::{
    ask, client, collection, config_for, poll_messages, publish, query_envelope, register,
    spawn_relay,
};
use frxd::crypto::Keypair;
use frxd::index::LocalIndex;
use frxd::message::{EXPOSURE_FULL, TYPE_RESPONSE};
use frxd::node::{self, Node, NodeHandle};
use serde_json::Value;

fn corpus_dir(root: &Path, name: &str, files: &[(&str, &str)]) -> std::path::PathBuf {
    let dir = root.join(format!("{name}-docs"));
    fs::create_dir_all(&dir).unwrap();
    for (file, text) in files {
        fs::write(dir.join(file), text).unwrap();
    }
    dir
}

async fn start_node_with_corpus(
    root: &Path,
    name: &str,
    relays: Vec<String>,
    files: &[(&str, &str)],
) -> NodeHandle {
    let docs = corpus_dir(root, name, files);
    let mut config = config_for(&root.join(name), name, relays.first().unwrap());
    config.node.relays = relays;
    {
        let index = LocalIndex::open(&config.index_dir()).unwrap();
        index
            .add_collection(&collection("docs", &docs, true, EXPOSURE_FULL))
            .unwrap();
    }
    Node::start(config).await.unwrap()
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn concurrent_queries_each_get_their_response() {
    let root = tempfile::tempdir().unwrap();
    let relay_url = spawn_relay().await;
    let _bob = start_node_with_corpus(
        root.path(),
        "bob",
        vec![relay_url.clone()],
        &[("doc.txt", "concurrent rust document")],
    )
    .await;
    let alice = start_node_with_corpus(
        root.path(),
        "alice",
        vec![relay_url.clone()],
        &[("mine.txt", "alice notes")],
    )
    .await;

    let http = client();
    let addr = alice.addr.to_string();
    let mut tasks = Vec::new();
    for _ in 0..4 {
        let http = http.clone();
        let addr = addr.clone();
        tasks.push(tokio::spawn(
            async move { ask(&http, &addr, "rust", 5).await },
        ));
    }
    let mut qids = BTreeSet::new();
    for task in tasks {
        let (_status, _raw, value) = task.await.unwrap();
        let qid = value
            .get("qid")
            .and_then(Value::as_str)
            .unwrap()
            .to_string();
        qids.insert(qid);
        let responses = value.get("responses").and_then(Value::as_array).unwrap();
        assert_eq!(responses.len(), 1, "query lost its response: {value}");
    }
    assert_eq!(qids.len(), 4);
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn multi_relay_deduplicates_and_responds_once() {
    let root = tempfile::tempdir().unwrap();
    let relay_one = spawn_relay().await;
    let relay_two = spawn_relay().await;
    let _bob = start_node_with_corpus(
        root.path(),
        "bob",
        vec![relay_one.clone(), relay_two.clone()],
        &[("doc.txt", "dual relay rust document")],
    )
    .await;

    let http = client();
    let alice = Keypair::generate();
    register(&http, &relay_one, &alice.public_hex()).await;
    register(&http, &relay_two, &alice.public_hex()).await;
    let envelope = query_envelope(&alice, "rust", 5);
    assert!(
        publish(&http, &relay_one, &envelope)
            .await
            .status()
            .is_success()
    );
    assert!(
        publish(&http, &relay_two, &envelope)
            .await
            .status()
            .is_success()
    );

    let deadline = Instant::now() + Duration::from_millis(1200);
    let mut responses = Vec::new();
    while responses.is_empty() && Instant::now() < deadline {
        for relay in [&relay_one, &relay_two] {
            let messages = poll_messages(&http, relay, &alice.public_hex(), 200).await;
            responses.extend(
                messages
                    .iter()
                    .filter(|m| m.get("type").and_then(Value::as_str) == Some(TYPE_RESPONSE))
                    .cloned(),
            );
        }
    }
    assert_eq!(
        responses.len(),
        1,
        "duplicate relay delivery was answered twice"
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn index_persists_across_node_restart() {
    let root = tempfile::tempdir().unwrap();
    let docs = corpus_dir(
        root.path(),
        "bob",
        &[("doc.txt", "persistent rust document")],
    );
    let config = config_for(&root.path().join("bob"), "bob", "http://127.0.0.1:1");
    {
        let index = LocalIndex::open(&config.index_dir()).unwrap();
        index
            .add_collection(&collection("docs", &docs, true, EXPOSURE_FULL))
            .unwrap();
    }

    let first = Node::start(config.clone()).await.unwrap();
    let outcome = node::control_query(
        &format!("http://{}", first.addr),
        "rust",
        Some(5),
        Some(100),
        false,
    )
    .await
    .unwrap();
    assert_eq!(
        outcome.pointer("/local/total").and_then(Value::as_u64),
        Some(1)
    );
    drop(first);

    let second = Node::start(config).await.unwrap();
    let outcome = node::control_query(
        &format!("http://{}", second.addr),
        "rust",
        Some(5),
        Some(100),
        false,
    )
    .await
    .unwrap();
    assert_eq!(
        outcome.pointer("/local/total").and_then(Value::as_u64),
        Some(1)
    );
}

#[tokio::test(flavor = "multi_thread", worker_threads = 4)]
async fn poll_returns_queued_batch_in_one_call() {
    let relay_url = spawn_relay().await;
    let http = client();
    let member = Keypair::generate();
    register(&http, &relay_url, &member.public_hex()).await;
    for index in 0..3 {
        let envelope = query_envelope(&member, &format!("query {index}"), 5);
        assert!(
            publish(&http, &relay_url, &envelope)
                .await
                .status()
                .is_success()
        );
    }
    let messages = poll_messages(&http, &relay_url, &member.public_hex(), 300).await;
    assert_eq!(messages.len(), 3);
    assert!(
        poll_messages(&http, &relay_url, &member.public_hex(), 50)
            .await
            .is_empty()
    );
}