rust_distributed_id/worker/
worker_assigner.rs

1use chrono::NaiveDateTime;
2use rbatis::core::value::DateTimeNow;
3use std::env;
4use rbatis::rbatis::Rbatis;
5
6use crate::worker::service::worker_service;
7use crate::worker::model::worker_node;
8
9const CONTAINER: i8               = 1; // container type
10const ACTUAL: i8                  = 2; // actual type
11
12const ENV_KEY_HOST: &str          = "JPAAS_HOST";           // env key host
13const ENV_KEY_PORT: &str          = "JPAAS_HTTP_PORT";      // env key port
14const ENV_KEY_PORT_ORIGINAL: &str = "JPAAS_HOST_PORT_8080"; // env key 8080 port
15const COMPUTER_NAME: &str         = "COMPUTERNAME";         //computer name
16
17pub struct Assigner {
18    worker_node_service: worker_service::WorkerService,
19    port: String,
20}
21
22impl Assigner {
23    // new_worker_id_assigner create worker id assigner instance
24    pub fn new(port: String, rb: &'static Rbatis) -> Self {
25        let worker_node_service = worker_service::WorkerService::new(rb);
26        Assigner { worker_node_service, port }
27    }
28
29    pub async fn assign_worker_id(&self) -> i64 {
30        let new_node = Assigner::build_worker_node(&self.port);
31        let node = self.worker_node_service.get_by_hostname(
32            &new_node.host_name.clone().unwrap()
33        ).await;
34
35        if let Ok(res) = node {
36            return res.id.unwrap() as i64;
37        }
38
39        let result = self.worker_node_service.save(new_node).await;
40        return result.unwrap();
41    }
42
43    fn build_worker_node(port: &String) -> worker_node::WorkerNode {
44        let mut node = worker_node::WorkerNode {
45            id: None,
46            host_name: None,
47            port: Some(port.clone()),
48            worker_type: None,
49            modified: NaiveDateTime::now().into(),
50            created: NaiveDateTime::now().into()
51        };
52
53        let host_var = env::var(ENV_KEY_HOST);
54        if let Ok(res) = host_var {
55            node.host_name = Some(res);
56            node.worker_type = Some(CONTAINER);
57            node.port = Assigner::get_docker_port();
58        } else {
59            node.worker_type = Some(ACTUAL);
60            node.host_name = Assigner::get_computer_hostname();
61            node.port = Some(port.clone());
62        }
63
64        return node
65    }
66
67    fn get_docker_port()-> Option<String> {
68        let port = env::var(ENV_KEY_PORT);
69        if let Ok(res) = port {
70            return Some(res)
71        }
72
73        env::var(ENV_KEY_PORT_ORIGINAL).ok()
74    }
75
76    fn get_computer_hostname() -> Option<String> {
77        let hostname = env::var(COMPUTER_NAME);
78        if let Ok(res) = hostname {
79            return Some(res)
80        }
81
82        let result = hostname::get().unwrap();
83        return result.into_string().ok();
84    }
85}
86