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(registry, hydrator);
126        let (shutdown, _) = watch::channel(false);
127        Ok(Self {
128            runtime,
129            service: RustWorkerService { dispatcher },
130            flow_service,
131            attribute_indexes,
132            attribute_index_sync_timeout: options.attribute_index_sync_timeout_value(),
133            bind_address,
134            worker_target,
135            shutdown,
136            state: AtomicU8::new(CREATED),
137        })
138    }
139
140    /// Returns the resolved WorkerService target that Clients should advertise to Dex.
141    pub fn worker_target(&self) -> &WorkerTarget {
142        &self.worker_target
143    }
144
145    /// Synchronizes Attribute indexes, serves WorkerService, and blocks until stopped or failed.
146    ///
147    /// # Errors
148    ///
149    /// Returns FlowDefinition if called more than once, or Service for index synchronization,
150    /// timeout, bind, and server failures.
151    pub fn start(&self) -> SdkResult<()> {
152        self.state
153            .compare_exchange(CREATED, RUNNING, Ordering::AcqRel, Ordering::Acquire)
154            .map_err(|_| SdkError::FlowDefinition {
155                message: "Worker can only be started once".to_string(),
156            })?;
157        let mut shutdown = self.shutdown.subscribe();
158        let mut flow_service = self.flow_service.clone();
159        let sync_result = self.runtime.block_on(async {
160            tokio::time::timeout(
161                self.attribute_index_sync_timeout,
162                flow_service.sync_attribute_indexes(SyncAttributeIndexRequest {
163                    attribute_indexes: self.attribute_indexes.clone(),
164                }),
165            )
166            .await
167        });
168        match sync_result {
169            Ok(Ok(_)) => {}
170            Ok(Err(status)) => {
171                self.state.store(STOPPED, Ordering::Release);
172                return Err(service_error(status));
173            }
174            Err(elapsed) => {
175                self.state.store(STOPPED, Ordering::Release);
176                return Err(service_error(elapsed));
177            }
178        }
179        let result = self.runtime.block_on(
180            Server::builder()
181                .add_service(WorkerServiceServer::new(self.service.clone()))
182                .serve_with_shutdown(self.bind_address, async move {
183                    while !*shutdown.borrow() && shutdown.changed().await.is_ok() {}
184                }),
185        );
186        self.state.store(STOPPED, Ordering::Release);
187        result.map_err(service_error)
188    }
189
190    /// Requests shutdown of a running Worker and returns immediately.
191    ///
192    /// Calling `stop` before or after `start` is safe; [`Self::start`] observes the shutdown signal.
193    pub fn stop(&self) {
194        let _ = self.shutdown.send(true);
195    }
196}
197
198#[derive(Clone)]
199struct RustWorkerService {
200    dispatcher: WorkerDispatcher,
201}
202
203#[tonic::async_trait]
204impl WorkerService for RustWorkerService {
205    async fn invoke_wait_for_method(
206        &self,
207        request: Request<InvokeWaitForMethodRequest>,
208    ) -> Result<Response<InvokeWaitForMethodResponse>, Status> {
209        self.dispatcher
210            .invoke_wait_for(request.into_inner())
211            .await
212            .map(Response::new)
213            .map_err(worker_status)
214    }
215
216    async fn invoke_execute_method(
217        &self,
218        request: Request<InvokeExecuteMethodRequest>,
219    ) -> Result<Response<InvokeExecuteMethodResponse>, Status> {
220        self.dispatcher
221            .invoke_execute(request.into_inner())
222            .await
223            .map(Response::new)
224            .map_err(worker_status)
225    }
226
227    async fn invoke_worker_rpc(
228        &self,
229        request: Request<InvokeWorkerRpcRequest>,
230    ) -> Result<Response<InvokeWorkerRpcResponse>, Status> {
231        self.dispatcher
232            .invoke_rpc(request.into_inner())
233            .await
234            .map(Response::new)
235            .map_err(worker_status)
236    }
237}
238
239#[derive(Clone, PartialEq, Message)]
240struct GoogleRpcStatus {
241    #[prost(int32, tag = "1")]
242    code: i32,
243    #[prost(string, tag = "2")]
244    message: String,
245    #[prost(message, repeated, tag = "3")]
246    details: Vec<Any>,
247}
248
249fn worker_status(error: HandlerError) -> Status {
250    let message = error.to_string();
251    let stack_trace = truncate_stack_trace(&format!("{message}\n{}", error.stack_trace()));
252    let worker_error = WorkerErrorResponse {
253        detail: message.clone(),
254        error_type: error.error_type().to_string(),
255        stack_trace,
256        retry_after_seconds: error.retry_after_seconds(),
257    };
258    let status = GoogleRpcStatus {
259        code: Code::Unknown as i32,
260        message: message.clone(),
261        details: vec![Any {
262            type_url: "type.googleapis.com/dex.WorkerErrorResponse".to_string(),
263            value: worker_error.encode_to_vec(),
264        }],
265    };
266    Status::with_details(Code::Unknown, message, status.encode_to_vec().into())
267}
268
269fn truncate_stack_trace(value: &str) -> String {
270    let encoded = value.as_bytes();
271    if encoded.len() <= MAX_WORKER_STACK_TRACE_BYTES {
272        return value.to_string();
273    }
274    let mut prefix_length = MAX_WORKER_STACK_TRACE_BYTES - STACK_TRACE_TRUNCATION_MARKER.len();
275    while prefix_length > 0 && encoded[prefix_length] & 0xc0 == 0x80 {
276        prefix_length -= 1;
277    }
278    format!(
279        "{}{}",
280        String::from_utf8_lossy(&encoded[..prefix_length]),
281        STACK_TRACE_TRUNCATION_MARKER
282    )
283}
284
285fn endpoint_address(address: &str) -> String {
286    if address.contains("://") {
287        address.to_string()
288    } else {
289        format!("http://{address}")
290    }
291}
292
293fn service_error(error: impl std::fmt::Display) -> SdkError {
294    SdkError::Service {
295        service: crate::ServiceError::local("worker", error.to_string()),
296    }
297}
dex_sdk/worker.rs

dex_sdk/
worker.rs
