wowlab_node/
work_context.rs

1//! Cached, resolved work context shared across the work items of a chunk.
2
3use std::sync::Arc;
4
5use moka::sync::Cache;
6use uuid::Uuid;
7use wowlab_common::{
8    ClaimToken, WorkContextHash,
9    node_http::NodeWorkContextResponse,
10    sim::intent::{IntentConfigError, SimConfigIntent, parse_sim_config},
11};
12use wowlab_types::proto::TournamentPayload;
13
14use crate::sentinel::SentinelClient;
15
16const MAX_CONTEXTS: u64 = 256;
17
18/// Immutable, resolved work context for one `(job_id, work_context_hash)`.
19#[derive(Clone, Debug)]
20pub struct WorkContext {
21    base_intent: Arc<SimConfigIntent>,
22    payload: Arc<TournamentPayload>,
23    sentinel_config: Arc<str>,
24}
25
26impl WorkContext {
27    pub fn new(
28        base_intent: SimConfigIntent,
29        payload: TournamentPayload,
30        sentinel_config: impl Into<String>,
31    ) -> Self {
32        let sentinel_config = sentinel_config.into();
33
34        Self {
35            base_intent: Arc::new(base_intent),
36            payload: Arc::new(payload),
37            sentinel_config: Arc::from(sentinel_config),
38        }
39    }
40
41    #[must_use]
42    pub fn base_intent(&self) -> &SimConfigIntent {
43        &self.base_intent
44    }
45
46    #[must_use]
47    pub fn payload(&self) -> &TournamentPayload {
48        &self.payload
49    }
50
51    #[must_use]
52    pub fn sentinel_config(&self) -> &str {
53        &self.sentinel_config
54    }
55}
56
57#[derive(thiserror::Error)]
58#[error("{kind}")]
59#[non_exhaustive]
60pub(crate) struct WorkContextError {
61    #[source]
62    kind: WorkContextErrorKind,
63}
64
65#[derive(Debug, thiserror::Error)]
66enum WorkContextErrorKind {
67    #[error("fetch failed: {0}")]
68    Fetch(#[source] crate::sentinel::SentinelError),
69    #[error("base sim config parse failed: {0}")]
70    ParseConfig(#[source] IntentConfigError),
71    #[error("tournament payload base64 decode failed: {0}")]
72    PayloadBase64(#[source] base64::DecodeError),
73    #[error("tournament payload protobuf decode failed: {0}")]
74    PayloadProto(#[source] prost::DecodeError),
75    #[error("work context response has invalid job id: {0}")]
76    ResponseJobId(#[source] uuid::Error),
77    #[error("work context response job id does not match the request")]
78    JobIdMismatch,
79    #[error("work context response hash does not match the request")]
80    ResponseHashMismatch,
81    #[error("work context payload does not match its hash")]
82    PayloadHashMismatch,
83}
84
85impl std::fmt::Debug for WorkContextError {
86    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
87        let category = match self.kind {
88            WorkContextErrorKind::Fetch(_) => "Fetch",
89            WorkContextErrorKind::ParseConfig(_) => "ParseConfig",
90            WorkContextErrorKind::PayloadBase64(_) => "PayloadBase64",
91            WorkContextErrorKind::PayloadProto(_) => "PayloadProto",
92            WorkContextErrorKind::ResponseJobId(_) => "ResponseJobId",
93            WorkContextErrorKind::JobIdMismatch => "JobIdMismatch",
94            WorkContextErrorKind::ResponseHashMismatch => "ResponseHashMismatch",
95            WorkContextErrorKind::PayloadHashMismatch => "PayloadHashMismatch",
96        };
97
98        f.debug_struct("WorkContextError")
99            .field("category", &category)
100            .finish_non_exhaustive()
101    }
102}
103
104impl From<crate::sentinel::SentinelError> for WorkContextError {
105    fn from(source: crate::sentinel::SentinelError) -> Self {
106        Self {
107            kind: WorkContextErrorKind::Fetch(source),
108        }
109    }
110}
111
112impl From<IntentConfigError> for WorkContextError {
113    fn from(source: IntentConfigError) -> Self {
114        Self {
115            kind: WorkContextErrorKind::ParseConfig(source),
116        }
117    }
118}
119
120impl From<base64::DecodeError> for WorkContextError {
121    fn from(source: base64::DecodeError) -> Self {
122        Self {
123            kind: WorkContextErrorKind::PayloadBase64(source),
124        }
125    }
126}
127
128impl From<prost::DecodeError> for WorkContextError {
129    fn from(source: prost::DecodeError) -> Self {
130        Self {
131            kind: WorkContextErrorKind::PayloadProto(source),
132        }
133    }
134}
135
136impl From<uuid::Error> for WorkContextError {
137    fn from(source: uuid::Error) -> Self {
138        Self {
139            kind: WorkContextErrorKind::ResponseJobId(source),
140        }
141    }
142}
143
144#[derive(Clone)]
145pub(crate) struct WorkContextCache {
146    contexts: Cache<(Uuid, WorkContextHash), WorkContext>,
147}
148
149impl WorkContextCache {
150    pub(crate) fn new() -> Self {
151        Self {
152            contexts: Cache::builder().max_capacity(MAX_CONTEXTS).build(),
153        }
154    }
155
156    pub(crate) fn get(&self, job_id: Uuid, hash: WorkContextHash) -> Option<WorkContext> {
157        self.contexts.get(&(job_id, hash))
158    }
159
160    pub(crate) fn insert(&self, job_id: Uuid, hash: WorkContextHash, context: WorkContext) {
161        self.contexts.insert((job_id, hash), context);
162    }
163
164    pub(crate) async fn get_or_fetch(
165        &self,
166        sentinel: &SentinelClient,
167        job_id: Uuid,
168        hash: WorkContextHash,
169        claim_token: &ClaimToken,
170    ) -> Result<WorkContext, WorkContextError> {
171        if let Some(context) = self.get(job_id, hash) {
172            return Ok(context);
173        }
174
175        let response = sentinel
176            .fetch_work_context(job_id, &hash.to_hex(), claim_token.as_str())
177            .await?;
178        let context = build_verified_context(&response, job_id, hash)?;
179
180        self.insert(job_id, hash, context.clone());
181
182        Ok(context)
183    }
184}
185
186impl std::fmt::Debug for WorkContextCache {
187    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
188        f.debug_struct("WorkContextCache")
189            .field("entry_count", &self.contexts.entry_count())
190            .finish()
191    }
192}
193
194impl Default for WorkContextCache {
195    fn default() -> Self {
196        Self::new()
197    }
198}
199
200#[cfg(test)]
201pub(crate) fn build_context(
202    response: &NodeWorkContextResponse,
203) -> Result<WorkContext, WorkContextError> {
204    use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
205
206    let payload_bytes = if response.tournament_payload_bytes.is_empty() {
207        Vec::new()
208    } else {
209        BASE64.decode(&response.tournament_payload_bytes)?
210    };
211
212    build_context_from_payload(response, &payload_bytes)
213}
214
215fn build_verified_context(
216    response: &NodeWorkContextResponse,
217    expected_job_id: Uuid,
218    expected_hash: WorkContextHash,
219) -> Result<WorkContext, WorkContextError> {
220    use base64::{Engine, engine::general_purpose::STANDARD as BASE64};
221
222    if response.job_id.parse::<Uuid>()? != expected_job_id {
223        return Err(WorkContextError {
224            kind: WorkContextErrorKind::JobIdMismatch,
225        });
226    }
227
228    if response.work_context_hash != expected_hash {
229        return Err(WorkContextError {
230            kind: WorkContextErrorKind::ResponseHashMismatch,
231        });
232    }
233
234    let payload_bytes = if response.tournament_payload_bytes.is_empty() {
235        Vec::new()
236    } else {
237        BASE64.decode(&response.tournament_payload_bytes)?
238    };
239    let computed_hash = WorkContextHash::compute(
240        &response.base_sim_config,
241        &payload_bytes,
242        &response.sentinel_config,
243    );
244
245    if computed_hash != expected_hash {
246        return Err(WorkContextError {
247            kind: WorkContextErrorKind::PayloadHashMismatch,
248        });
249    }
250
251    build_context_from_payload(response, &payload_bytes)
252}
253
254fn build_context_from_payload(
255    response: &NodeWorkContextResponse,
256    payload_bytes: &[u8],
257) -> Result<WorkContext, WorkContextError> {
258    use prost::Message;
259
260    let base_intent = parse_sim_config(&response.base_sim_config)?;
261    let payload = if payload_bytes.is_empty() {
262        TournamentPayload::default()
263    } else {
264        TournamentPayload::decode(payload_bytes)?
265    };
266
267    Ok(WorkContext::new(
268        base_intent,
269        payload,
270        response.sentinel_config.clone(),
271    ))
272}
273
274#[cfg(test)]
275mod tests;
wowlab_node/work_context.rs

wowlab_node/
work_context.rs
