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        413 | 429 => tonic::Code::ResourceExhausted,
77        _ => tonic::Code::Internal,
78    };
79    let mut s = Status::new(code, format!("{}: {}", e.code, e.message));
80    if let Ok(v) = e.code.parse() {
81        s.metadata_mut().insert("orbit-error-code", v);
82    }
83    s
84}
85
86fn api(e: ApiError) -> Status {
87    status(e.0)
88}
89
90// ------------------------------------------------------------- service
91
92pub struct GrpcWorkers {
93    state: Shared,
94}
95
96impl GrpcWorkers {
97    pub fn new(state: Shared) -> WorkerServiceServer<GrpcWorkers> {
98        WorkerServiceServer::new(GrpcWorkers { state })
99    }
100
101    /// Authenticate a call from its metadata and client certificate.
102    fn caller<T>(&self, req: &Request<T>) -> Result<Caller, Status> {
103        let bearer = req
104            .metadata()
105            .get("authorization")
106            .and_then(|v| v.to_str().ok())
107            .and_then(|v| v.strip_prefix("Bearer "));
108        let certs = req.peer_certs();
109        let der = certs
110            .as_ref()
111            .and_then(|c| c.first())
112            .map(|c| c.as_ref().to_vec());
113        match identify(&self.state, bearer, der.as_deref()).map_err(api)? {
114            // The gRPC binding is the worker protocol: applications use HTTP.
115            Some(Caller::Client(c)) => Err(status(EngineError::new(
116                403,
117                "FORBIDDEN",
118                format!("client `{}` may not use the worker protocol", c.id),
119            ))),
120            Some(c) => Ok(c),
121            None => Err(Status::unauthenticated(
122                "UNAUTHENTICATED: missing or invalid credentials",
123            )),
124        }
125    }
126
127    fn identity_check_register(caller: &Caller, req: &RegisterWorker) -> Result<(), Status> {
128        if let Caller::Worker(id) = caller {
129            let deny = |m: String| Err(status(EngineError::new(403, "FORBIDDEN", m)));
130            if req.worker_id != id.id {
131                return deny(format!(
132                    "this credential belongs to worker `{}`, not `{}`",
133                    id.id, req.worker_id
134                ));
135            }
136            if let Some(rt) = &id.runtime {
137                if !rt.eq_ignore_ascii_case(&req.runtime) {
138                    return deny(format!("worker `{}` must announce runtime `{rt}`", id.id));
139                }
140            }
141            let denied: Vec<&str> = req
142                .contracts
143                .iter()
144                .map(|c| c.task_name.as_str())
145                .filter(|t| !id.may_serve(t))
146                .collect();
147            if !denied.is_empty() {
148                return deny(format!(
149                    "worker `{}` may not serve: {} (allowed: {})",
150                    id.id,
151                    denied.join(", "),
152                    id.tasks.join(", ")
153                ));
154            }
155        }
156        Ok(())
157    }
158}
159
160#[tonic::async_trait]
161impl WorkerService for GrpcWorkers {
162    async fn register_worker(
163        &self,
164        req: Request<pb::RegisterWorkerRequest>,
165    ) -> Result<Response<pb::RegisterWorkerResponse>, Status> {
166        let caller = self.caller(&req)?;
167        let r = req.into_inner();
168        let reg = RegisterWorker {
169            worker_id: r.worker_id,
170            runtime: r.runtime,
171            sdk: (!r.sdk.is_empty()).then_some(r.sdk),
172            contracts: r
173                .contracts
174                .into_iter()
175                .map(|c| ContractAnnouncement {
176                    task_name: c.task_name,
177                    contract_version: c.contract_version,
178                })
179                .collect(),
180        };
181        Self::identity_check_register(&caller, &reg)?;
182        let out = engine(&self.state).register_worker(reg).map_err(status)?;
183        Ok(Response::new(pb::RegisterWorkerResponse {
184            accepted: out.accepted,
185            worker_session_id: out.worker_session_id,
186            heartbeat_interval_ms: out.heartbeat_interval_ms,
187            session_timeout_ms: out.session_timeout_ms,
188            warnings: out.warnings,
189        }))
190    }
191
192    async fn heartbeat(
193        &self,
194        req: Request<pb::HeartbeatRequest>,
195    ) -> Result<Response<pb::HeartbeatResponse>, Status> {
196        let caller = self.caller(&req)?;
197        let session = req.into_inner().worker_session_id;
198        authorize_session(&self.state, &caller, &session).map_err(api)?;
199        let ack = engine(&self.state).heartbeat(&session).map_err(status)?;
200        Ok(Response::new(pb::HeartbeatResponse {
201            accepted: ack.accepted,
202            cancel: ack.cancel,
203        }))
204    }
205
206    async fn poll_task(
207        &self,
208        req: Request<pb::PollTaskRequest>,
209    ) -> Result<Response<pb::ExecuteTaskRequest>, Status> {
210        let caller = self.caller(&req)?;
211        let r = req.into_inner();
212        authorize_session(&self.state, &caller, &r.worker_session_id).map_err(api)?;
213        let wait = Duration::from_millis(if r.wait_ms == 0 {
214            20_000
215        } else {
216            r.wait_ms.min(60_000)
217        });
218        let deadline = tokio::time::Instant::now() + wait;
219        loop {
220            let notified = self.state.changed.notified();
221            tokio::pin!(notified);
222            notified.as_mut().enable();
223            if let Some(t) =
224                crate::shared::poll(&self.state, &r.worker_session_id).map_err(status)?
225            {
226                return Ok(Response::new(pb::ExecuteTaskRequest {
227                    token: t.token,
228                    execution_id: t.execution_id,
229                    workflow: t.workflow,
230                    workflow_version: t.workflow_version,
231                    step_id: t.step_id,
232                    task_name: t.task_name,
233                    contract_version: t.contract_version,
234                    input: Some(from_json(t.input)),
235                    idempotency_key: t.idempotency_key,
236                    attempt: t.attempt,
237                    redelivery: t.redelivery,
238                    deadline_ms: t.deadline_ms,
239                    timeout_ms: t.timeout_ms,
240                }));
241            }
242            if tokio::time::Instant::now() >= deadline {
243                // No task: an empty token (HTTP answers 204).
244                return Ok(Response::new(pb::ExecuteTaskRequest::default()));
245            }
246            let _ = tokio::time::timeout_at(deadline, notified).await;
247        }
248    }
249
250    async fn complete_task(
251        &self,
252        req: Request<pb::CompleteTaskRequest>,
253    ) -> Result<Response<pb::TaskResultAck>, Status> {
254        let caller = self.caller(&req)?;
255        let r = req.into_inner();
256        if r.token.is_empty() {
257            return Err(Status::invalid_argument(
258                "INVALID_REQUEST: token is required",
259            ));
260        }
261        authorize_token(&self.state, &caller, &r.token).map_err(api)?;
262        let output = r.output.map(to_json).unwrap_or(Json::Null);
263        let ack = engine(&self.state)
264            .complete(&r.token, output)
265            .map_err(status)?;
266        Ok(Response::new(pb::TaskResultAck {
267            accepted: ack.accepted,
268            reason: ack.reason.unwrap_or_default(),
269        }))
270    }
271
272    async fn fail_task(
273        &self,
274        req: Request<pb::FailTaskRequest>,
275    ) -> Result<Response<pb::TaskResultAck>, Status> {
276        let caller = self.caller(&req)?;
277        let r = req.into_inner();
278        if r.error_code.is_empty() {
279            return Err(Status::invalid_argument(
280                "INVALID_REQUEST: error_code is required",
281            ));
282        }
283        authorize_token(&self.state, &caller, &r.token).map_err(api)?;
284        let error = TaskError {
285            code: r.error_code,
286            message: r.error_message,
287            retryable: r.retryable,
288            details: r.details.map(to_json).filter(|d| !d.is_null()),
289        };
290        let ack = engine(&self.state).fail(&r.token, error).map_err(status)?;
291        Ok(Response::new(pb::TaskResultAck {
292            accepted: ack.accepted,
293            reason: ack.reason.unwrap_or_default(),
294        }))
295    }
296
297    async fn unregister_worker(
298        &self,
299        req: Request<pb::UnregisterWorkerRequest>,
300    ) -> Result<Response<pb::UnregisterWorkerResponse>, Status> {
301        let caller = self.caller(&req)?;
302        let session = req.into_inner().worker_session_id;
303        authorize_session(&self.state, &caller, &session).map_err(api)?;
304        engine(&self.state)
305            .unregister_worker(&session)
306            .map_err(status)?;
307        Ok(Response::new(pb::UnregisterWorkerResponse {
308            accepted: true,
309        }))
310    }
311}
orbit_runtime/grpc.rs

orbit_runtime/
grpc.rs
