maelstrom-client-process 0.14.0

Client library background process for Maelstrom.
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

use crate::artifact_pusher::{construct_upload_name, start_task_inner, Receiver, SuccessCb};
use crate::progress::{ProgressTracker, UploadProgressReader};
use anyhow::{anyhow, Context as _, Result};
use maelstrom_base::{
    proto::{ArtifactPusherToBroker, BrokerToArtifactPusher, Hello},
    ArtifactUploadLocation, Sha256Digest,
};
use maelstrom_util::{async_fs::Fs, config::common::BrokerAddr, net, net::AsRawFdExt as _};
use slog::Logger;
use std::path::{Path, PathBuf};
use tokio::{
    io::{self, AsyncReadExt as _},
    net::TcpStream,
    task::JoinSet,
};

pub async fn push_one_artifact(
    stream: Option<TcpStream>,
    upload_tracker: ProgressTracker,
    broker_addr: BrokerAddr,
    path: PathBuf,
    digest: Sha256Digest,
    success_callback: SuccessCb,
    log: Logger,
) -> Result<(TcpStream, ())> {
    push_one_artifact_inner(
        stream,
        upload_tracker,
        broker_addr,
        &path,
        digest,
        success_callback,
        log,
    )
    .await
    .with_context(|| format!("pushing artifact {}", path.display()))
}

async fn push_one_artifact_inner(
    stream: Option<TcpStream>,
    upload_tracker: ProgressTracker,
    broker_addr: BrokerAddr,
    path: &Path,
    digest: Sha256Digest,
    success_callback: SuccessCb,
    log: Logger,
) -> Result<(TcpStream, ())> {
    let mut stream = match stream {
        Some(stream) => stream,
        None => {
            let mut stream = TcpStream::connect(broker_addr.inner())
                .await?
                .set_socket_options()?;
            net::write_message_to_async_socket(&mut stream, &Hello::ArtifactPusher, &log).await?;
            stream
        }
    };

    let fs = Fs::new();
    let file = fs.open_file(&path).await?;
    let size = file.metadata().await?.len();

    let upload_name = construct_upload_name(&digest, path);
    let prog = upload_tracker.new_task(&upload_name, size);

    let mut file = UploadProgressReader::new(prog, file.chain(io::repeat(0)).take(size));

    net::write_message_to_async_socket(&mut stream, &ArtifactPusherToBroker(digest, size), &log)
        .await?;
    let copied = io::copy(&mut file, &mut stream).await?;
    assert_eq!(copied, size);

    let BrokerToArtifactPusher(resp) =
        net::read_message_from_async_socket(&mut stream, &log).await?;

    resp.map_err(|e| anyhow!("Error from broker: {e}"))?;

    success_callback(ArtifactUploadLocation::TcpUpload);

    Ok((stream, ()))
}

pub fn start_task(
    join_set: &mut JoinSet<Result<()>>,
    receiver: Receiver,
    broker_addr: BrokerAddr,
    upload_tracker: ProgressTracker,
    log: Logger,
) {
    start_task_inner(
        join_set,
        receiver,
        move |stream, path, digest, success_callback| {
            push_one_artifact(
                stream,
                upload_tracker.clone(),
                broker_addr,
                path,
                digest,
                success_callback,
                log.clone(),
            )
        },
    )
}