stigmerge_peer 0.6.2

stigmerge p2p protocol and agent components
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130

//! Example: fetch a single block from a share
#![recursion_limit = "256"]

use std::path::PathBuf;
use std::sync::Arc;

use clap::Parser;
use stigmerge_peer::new_connection;
use tokio::task::JoinSet;
use tokio_util::sync::CancellationToken;
use tracing::info;
use veilid_core::RecordKey;

use stigmerge_peer::block_fetcher::BlockFetcher;
use stigmerge_peer::share_resolver::ShareResolver;
use stigmerge_peer::types::FileBlockFetch;
use stigmerge_peer::Error;
use stigmerge_peer::Retry;
use veilnet::Connection;

/// Fetch block CLI arguments
#[derive(Parser, Debug)]
#[command(author, version, about, long_about = None)]
struct Args {
    /// The share key to fetch from
    #[arg(help = "The share key to fetch from")]
    share_key: String,

    /// File index in the share
    #[arg(long, default_value = "0", help = "File index in the share")]
    file_index: usize,

    /// Piece index in the file
    #[arg(long, default_value = "0", help = "Piece index in the file")]
    piece_index: usize,

    /// Block index in the piece
    #[arg(long, default_value = "0", help = "Block index in the piece")]
    block_index: usize,

    /// Piece offset
    #[arg(long, default_value = "0", help = "Piece offset")]
    piece_offset: usize,

    /// Output file path
    #[arg(short, long, help = "Output file path")]
    output: Option<PathBuf>,
}

#[tokio::main]
async fn main() -> std::result::Result<(), Error> {
    tracing_subscriber::fmt::init();

    // Parse command line arguments
    let args = Args::parse();
    let share_key = args.share_key;
    let file_index = args.file_index;
    let piece_index = args.piece_index;
    let block_index = args.block_index;
    let piece_offset = args.piece_offset;

    // Determine output file path
    let output_path = args.output.unwrap_or_else(|| {
        PathBuf::from(format!(
            "block_f{}_p{}_b{}.bin",
            file_index, piece_index, block_index
        ))
    });

    // Set up Veilid peer
    let state_dir = tempfile::tempdir()?;
    let conn = new_connection(state_dir.path().to_str().unwrap(), None).await?;
    let cancel = CancellationToken::new();
    let retry = Retry::default();

    let mut tasks = JoinSet::new();

    // Set up share resolver
    let (share_resolver, resolver_task) = ShareResolver::new_task(
        cancel.clone(),
        retry.clone(),
        conn.clone(),
        &output_path.clone(),
    );
    tasks.spawn(async move {
        resolver_task.await??;
        Ok(())
    });

    // Parse share key
    let key: RecordKey = share_key.parse()?;

    // Resolve the header and index
    let remote_share = Arc::new(share_resolver.add_share(&key).await?);

    info!(?key, "resolved remote share");

    // Create the block fetcher
    let mut block_fetcher = BlockFetcher::new(conn.clone(), output_path.clone());

    // Create the block fetch request
    let block = FileBlockFetch {
        file_index,
        piece_index,
        block_index,
        piece_offset,
    };

    info!(
        "Fetching block: file_index={}, piece_index={}, block_index={}, piece_offset={}",
        file_index, piece_index, block_index, piece_offset
    );

    let (_, length) = block_fetcher
        .fetch_block(remote_share, &block, true)
        .await?;
    info!("Successfully fetched block of length: {} bytes", length);

    info!("Block written to: {}", output_path.display());

    // Clean up
    cancel.cancel();
    tasks.spawn(async { conn.close().await.map_err(|e| e.into()) });
    tasks
        .join_all()
        .await
        .into_iter()
        .collect::<anyhow::Result<()>>()?;
    Ok(())
}