aboutsummaryrefslogtreecommitdiff
path: root/src/context.rs
blob: f67c3025cc0051953ab5cc664c69fe924ec28897 (plain)
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
use eyre::Result;

use crate::machine::Machine;

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
pub enum ExecutionNode {
    Frontend,
    Machine(Machine),
    Unknown,
}

#[derive(Debug, Clone)]
pub struct Context {
    pub node: ExecutionNode,
    job_id: Option<u32>,
    infer_job_id: bool,
    frontend_hostname: Option<String>,
}

impl Context {
    pub async fn new(
        job_id: Option<u32>,
        infer_job_id: bool,
        frontend_hostname: Option<String>,
    ) -> Result<Self> {
        Ok(Self {
            node: get_execution_node().await?,
            job_id,
            infer_job_id,
            frontend_hostname,
        })
    }

    pub async fn job_id(&self) -> Result<u32> {
        tracing::debug!("obtaining job id");
        if let Some(job_id) = self.job_id {
            tracing::debug!("job id was set, using {job_id}");
            Ok(job_id)
        } else if self.infer_job_id {
            tracing::debug!("job id was not set but inference is enabled, finding job id");
            let job_ids = crate::oar::list_user_job_ids(self).await?;
            match job_ids.len() {
                0 => Err(eyre::eyre!("cannot infer job id, no jobs are running")),
                1 => Ok(job_ids[0]),
                _ => Err(eyre::eyre!(
                    "cannot infer job id, multiple jobs are running"
                )),
            }
        } else {
            tracing::debug!("inference was disabled and job id is not set");
            Err(eyre::eyre!("missing job id"))
        }
    }

    pub fn frontend_hostname(&self) -> Result<&str> {
        self.frontend_hostname
            .as_deref()
            .ok_or_else(|| eyre::eyre!("missing frontend hostname"))
    }
}

async fn get_execution_node() -> Result<ExecutionNode> {
    let hostname = get_hostname().await;
    let node = match hostname.as_str() {
        "frontend" => ExecutionNode::Frontend,
        _ => match Machine::from_hostname(&hostname) {
            Some(machine) => ExecutionNode::Machine(machine),
            _ => ExecutionNode::Unknown,
        },
    };
    Ok(node)
}

async fn get_hostname() -> String {
    let hostname = if let Ok(hostname) = tokio::fs::read_to_string("/etc/hostname").await {
        hostname
    } else if let Ok(hostname) = std::env::var("HOSTNAME") {
        hostname
    } else {
        tracing::warn!("unable to obtain hostname, using empty string");
        String::default()
    };
    hostname.trim().to_string()
}