dex_sdk/
worker.rs

1// Copyright (c) 2026 Super Durable, Inc.
2//
3// Licensed under the Super Durable Source License 1.0.
4// You may not use this file except in compliance with the License.
5// See the LICENSE file in the repository root.
6//
7// SPDX-License-Identifier: LicenseRef-Super-Durable-1.0
8
9use std::collections::HashMap;
10use std::net::SocketAddr;
11use std::sync::Arc;
12use std::sync::atomic::{AtomicU8, Ordering};
13
14use dex_protocol::dex::WorkerErrorResponse;
15use dex_protocol::dex::flow_service_client::FlowServiceClient;
16use dex_protocol::dex::worker_service_server::{WorkerService, WorkerServiceServer};
17use dex_protocol::dex::{
18    InvokeExecuteMethodRequest, InvokeExecuteMethodResponse, InvokeWaitForMethodRequest,
19    InvokeWaitForMethodResponse, InvokeWorkerRpcRequest, InvokeWorkerRpcResponse,
20    SyncAttributeIndexRequest,
21};
22use prost::Message;
23use prost_types::Any;
24use tokio::runtime::Runtime;
25use tokio::sync::watch;
26use tonic::transport::{Channel, Endpoint, Server};
27use tonic::{Code, Request, Response, Status};
28
29use crate::value_hydrator::ValueHydrator;
30use crate::worker_dispatcher::WorkerDispatcher;
31use crate::{BlobCache, HandlerError, Registry, SdkError, SdkResult, WorkerOptions, WorkerTarget};
32
33// Keep worker error status small enough for default gRPC trailer limits after
34// the server wraps WorkerErrorResponse in ServiceErrorResponse.
35const MAX_WORKER_STACK_TRACE_BYTES: usize = 4 * 1024;
36const STACK_TRACE_TRUNCATION_MARKER: &str = "\n... stack trace truncated by Dex Rust SDK ...";
37
38const CREATED: u8 = 0;
39const RUNNING: u8 = 1;
40const STOPPED: u8 = 2;
41
42/// Hosts registered Flow Step and RPC handlers over application WorkerService.
43///
44/// A Worker owns an internal Tokio runtime and is one-shot. [`Self::start`] first synchronizes
45/// Attribute indexes with Dex, then blocks the calling thread while serving gRPC. Call [`Self::stop`]
46/// from another thread to request graceful server shutdown.
47///
48/// # Examples
49///
50/// ```no_run
51/// use dex_sdk::{BlobCache, BlobCacheConfig, Registry, Worker, WorkerOptions};
52/// use std::sync::Arc;
53///
54/// let cache = Arc::new(BlobCache::open(BlobCacheConfig::new(
55///     "/tmp/dex-blobs", 64 * 1024 * 1024, 0,
56/// )?)?);
57/// let worker = Worker::try_new(Registry::new(), cache, WorkerOptions::new())?;
58/// println!("advertise {}", worker.worker_target().address());
59/// worker.start()?;
60/// # Ok::<(), Box<dyn std::error::Error>>(())
61/// ```
62pub struct Worker {
63    runtime: Runtime,
64    service: RustWorkerService,
65    flow_service: FlowServiceClient<Channel>,
66    attribute_indexes: HashMap<String, i32>,
67    attribute_index_sync_timeout: std::time::Duration,
68    bind_address: SocketAddr,
69    worker_target: WorkerTarget,
70    shutdown: watch::Sender<bool>,
71    state: AtomicU8,
72}
73
74impl Worker {
75    /// Creates a Worker or panics when options or runtime initialization are invalid.
76    ///
77    /// # Panics
78    ///
79    /// Panics for an invalid bind address, zero index-sync timeout, invalid Dex endpoint, or Tokio
80    /// runtime creation failure. Prefer [`Self::try_new`] for recoverable startup.
81    pub fn new(registry: Registry, blob_cache: Arc<BlobCache>, options: WorkerOptions) -> Self {
82        Self::try_new(registry, blob_cache, options)
83            .unwrap_or_else(|error| panic!("cannot create Rust Worker: {error}"))
84    }
85
86    /// Creates a one-shot Worker without binding its listener.
87    ///
88    /// # Errors
89    ///
90    /// Returns [`SdkError::FlowDefinition`] for invalid Worker options or Service for runtime and
91    /// endpoint initialization failures. The Dex connection is lazy.
92    pub fn try_new(
93        registry: Registry,
94        blob_cache: Arc<BlobCache>,
95        options: WorkerOptions,
96    ) -> SdkResult<Self> {
97        let runtime = Runtime::new().map_err(service_error)?;
98        if options.attribute_index_sync_timeout_value().is_zero() {
99            return Err(SdkError::FlowDefinition {
100                message: "attribute index sync timeout must be positive".to_string(),
101            });
102        }
103        let bind_address = options
104            .bind_address_value()
105            .parse::<SocketAddr>()
106            .map_err(|error| SdkError::FlowDefinition {
107                message: format!("invalid Worker bind address: {error}"),
108            })?;
109        let worker_target = options.worker_target_value().cloned().unwrap_or_else(|| {
110            let host = if bind_address.ip().is_unspecified() {
111                "127.0.0.1".to_string()
112            } else {
113                bind_address.ip().to_string()
114            };
115            WorkerTarget::new(format!("{host}:{}", bind_address.port()))
116        });
117        let endpoint = Endpoint::from_shared(endpoint_address(options.server_address_value()))
118            .map_err(service_error)?;
119        let flow_service = {
120            let _runtime_guard = runtime.enter();
121            FlowServiceClient::new(endpoint.connect_lazy())
122        };
123        let attribute_indexes = registry.attribute_indexes().clone();
124        let hydrator = ValueHydrator::new(flow_service.clone(), blob_cache);
125        let dispatcher = WorkerDispatcher::new(
126            registry,
127            hydrator,
128            flow_service.clone(),
129            runtime.handle().clone(),
130        );
131        let (shutdown, _) = watch::channel(false);
132        Ok(Self {
133            runtime,
134            service: RustWorkerService { dispatcher },
135            flow_service,
136            attribute_indexes,
137            attribute_index_sync_timeout: options.attribute_index_sync_timeout_value(),
138            bind_address,
139            worker_target,
140            shutdown,
141            state: AtomicU8::new(CREATED),
142        })
143    }
144
145    /// Returns the resolved WorkerService target that Clients should advertise to Dex.
146    pub fn worker_target(&self) -> &WorkerTarget {
147        &self.worker_target
148    }
149
150    /// Synchronizes Attribute indexes, serves WorkerService, and blocks until stopped or failed.
151    ///
152    /// # Errors
153    ///
154    /// Returns FlowDefinition if called more than once, or Service for index synchronization,
155    /// timeout, bind, and server failures.
156    pub fn start(&self) -> SdkResult<()> {
157        self.state
158            .compare_exchange(CREATED, RUNNING, Ordering::AcqRel, Ordering::Acquire)
159            .map_err(|_| SdkError::FlowDefinition {
160                message: "Worker can only be started once".to_string(),
161            })?;
162        let mut shutdown = self.shutdown.subscribe();
163        let mut flow_service = self.flow_service.clone();
164        let sync_result = self.runtime.block_on(async {
165            tokio::time::timeout(
166                self.attribute_index_sync_timeout,
167                flow_service.sync_attribute_indexes(SyncAttributeIndexRequest {
168                    attribute_indexes: self.attribute_indexes.clone(),
169                }),
170            )
171            .await
172        });
173        match sync_result {
174            Ok(Ok(_)) => {}
175            Ok(Err(status)) => {
176                self.state.store(STOPPED, Ordering::Release);
177                return Err(service_error(status));
178            }
179            Err(elapsed) => {
180                self.state.store(STOPPED, Ordering::Release);
181                return Err(service_error(elapsed));
182            }
183        }
184        let result = self.runtime.block_on(
185            Server::builder()
186                .add_service(WorkerServiceServer::new(self.service.clone()))
187                .serve_with_shutdown(self.bind_address, async move {
188                    while !*shutdown.borrow() && shutdown.changed().await.is_ok() {}
189                }),
190        );
191        self.state.store(STOPPED, Ordering::Release);
192        result.map_err(service_error)
193    }
194
195    /// Requests shutdown of a running Worker and returns immediately.
196    ///
197    /// Calling `stop` before or after `start` is safe; [`Self::start`] observes the shutdown signal.
198    pub fn stop(&self) {
199        let _ = self.shutdown.send(true);
200    }
201}
202
203#[derive(Clone)]
204struct RustWorkerService {
205    dispatcher: WorkerDispatcher,
206}
207
208#[tonic::async_trait]
209impl WorkerService for RustWorkerService {
210    async fn invoke_wait_for_method(
211        &self,
212        request: Request<InvokeWaitForMethodRequest>,
213    ) -> Result<Response<InvokeWaitForMethodResponse>, Status> {
214        self.dispatcher
215            .invoke_wait_for(request.into_inner())
216            .await
217            .map(Response::new)
218            .map_err(worker_status)
219    }
220
221    async fn invoke_execute_method(
222        &self,
223        request: Request<InvokeExecuteMethodRequest>,
224    ) -> Result<Response<InvokeExecuteMethodResponse>, Status> {
225        self.dispatcher
226            .invoke_execute(request.into_inner())
227            .await
228            .map(Response::new)
229            .map_err(worker_status)
230    }
231
232    async fn invoke_worker_rpc(
233        &self,
234        request: Request<InvokeWorkerRpcRequest>,
235    ) -> Result<Response<InvokeWorkerRpcResponse>, Status> {
236        self.dispatcher
237            .invoke_rpc(request.into_inner())
238            .await
239            .map(Response::new)
240            .map_err(worker_status)
241    }
242}
243
244#[derive(Clone, PartialEq, Message)]
245struct GoogleRpcStatus {
246    #[prost(int32, tag = "1")]
247    code: i32,
248    #[prost(string, tag = "2")]
249    message: String,
250    #[prost(message, repeated, tag = "3")]
251    details: Vec<Any>,
252}
253
254fn worker_status(error: HandlerError) -> Status {
255    let message = error.to_string();
256    let stack_trace = truncate_stack_trace(&format!("{message}\n{}", error.stack_trace()));
257    let worker_error = WorkerErrorResponse {
258        detail: message.clone(),
259        error_type: error.error_type().to_string(),
260        stack_trace,
261        retry_after_seconds: error.retry_after_seconds(),
262    };
263    let status = GoogleRpcStatus {
264        code: Code::Unknown as i32,
265        message: message.clone(),
266        details: vec![Any {
267            type_url: "type.googleapis.com/dex.WorkerErrorResponse".to_string(),
268            value: worker_error.encode_to_vec(),
269        }],
270    };
271    Status::with_details(Code::Unknown, message, status.encode_to_vec().into())
272}
273
274fn truncate_stack_trace(value: &str) -> String {
275    let encoded = value.as_bytes();
276    if encoded.len() <= MAX_WORKER_STACK_TRACE_BYTES {
277        return value.to_string();
278    }
279    let mut prefix_length = MAX_WORKER_STACK_TRACE_BYTES - STACK_TRACE_TRUNCATION_MARKER.len();
280    while prefix_length > 0 && encoded[prefix_length] & 0xc0 == 0x80 {
281        prefix_length -= 1;
282    }
283    format!(
284        "{}{}",
285        String::from_utf8_lossy(&encoded[..prefix_length]),
286        STACK_TRACE_TRUNCATION_MARKER
287    )
288}
289
290fn endpoint_address(address: &str) -> String {
291    if address.contains("://") {
292        address.to_string()
293    } else {
294        format!("http://{address}")
295    }
296}
297
298fn service_error(error: impl std::fmt::Display) -> SdkError {
299    SdkError::Service {
300        service: crate::ServiceError::local("worker", error.to_string()),
301    }
302}
dex_sdk/worker.rs

dex_sdk/
worker.rs
