laminar_db/process_function/remote/
client.rs

1use std::collections::BTreeSet;
2use std::num::NonZeroU32;
3use std::sync::Arc;
4use std::time::{Duration, SystemTime, UNIX_EPOCH};
5
6use laminar_core::state::PartitionKeyCodecV1;
7use prost::Message;
8use sha2::{Digest, Sha256};
9use tokio::sync::Semaphore;
10use tokio_util::sync::CancellationToken;
11use tonic::transport::{Channel, Endpoint};
12use tonic::Request;
13
14use super::codec::encode_activation;
15use super::response::read_response;
16use super::wire::{self, host_frame, worker_frame};
17use super::{
18    RemoteInvocationScope, MAX_FRAME_BYTES, MAX_INVOCATION_WIRE_BYTES, MAX_IN_FLIGHT,
19    PROTOCOL_VERSION,
20};
21use crate::error::DbError;
22use crate::process_function::{
23    ProcessActivation, ProcessActivationResult, ProcessFunctionDescriptor, ProcessRuntime,
24};
25
26/// Bounded v1 gRPC/Arrow client. Clones share one connection pool and one invocation credit pool.
27/// Only explicit IP loopback endpoints are accepted by this development transport.
28#[derive(Clone)]
29pub struct RemoteProcessClient {
30    channel: Channel,
31    descriptor: Arc<ProcessFunctionDescriptor>,
32    digest: [u8; 32],
33    credits: Arc<Semaphore>,
34    max_in_flight: usize,
35    timeout: Duration,
36    python_replay_lifetime: Option<CancellationToken>,
37}
38
39impl RemoteProcessClient {
40    pub(super) fn bind_python_replay_lifetime(&mut self, exited: CancellationToken) {
41        self.python_replay_lifetime = Some(exited);
42    }
43
44    pub(crate) fn has_python_replay_binding(&self) -> bool {
45        self.python_replay_lifetime
46            .as_ref()
47            .is_some_and(|exited| !exited.is_cancelled())
48    }
49
50    /// Immutable descriptor negotiated when this client connected.
51    #[must_use]
52    pub fn descriptor(&self) -> &ProcessFunctionDescriptor {
53        &self.descriptor
54    }
55
56    /// Number of invocation credits shared by clones of this client.
57    #[must_use]
58    pub const fn max_in_flight(&self) -> usize {
59        self.max_in_flight
60    }
61
62    /// Connect to a reference worker over an explicit loopback address.
63    ///
64    /// # Errors
65    /// Rejects nonloopback/plaintext deployment, invalid bounds, or connection failure.
66    pub async fn connect_loopback(
67        endpoint: &str,
68        descriptor: ProcessFunctionDescriptor,
69        max_in_flight: usize,
70        timeout: Duration,
71    ) -> Result<Self, DbError> {
72        let address = endpoint
73            .strip_prefix("http://")
74            .and_then(|value| value.parse::<std::net::SocketAddr>().ok())
75            .filter(|address| address.ip().is_loopback())
76            .ok_or_else(|| {
77                DbError::Unsupported("process worker plaintext endpoint must be loopback".into())
78            })?;
79        if max_in_flight == 0
80            || max_in_flight > MAX_IN_FLIGHT
81            || timeout < Duration::from_millis(1)
82            || timeout > Duration::from_secs(30)
83            || descriptor.runtime == ProcessRuntime::NativeRust
84        {
85            return Err(DbError::InvalidOperation(
86                "remote process client requires a remote runtime, deadline, and 1..=32 credits"
87                    .into(),
88            ));
89        }
90        let manifest = descriptor.to_manifest_json()?;
91        let digest: [u8; 32] = Sha256::digest(manifest).into();
92        let uri = format!("http://{address}");
93        let channel = Endpoint::from_shared(uri)
94            .map_err(|error| DbError::Config(format!("process worker endpoint: {error}")))?
95            .connect_timeout(timeout)
96            .connect()
97            .await
98            .map_err(|error| DbError::Pipeline(format!("connect process worker: {error}")))?;
99        Ok(Self {
100            channel,
101            descriptor: Arc::new(descriptor),
102            digest,
103            credits: Arc::new(Semaphore::new(max_in_flight)),
104            max_in_flight,
105            timeout,
106            python_replay_lifetime: None,
107        })
108    }
109
110    /// Invoke one key-distinct, single-vnode batch. A complete validated response is returned;
111    /// the caller must still validate and apply it under its authoritative ownership fence.
112    ///
113    /// # Errors
114    /// Rejects malformed inputs/responses, budget excess, timeout, or transport failure.
115    pub async fn invoke(
116        &self,
117        scope: &RemoteInvocationScope,
118        activations: &[ProcessActivation],
119    ) -> Result<Vec<ProcessActivationResult>, DbError> {
120        if self
121            .python_replay_lifetime
122            .as_ref()
123            .is_some_and(CancellationToken::is_cancelled)
124        {
125            return Err(DbError::Pipeline(
126                "replay-bound Python worker has exited".into(),
127            ));
128        }
129        // Bounded admission: callers retry on backpressure rather than queueing unbounded RPCs.
130        let _permit = self.credits.try_acquire().map_err(|_| {
131            DbError::BackpressureFail("process worker invocation slots full".into())
132        })?;
133        let request = self.encode_request(scope, activations)?;
134        tokio::time::timeout(self.timeout, self.exchange(scope, activations, request))
135            .await
136            .map_err(|_| DbError::Pipeline("process worker invocation timed out".into()))?
137    }
138
139    fn encode_request(
140        &self,
141        scope: &RemoteInvocationScope,
142        activations: &[ProcessActivation],
143    ) -> Result<Vec<wire::HostFrame>, DbError> {
144        let vnode_count = NonZeroU32::new(scope.vnode_count)
145            .filter(|count| scope.vnode < count.get())
146            .ok_or_else(|| DbError::InvalidOperation("invalid process invocation vnode".into()))?;
147        if scope.operator_id.is_empty()
148            || scope.operator_id.len() > 128
149            || scope.batch_id == scope.attempt_id
150            || activations.is_empty()
151            || activations.len() > self.descriptor.limits.max_batch_rows
152        {
153            return Err(DbError::InvalidOperation(
154                "invalid process invocation scope or size".into(),
155            ));
156        }
157        let deadline_unix_ms = SystemTime::now()
158            .duration_since(UNIX_EPOCH)
159            .map_err(|_| DbError::Pipeline("system clock precedes Unix epoch".into()))?
160            .as_millis()
161            .checked_add(self.timeout.as_millis())
162            .and_then(|value| u64::try_from(value).ok())
163            .ok_or_else(|| DbError::Pipeline("process invocation deadline overflow".into()))?;
164        let mut frames = vec![wire::HostFrame {
165            kind: Some(host_frame::Kind::Open(wire::Open {
166                protocol_version: PROTOCOL_VERSION,
167                descriptor_sha256: self.digest.to_vec(),
168                operator_id: scope.operator_id.clone(),
169                vnode: scope.vnode,
170                owner_generation: scope.owner_generation,
171                recovery_generation: scope.recovery_generation,
172                batch_id: scope.batch_id.as_bytes().to_vec(),
173                attempt_id: scope.attempt_id.as_bytes().to_vec(),
174                input_watermark_us: scope.input_watermark_us,
175                deadline_unix_ms,
176                vnode_count: scope.vnode_count,
177            })),
178        }];
179        let mut ids = BTreeSet::new();
180        let mut keys = BTreeSet::new();
181        let mut input_bytes = 0usize;
182        for activation in activations {
183            if !ids.insert(activation.id)
184                || !keys.insert(activation.key.as_ref())
185                || PartitionKeyCodecV1::vnode_for_encoded(&activation.key, vnode_count)
186                    != scope.vnode
187            {
188                return Err(DbError::InvalidOperation(
189                    "process invocation has another vnode, repeated key, or repeated ID".into(),
190                ));
191            }
192            let encoded = encode_activation(activation, &self.descriptor)?;
193            if let Some(wire::activation::Callback::InputIpc(ipc)) = &encoded.callback {
194                // A one-row Arrow slice can retain the parent batch's full buffers. Charge the
195                // bytes actually sent rather than multiplying that allocation by row count.
196                input_bytes = input_bytes
197                    .checked_add(ipc.len())
198                    .filter(|bytes| *bytes <= self.descriptor.limits.max_input_bytes)
199                    .ok_or_else(|| {
200                        DbError::BackpressureFail("process input IPC byte limit exceeded".into())
201                    })?;
202            }
203            frames.push(wire::HostFrame {
204                kind: Some(host_frame::Kind::Activation(encoded)),
205            });
206        }
207        frames.push(wire::HostFrame {
208            kind: Some(host_frame::Kind::End(wire::End {
209                count: u32::try_from(activations.len()).map_err(|_| {
210                    DbError::BackpressureFail("process batch count overflow".into())
211                })?,
212            })),
213        });
214        frames.iter().try_fold(0usize, |sum, frame| {
215            sum.checked_add(frame.encoded_len())
216                .filter(|bytes| *bytes <= MAX_INVOCATION_WIRE_BYTES)
217                .ok_or_else(|| {
218                    DbError::BackpressureFail("process request wire budget exceeded".into())
219                })
220        })?;
221        Ok(frames)
222    }
223
224    async fn exchange(
225        &self,
226        scope: &RemoteInvocationScope,
227        activations: &[ProcessActivation],
228        frames: Vec<wire::HostFrame>,
229    ) -> Result<Vec<ProcessActivationResult>, DbError> {
230        let mut client =
231            wire::process_worker_client::ProcessWorkerClient::new(self.channel.clone())
232                .max_decoding_message_size(MAX_FRAME_BYTES + 64 * 1024)
233                .max_encoding_message_size(MAX_FRAME_BYTES + 64 * 1024);
234        let mut request = Request::new(tokio_stream::iter(frames));
235        request.set_timeout(self.timeout);
236        let mut inbound = client
237            .exchange(request)
238            .await
239            .map_err(|error| DbError::Pipeline(format!("process worker exchange: {error}")))?
240            .into_inner();
241        let ack = inbound
242            .message()
243            .await
244            .map_err(|error| DbError::Pipeline(format!("process worker response: {error}")))?
245            .ok_or_else(|| DbError::Pipeline("process worker closed without response".into()))?;
246        let ack_bytes = ack.encoded_len();
247        match ack.kind {
248            Some(worker_frame::Kind::Ack(ack))
249                if ack.protocol_version == PROTOCOL_VERSION
250                    && ack.descriptor_sha256 == self.digest
251                    && ack.attempt_id == scope.attempt_id.as_bytes() => {}
252            Some(worker_frame::Kind::Failure(failure)) => {
253                return Err(DbError::Pipeline(format!(
254                    "process worker failed: {}",
255                    failure.message
256                )));
257            }
258            _ => {
259                return Err(DbError::InvalidOperation(
260                    "invalid process worker acknowledgement".into(),
261                ))
262            }
263        }
264        read_response(&self.descriptor, activations, ack_bytes, &mut inbound).await
265    }
266}
laminar_db/process_function/remote/client.rs

laminar_db/process_function/remote/
client.rs
