dex_sdk/
worker.rs

1// Copyright (c) 2026 Super Durable, Inc.
2//
3// Licensed under the Sustainable Use 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-Sustainable-Use-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    InvokeExecuteMethodOutput, InvokeExecuteMethodRequest, InvokeWaitForMethodOutput,
19    InvokeWaitForMethodRequest, 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::worker_output::WorkerResponseStream;
32use crate::{BlobCache, HandlerError, Registry, SdkError, SdkResult, WorkerOptions, WorkerTarget};
33
34// Keep worker error status small enough for default gRPC trailer limits after
35// the server wraps WorkerErrorResponse in ServiceErrorResponse.
36const MAX_WORKER_STACK_TRACE_BYTES: usize = 4 * 1024;
37const STACK_TRACE_TRUNCATION_MARKER: &str = "\n... stack trace truncated by Dex Rust SDK ...";
38
39const CREATED: u8 = 0;
40const RUNNING: u8 = 1;
41const STOPPED: u8 = 2;
42
43/// Hosts registered Flow Step and RPC handlers over application WorkerService.
44///
45/// A Worker owns an internal Tokio runtime and is one-shot. [`Self::start`] first synchronizes
46/// Attribute indexes with Dex, then blocks the calling thread while serving gRPC. Call [`Self::stop`]
47/// from another thread to request graceful server shutdown.
48///
49/// # Examples
50///
51/// ```no_run
52/// use dex_sdk::{BlobCache, BlobCacheConfig, Registry, Worker, WorkerOptions};
53/// use std::sync::Arc;
54///
55/// let cache = Arc::new(BlobCache::open(BlobCacheConfig::new(
56///     "/tmp/dex-blobs", 64 * 1024 * 1024, 0,
57/// )?)?);
58/// let worker = Worker::try_new(Registry::new(), cache, WorkerOptions::new())?;
59/// println!("advertise {}", worker.worker_target().address());
60/// worker.start()?;
61/// # Ok::<(), Box<dyn std::error::Error>>(())
62/// ```
63pub struct Worker {
64    runtime: Runtime,
65    service: RustWorkerService,
66    flow_service: FlowServiceClient<Channel>,
67    attribute_indexes: HashMap<String, i32>,
68    attribute_index_sync_timeout: std::time::Duration,
69    bind_address: SocketAddr,
70    worker_target: WorkerTarget,
71    shutdown: watch::Sender<bool>,
72    state: AtomicU8,
73}
74
75impl Worker {
76    /// Creates a Worker or panics when options or runtime initialization are invalid.
77    ///
78    /// # Panics
79    ///
80    /// Panics for an invalid bind address, zero index-sync timeout, invalid Dex endpoint, or Tokio
81    /// runtime creation failure. Prefer [`Self::try_new`] for recoverable startup.
82    pub fn new(registry: Registry, blob_cache: Arc<BlobCache>, options: WorkerOptions) -> Self {
83        Self::try_new(registry, blob_cache, options)
84            .unwrap_or_else(|error| panic!("cannot create Rust Worker: {error}"))
85    }
86
87    /// Creates a one-shot Worker without binding its listener.
88    ///
89    /// # Errors
90    ///
91    /// Returns [`SdkError::FlowDefinition`] for invalid Worker options or Service for runtime and
92    /// endpoint initialization failures. The Dex connection is lazy.
93    pub fn try_new(
94        registry: Registry,
95        blob_cache: Arc<BlobCache>,
96        options: WorkerOptions,
97    ) -> SdkResult<Self> {
98        let runtime = Runtime::new().map_err(service_error)?;
99        if options.attribute_index_sync_timeout_value().is_zero() {
100            return Err(SdkError::FlowDefinition {
101                message: "attribute index sync timeout must be positive".to_string(),
102            });
103        }
104        let bind_address = options
105            .bind_address_value()
106            .parse::<SocketAddr>()
107            .map_err(|error| SdkError::FlowDefinition {
108                message: format!("invalid Worker bind address: {error}"),
109            })?;
110        let worker_target = options.worker_target_value().cloned().unwrap_or_else(|| {
111            let host = if bind_address.ip().is_unspecified() {
112                "127.0.0.1".to_string()
113            } else {
114                bind_address.ip().to_string()
115            };
116            WorkerTarget::new(format!("{host}:{}", bind_address.port()))
117        });
118        let endpoint = Endpoint::from_shared(endpoint_address(options.server_address_value()))
119            .map_err(service_error)?;
120        let flow_service = {
121            let _runtime_guard = runtime.enter();
122            FlowServiceClient::new(endpoint.connect_lazy())
123        };
124        let attribute_indexes = registry.attribute_indexes().clone();
125        let hydrator = ValueHydrator::new(flow_service.clone(), blob_cache);
126        let dispatcher = WorkerDispatcher::new(registry, hydrator);
127        let (shutdown, _) = watch::channel(false);
128        Ok(Self {
129            runtime,
130            service: RustWorkerService { dispatcher },
131            flow_service,
132            attribute_indexes,
133            attribute_index_sync_timeout: options.attribute_index_sync_timeout_value(),
134            bind_address,
135            worker_target,
136            shutdown,
137            state: AtomicU8::new(CREATED),
138        })
139    }
140
141    /// Returns the resolved WorkerService target that Clients should advertise to Dex.
142    pub fn worker_target(&self) -> &WorkerTarget {
143        &self.worker_target
144    }
145
146    /// Synchronizes Attribute indexes, serves WorkerService, and blocks until stopped or failed.
147    ///
148    /// # Errors
149    ///
150    /// Returns FlowDefinition if called more than once, or Service for index synchronization,
151    /// timeout, bind, and server failures.
152    pub fn start(&self) -> SdkResult<()> {
153        self.state
154            .compare_exchange(CREATED, RUNNING, Ordering::AcqRel, Ordering::Acquire)
155            .map_err(|_| SdkError::FlowDefinition {
156                message: "Worker can only be started once".to_string(),
157            })?;
158        let mut shutdown = self.shutdown.subscribe();
159        let mut flow_service = self.flow_service.clone();
160        let sync_result = self.runtime.block_on(async {
161            tokio::time::timeout(
162                self.attribute_index_sync_timeout,
163                flow_service.sync_attribute_indexes(SyncAttributeIndexRequest {
164                    attribute_indexes: self.attribute_indexes.clone(),
165                }),
166            )
167            .await
168        });
169        match sync_result {
170            Ok(Ok(_)) => {}
171            Ok(Err(status)) => {
172                self.state.store(STOPPED, Ordering::Release);
173                return Err(service_error(status));
174            }
175            Err(elapsed) => {
176                self.state.store(STOPPED, Ordering::Release);
177                return Err(service_error(elapsed));
178            }
179        }
180        let result = self.runtime.block_on(
181            Server::builder()
182                .add_service(WorkerServiceServer::new(self.service.clone()))
183                .serve_with_shutdown(self.bind_address, async move {
184                    while !*shutdown.borrow() && shutdown.changed().await.is_ok() {}
185                }),
186        );
187        self.state.store(STOPPED, Ordering::Release);
188        result.map_err(service_error)
189    }
190
191    /// Requests shutdown of a running Worker and returns immediately.
192    ///
193    /// Calling `stop` before or after `start` is safe; [`Self::start`] observes the shutdown signal.
194    pub fn stop(&self) {
195        let _ = self.shutdown.send(true);
196    }
197}
198
199#[derive(Clone)]
200struct RustWorkerService {
201    dispatcher: WorkerDispatcher,
202}
203
204#[tonic::async_trait]
205impl WorkerService for RustWorkerService {
206    type InvokeWaitForMethodStream = WorkerResponseStream<InvokeWaitForMethodOutput>;
207
208    async fn invoke_wait_for_method(
209        &self,
210        request: Request<InvokeWaitForMethodRequest>,
211    ) -> Result<Response<Self::InvokeWaitForMethodStream>, Status> {
212        Ok(Response::new(
213            self.dispatcher
214                .invoke_wait_for(request.into_inner())
215                .into_stream(worker_status),
216        ))
217    }
218
219    type InvokeExecuteMethodStream = WorkerResponseStream<InvokeExecuteMethodOutput>;
220
221    async fn invoke_execute_method(
222        &self,
223        request: Request<InvokeExecuteMethodRequest>,
224    ) -> Result<Response<Self::InvokeExecuteMethodStream>, Status> {
225        Ok(Response::new(
226            self.dispatcher
227                .invoke_execute(request.into_inner())
228                .into_stream(worker_status),
229        ))
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
