rust_distributed_id/worker/
worker_assigner.rs1use 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; const ACTUAL: i8 = 2; const ENV_KEY_HOST: &str = "JPAAS_HOST"; const ENV_KEY_PORT: &str = "JPAAS_HTTP_PORT"; const ENV_KEY_PORT_ORIGINAL: &str = "JPAAS_HOST_PORT_8080"; const COMPUTER_NAME: &str = "COMPUTERNAME"; pub struct Assigner {
18 worker_node_service: worker_service::WorkerService,
19 port: String,
20}
21
22impl Assigner {
23 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