maelstrom-worker 0.14.0

The Maelstrom worker. This process executes jobs as directed by the broker.
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

use crate::{
    dispatcher::{self, Message},
    types::{DispatcherSender, TempFile, TempFileFactory},
};
use anyhow::{anyhow, Result};
use maelstrom_base::Sha256Digest;
use maelstrom_github::GitHubClient;
use maelstrom_util::{
    async_fs::Fs,
    cache::{fs::TempFile as _, GotArtifact},
    r#async::Pool,
};
use slog::{debug, o, warn, Logger};
use std::{num::NonZeroU32, sync::Arc};
use tokio::task;

pub struct GitHubArtifactFetcher {
    github_client: Arc<GitHubClient>,
    dispatcher_sender: DispatcherSender,
    log: Logger,
    pool: Arc<Pool<()>>,
    temp_file_factory: TempFileFactory,
}

impl GitHubArtifactFetcher {
    pub fn new(
        max_simultaneous_fetches: NonZeroU32,
        github_client: GitHubClient,
        dispatcher_sender: DispatcherSender,
        log: Logger,
        temp_file_factory: TempFileFactory,
    ) -> Self {
        let github_client = Arc::new(github_client);
        Self {
            github_client,
            dispatcher_sender,
            log,
            pool: Arc::new(Pool::new(max_simultaneous_fetches)),
            temp_file_factory,
        }
    }
}

impl dispatcher::ArtifactFetcher for GitHubArtifactFetcher {
    fn start_artifact_fetch(&mut self, digest: Sha256Digest) {
        let log = self.log.new(o!(
            "digest" => digest.to_string(),
        ));
        debug!(log, "artifact fetcher request enqueued");
        let github_client = self.github_client.clone();
        let dispatcher_sender = self.dispatcher_sender.clone();
        let pool = self.pool.clone();
        let temp_file_factory = self.temp_file_factory.clone();
        task::spawn(async move {
            let result = pool
                .call_with_item(|_: Option<()>| {
                    main(github_client, &digest, &log, temp_file_factory)
                })
                .await;
            debug!(log, "artifact fetcher request completed"; "result" => ?result);
            let _ = dispatcher_sender.send(Message::ArtifactFetchCompleted(
                digest,
                result.map(GotArtifact::file),
            ));
        });
    }
}

async fn main(
    github_client: Arc<GitHubClient>,
    digest: &Sha256Digest,
    log: &Logger,
    temp_file_factory: TempFileFactory,
) -> Result<((), TempFile)> {
    let temp_file = task::spawn_blocking(move || temp_file_factory.temp_file())
        .await
        .unwrap()
        .inspect_err(|err| {
            warn!(log, "artifact fetcher failed to create a temporary file"; "error" => %err);
        })?;

    let artifact_name = format!("maelstrom-cache-sha256-{digest}");
    let artifact = github_client
        .get(&artifact_name)
        .await?
        .ok_or_else(|| anyhow!("artifact {digest} not found"))?;
    let mut stream = github_client
        .download(artifact.backend_ids, &artifact_name)
        .await?;

    let fs = Fs::new();
    let mut file = fs.create_file(temp_file.path()).await?;
    tokio::io::copy(&mut stream, &mut file).await?;

    Ok(((), temp_file))
}