Skip to main content

wowlab_node/core/
lifecycle.rs

1use std::{
2    sync::atomic::{AtomicBool, Ordering},
3    time::Duration,
4};
5
6use tokio::runtime::Handle;
7use wowlab_common::NodePublicKey;
8
9use super::{NodeCore, NodeCoreEvent, NodeEvents};
10use crate::{ConnectionStatus, NodeState, NodeStats, claim};
11
12const DRIVER_INTERVAL: Duration = Duration::from_millis(100);
13
14impl NodeCore {
15    /// Start async tasks. Call once after creation.
16    pub fn start(&mut self) {
17        if self.started {
18            return;
19        }
20
21        self.started = true;
22
23        self.worker_pool.start(self.runtime.handle());
24        self.result_rx = self.worker_pool.result_rx();
25
26        match self.state {
27            NodeState::Verifying => self.start_verification(),
28            NodeState::Registering => self.start_registration(),
29            NodeState::Running => {
30                if self.registered {
31                    self.start_realtime();
32                }
33            }
34            NodeState::Setup
35            | NodeState::NotFound
36            | NodeState::Unavailable
37            | NodeState::Unlinking
38            | NodeState::Unlinked => {}
39        }
40    }
41
42    /// Set claim token and start registration.
43    pub fn set_token_claim(&mut self, token: String) {
44        if let Err(msg) = claim::validate_token(&token) {
45            tracing::error!(error = %msg, "Invalid claim token");
46
47            return;
48        }
49
50        self.config.token_claim = Some(wowlab_types::sensitive::Sensitive::new(token));
51        self.set_state(NodeState::Registering);
52        self.start_registration();
53    }
54
55    /// Poll for updates. Call periodically.
56    pub fn poll(&mut self) {
57        self.check_unlink();
58        self.check_verification();
59        self.check_registration();
60        self.check_realtime_events();
61        self.check_work_results();
62        self.check_retry();
63    }
64
65    /// Drive the node and its events on the current thread until `running` becomes false.
66    pub fn drive_until_stopped(
67        &mut self,
68        events: &mut NodeEvents,
69        running: &AtomicBool,
70        mut on_event: impl FnMut(&NodeCoreEvent, &Self),
71    ) {
72        self.start();
73
74        while running.load(Ordering::SeqCst) {
75            self.poll();
76
77            while let Some(event) = events.try_recv() {
78                on_event(&event, self);
79                std::thread::yield_now();
80            }
81
82            std::thread::park_timeout(DRIVER_INTERVAL);
83            std::thread::yield_now();
84        }
85    }
86
87    pub fn state(&self) -> &NodeState {
88        &self.state
89    }
90
91    pub fn connection_status(&self) -> ConnectionStatus {
92        self.connection_status
93    }
94
95    pub fn public_key(&self) -> &NodePublicKey {
96        &self.config.public_key
97    }
98
99    pub fn is_registered(&self) -> bool {
100        self.registered
101    }
102
103    pub fn node_name(&self) -> &str {
104        &self.node_name
105    }
106
107    pub fn stats(&self) -> NodeStats {
108        let mut stats = self.worker_pool.stats();
109
110        stats.total_cores = self.total_cores;
111
112        stats
113    }
114
115    pub fn time_until_retry(&self) -> Option<Duration> {
116        self.backoff.time_until_retry()
117    }
118
119    pub fn disconnect(&mut self) {
120        if let Some(token) = self.realtime_shutdown.take() {
121            token.cancel();
122        }
123
124        self.realtime_rx = None;
125        self.set_connection(ConnectionStatus::Disconnected);
126    }
127
128    pub fn reconnect(&mut self) {
129        if matches!(self.state, NodeState::Running) && self.registered {
130            self.start_realtime();
131        }
132    }
133
134    pub fn runtime_handle(&self) -> &Handle {
135        self.runtime.handle()
136    }
137
138    pub(super) fn set_state(&mut self, state: NodeState) {
139        self.state = state;
140        let _ = self
141            .event_tx
142            .try_send(NodeCoreEvent::StateChanged(self.state.clone()));
143    }
144
145    pub(super) fn set_connection(&mut self, status: ConnectionStatus) {
146        self.connection_status = status;
147        let _ = self
148            .event_tx
149            .try_send(NodeCoreEvent::ConnectionChanged(self.connection_status));
150    }
151}