orbit_runtime/
grpc.rs

1//! gRPC binding of the worker protocol (`orbit.worker.v1.WorkerService`).
2//!
3//! It shares the engine, the credentials, per-worker identities and the
4//! authorization rules with the HTTP API, so the two bindings behave the
5//! same. Errors carry the ORBIT error code in the `orbit-error-code`
6//! metadata entry and as a `CODE: message` prefix.
7
8use crate::engine::{ContractAnnouncement, EngineError, RegisterWorker};
9use crate::model::TaskError;
10use crate::server::{
11    authorize_session, authorize_token, engine, identify, ApiError, Caller, Shared,
12};
13use serde_json::{Map, Number, Value as Json};
14use std::time::Duration;
15use tonic::{Request, Response, Status};
16
17/// Generated protocol types and the client/server stubs.
18pub mod pb {
19    tonic::include_proto!("orbit.worker.v1");
20}
21
22use pb::worker_service_server::{WorkerService, WorkerServiceServer};
23
24// ------------------------------------------------------- value conversion
25
26/// `google.protobuf.Value` -> JSON. Whole numbers within ±2^53 become
27/// integers, since protobuf carries every number as a double.
28pub fn to_json(v: prost_types::Value) -> Json {
29    use prost_types::value::Kind;
30    match v.kind {
31        None | Some(Kind::NullValue(_)) => Json::Null,
32        Some(Kind::BoolValue(b)) => Json::Bool(b),
33        Some(Kind::NumberValue(n)) => {
34            if n.fract() == 0.0 && n.abs() <= 9_007_199_254_740_992.0 {
35                Json::from(n as i64)
36            } else {
37                Number::from_f64(n).map(Json::Number).unwrap_or(Json::Null)
38            }
39        }
40        Some(Kind::StringValue(s)) => Json::String(s),
41        Some(Kind::ListValue(l)) => Json::Array(l.values.into_iter().map(to_json).collect()),
42        Some(Kind::StructValue(s)) => Json::Object(
43            s.fields
44                .into_iter()
45                .map(|(k, v)| (k, to_json(v)))
46                .collect::<Map<_, _>>(),
47        ),
48    }
49}
50
51/// JSON -> `google.protobuf.Value`.
52pub fn from_json(v: Json) -> prost_types::Value {
53    use prost_types::value::Kind;
54    let kind = match v {
55        Json::Null => Kind::NullValue(0),
56        Json::Bool(b) => Kind::BoolValue(b),
57        Json::Number(n) => Kind::NumberValue(n.as_f64().unwrap_or(0.0)),
58        Json::String(s) => Kind::StringValue(s),
59        Json::Array(a) => Kind::ListValue(prost_types::ListValue {
60            values: a.into_iter().map(from_json).collect(),
61        }),
62        Json::Object(o) => Kind::StructValue(prost_types::Struct {
63            fields: o.into_iter().map(|(k, v)| (k, from_json(v))).collect(),
64        }),
65    };
66    prost_types::Value { kind: Some(kind) }
67}
68
69fn status(e: EngineError) -> Status {
70    let code = match e.status {
71        400 => tonic::Code::InvalidArgument,
72        401 => tonic::Code::Unauthenticated,
73        403 => tonic::Code::PermissionDenied,
74        404 => tonic::Code::NotFound,
75        409 => tonic::Code::FailedPrecondition,
76        _ => tonic::Code::Internal,
77    };
78    let mut s = Status::new(code, format!("{}: {}", e.code, e.message));
79    if let Ok(v) = e.code.parse() {
80        s.metadata_mut().insert("orbit-error-code", v);
81    }
82    s
83}
84
85fn api(e: ApiError) -> Status {
86    status(e.0)
87}
88
89// ------------------------------------------------------------- service
90
91pub struct GrpcWorkers {
92    state: Shared,
93}
94
95impl GrpcWorkers {
96    pub fn new(state: Shared) -> WorkerServiceServer<GrpcWorkers> {
97        WorkerServiceServer::new(GrpcWorkers { state })
98    }
99
100    /// Authenticate a call from its metadata and client certificate.
101    fn caller<T>(&self, req: &Request<T>) -> Result<Caller, Status> {
102        let bearer = req
103            .metadata()
104            .get("authorization")
105            .and_then(|v| v.to_str().ok())
106            .and_then(|v| v.strip_prefix("Bearer "));
107        let certs = req.peer_certs();
108        let der = certs
109            .as_ref()
110            .and_then(|c| c.first())
111            .map(|c| c.as_ref().to_vec());
112        match identify(&self.state, bearer, der.as_deref()).map_err(api)? {
113            Some(c) => Ok(c),
114            None => Err(Status::unauthenticated(
115                "UNAUTHENTICATED: missing or invalid credentials",
116            )),
117        }
118    }
119
120    fn identity_check_register(caller: &Caller, req: &RegisterWorker) -> Result<(), Status> {
121        if let Caller::Worker(id) = caller {
122            let deny = |m: String| Err(status(EngineError::new(403, "FORBIDDEN", m)));
123            if req.worker_id != id.id {
124                return deny(format!(
125                    "this credential belongs to worker `{}`, not `{}`",
126                    id.id, req.worker_id
127                ));
128            }
129            if let Some(rt) = &id.runtime {
130                if !rt.eq_ignore_ascii_case(&req.runtime) {
131                    return deny(format!("worker `{}` must announce runtime `{rt}`", id.id));
132                }
133            }
134            let denied: Vec<&str> = req
135                .contracts
136                .iter()
137                .map(|c| c.task_name.as_str())
138                .filter(|t| !id.may_serve(t))
139                .collect();
140            if !denied.is_empty() {
141                return deny(format!(
142                    "worker `{}` may not serve: {} (allowed: {})",
143                    id.id,
144                    denied.join(", "),
145                    id.tasks.join(", ")
146                ));
147            }
148        }
149        Ok(())
150    }
151}
152
153#[tonic::async_trait]
154impl WorkerService for GrpcWorkers {
155    async fn register_worker(
156        &self,
157        req: Request<pb::RegisterWorkerRequest>,
158    ) -> Result<Response<pb::RegisterWorkerResponse>, Status> {
159        let caller = self.caller(&req)?;
160        let r = req.into_inner();
161        let reg = RegisterWorker {
162            worker_id: r.worker_id,
163            runtime: r.runtime,
164            sdk: (!r.sdk.is_empty()).then_some(r.sdk),
165            contracts: r
166                .contracts
167                .into_iter()
168                .map(|c| ContractAnnouncement {
169                    task_name: c.task_name,
170                    contract_version: c.contract_version,
171                })
172                .collect(),
173        };
174        Self::identity_check_register(&caller, &reg)?;
175        let out = engine(&self.state).register_worker(reg).map_err(status)?;
176        Ok(Response::new(pb::RegisterWorkerResponse {
177            accepted: out.accepted,
178            worker_session_id: out.worker_session_id,
179            heartbeat_interval_ms: out.heartbeat_interval_ms,
180            session_timeout_ms: out.session_timeout_ms,
181            warnings: out.warnings,
182        }))
183    }
184
185    async fn heartbeat(
186        &self,
187        req: Request<pb::HeartbeatRequest>,
188    ) -> Result<Response<pb::HeartbeatResponse>, Status> {
189        let caller = self.caller(&req)?;
190        let session = req.into_inner().worker_session_id;
191        authorize_session(&self.state, &caller, &session).map_err(api)?;
192        let ack = engine(&self.state).heartbeat(&session).map_err(status)?;
193        Ok(Response::new(pb::HeartbeatResponse {
194            accepted: ack.accepted,
195            cancel: ack.cancel,
196        }))
197    }
198
199    async fn poll_task(
200        &self,
201        req: Request<pb::PollTaskRequest>,
202    ) -> Result<Response<pb::ExecuteTaskRequest>, Status> {
203        let caller = self.caller(&req)?;
204        let r = req.into_inner();
205        authorize_session(&self.state, &caller, &r.worker_session_id).map_err(api)?;
206        let wait = Duration::from_millis(if r.wait_ms == 0 {
207            20_000
208        } else {
209            r.wait_ms.min(60_000)
210        });
211        let deadline = tokio::time::Instant::now() + wait;
212        loop {
213            let notified = self.state.changed.notified();
214            tokio::pin!(notified);
215            notified.as_mut().enable();
216            if let Some(t) = engine(&self.state)
217                .poll(&r.worker_session_id)
218                .map_err(status)?
219            {
220                return Ok(Response::new(pb::ExecuteTaskRequest {
221                    token: t.token,
222                    execution_id: t.execution_id,
223                    workflow: t.workflow,
224                    workflow_version: t.workflow_version,
225                    step_id: t.step_id,
226                    task_name: t.task_name,
227                    contract_version: t.contract_version,
228                    input: Some(from_json(t.input)),
229                    idempotency_key: t.idempotency_key,
230                    attempt: t.attempt,
231                    redelivery: t.redelivery,
232                    deadline_ms: t.deadline_ms,
233                    timeout_ms: t.timeout_ms,
234                }));
235            }
236            if tokio::time::Instant::now() >= deadline {
237                // No task: an empty token (HTTP answers 204).
238                return Ok(Response::new(pb::ExecuteTaskRequest::default()));
239            }
240            let _ = tokio::time::timeout_at(deadline, notified).await;
241        }
242    }
243
244    async fn complete_task(
245        &self,
246        req: Request<pb::CompleteTaskRequest>,
247    ) -> Result<Response<pb::TaskResultAck>, Status> {
248        let caller = self.caller(&req)?;
249        let r = req.into_inner();
250        if r.token.is_empty() {
251            return Err(Status::invalid_argument(
252                "INVALID_REQUEST: token is required",
253            ));
254        }
255        authorize_token(&self.state, &caller, &r.token).map_err(api)?;
256        let output = r.output.map(to_json).unwrap_or(Json::Null);
257        let ack = engine(&self.state)
258            .complete(&r.token, output)
259            .map_err(status)?;
260        Ok(Response::new(pb::TaskResultAck {
261            accepted: ack.accepted,
262            reason: ack.reason.unwrap_or_default(),
263        }))
264    }
265
266    async fn fail_task(
267        &self,
268        req: Request<pb::FailTaskRequest>,
269    ) -> Result<Response<pb::TaskResultAck>, Status> {
270        let caller = self.caller(&req)?;
271        let r = req.into_inner();
272        if r.error_code.is_empty() {
273            return Err(Status::invalid_argument(
274                "INVALID_REQUEST: error_code is required",
275            ));
276        }
277        authorize_token(&self.state, &caller, &r.token).map_err(api)?;
278        let error = TaskError {
279            code: r.error_code,
280            message: r.error_message,
281            retryable: r.retryable,
282            details: r.details.map(to_json).filter(|d| !d.is_null()),
283        };
284        let ack = engine(&self.state).fail(&r.token, error).map_err(status)?;
285        Ok(Response::new(pb::TaskResultAck {
286            accepted: ack.accepted,
287            reason: ack.reason.unwrap_or_default(),
288        }))
289    }
290
291    async fn unregister_worker(
292        &self,
293        req: Request<pb::UnregisterWorkerRequest>,
294    ) -> Result<Response<pb::UnregisterWorkerResponse>, Status> {
295        let caller = self.caller(&req)?;
296        let session = req.into_inner().worker_session_id;
297        authorize_session(&self.state, &caller, &session).map_err(api)?;
298        engine(&self.state)
299            .unregister_worker(&session)
300            .map_err(status)?;
301        Ok(Response::new(pb::UnregisterWorkerResponse {
302            accepted: true,
303        }))
304    }
305}
orbit_runtime/grpc.rs

orbit_runtime/
grpc.rs
