laminar_db/process_function/
benchmark.rs

1//! Benchmark fixtures for native processing and private cluster execution hooks.
2
3use std::sync::Arc;
4use std::time::Duration;
5
6use arrow::array::RecordBatch;
7use laminar_core::checkpoint::{CheckpointAssignmentFence, CheckpointParticipant};
8use laminar_core::cluster::control::LeaseDeadline;
9use laminar_core::shuffle::{ShuffleReceiver, ShuffleSender};
10use laminar_core::state::{NodeId, VnodeRegistry};
11use uuid::Uuid;
12
13use super::{NativeProcessFunction, ProcessFunctionDescriptor, ProcessFunctionOperator};
14use crate::error::DbError;
15use crate::operator::sql_query::ClusterShuffleConfig;
16use crate::operator_graph::{GraphOperator, InputFrontier};
17
18const VNODE_COUNT: u32 = 1_024;
19
20/// Execution path measured with the same prepared input and native handler.
21#[derive(Clone, Copy)]
22pub enum NativeProcessBenchmarkMode {
23    /// Existing in-process operator execution without cluster authority.
24    Local,
25    /// Single-owner routing and result fencing under a cluster assignment.
26    SingleOwner,
27    /// Private input routing across two loopback shuffle endpoints. Admission remains closed.
28    TwoOwners,
29}
30
31/// A production native operator with deterministic fixture authority and no connector I/O.
32pub struct NativeProcessBenchmark {
33    operator: ProcessFunctionOperator,
34    peer: Option<BenchmarkPeer>,
35}
36
37struct BenchmarkPeer {
38    operator: ProcessFunctionOperator,
39    receiver: Arc<ShuffleReceiver>,
40}
41
42impl NativeProcessBenchmark {
43    /// Prepare the same operator in the selected execution mode before timing begins.
44    ///
45    /// # Errors
46    /// Rejects invalid descriptors or inability to prepare the loopback shuffle fixture.
47    pub async fn new(
48        descriptor: ProcessFunctionDescriptor,
49        handler: Arc<dyn NativeProcessFunction>,
50        mode: NativeProcessBenchmarkMode,
51    ) -> Result<Self, DbError> {
52        if matches!(mode, NativeProcessBenchmarkMode::Local) {
53            return Ok(Self {
54                operator: ProcessFunctionOperator::new(descriptor, handler, VNODE_COUNT)?,
55                peer: None,
56            });
57        }
58        let owners = match mode {
59            NativeProcessBenchmarkMode::SingleOwner => vec![7; VNODE_COUNT as usize],
60            NativeProcessBenchmarkMode::TwoOwners => (0..VNODE_COUNT)
61                .map(|vnode| 7 + u64::from(vnode % 2))
62                .collect(),
63            NativeProcessBenchmarkMode::Local => unreachable!("local fixture returned above"),
64        };
65        let assignment = CheckpointAssignmentFence::from_owner_map(
66            7,
67            &owners,
68            owners
69                .iter()
70                .copied()
71                .collect::<std::collections::BTreeSet<_>>()
72                .into_iter()
73                .map(|node_id| CheckpointParticipant {
74                    node_id,
75                    boot_incarnation: Uuid::from_u128(u128::from(node_id)),
76                })
77                .collect(),
78        )
79        .map_err(DbError::Config)?;
80        let (operator, scope) = cluster_operator(
81            descriptor.clone(),
82            Arc::clone(&handler),
83            7,
84            &owners,
85            &assignment,
86        )
87        .await?;
88        let peer = if matches!(mode, NativeProcessBenchmarkMode::TwoOwners) {
89            let (operator, remote) =
90                cluster_operator(descriptor, handler, 8, &owners, &assignment).await?;
91            scope.sender.register_peer(8, remote.receiver.local_addr());
92            remote.sender.register_peer(7, scope.receiver.local_addr());
93            Some(BenchmarkPeer {
94                operator,
95                receiver: remote.receiver,
96            })
97        } else {
98            None
99        };
100        Ok(Self { operator, peer })
101    }
102
103    /// Process a prepared Arrow batch through production routing, state validation and application.
104    /// Two-owner measurements drain both operators and transport rows assigned to the peer.
105    ///
106    /// # Errors
107    /// Returns the ordinary operator errors, including lost fixture authority and resource limits.
108    pub async fn step(&mut self, batch: &RecordBatch) -> Result<Vec<RecordBatch>, DbError> {
109        if self.peer.is_some() {
110            return self.step_two_owners(batch).await;
111        }
112        self.operator
113            .process_with_frontiers(
114                &[vec![batch.clone()]],
115                &[InputFrontier {
116                    watermark: None,
117                    idle: false,
118                }],
119            )
120            .await
121    }
122
123    async fn step_two_owners(&mut self, batch: &RecordBatch) -> Result<Vec<RecordBatch>, DbError> {
124        let peer = self
125            .peer
126            .as_mut()
127            .ok_or_else(|| DbError::Config("benchmark peer is missing".into()))?;
128        let before = self
129            .operator
130            .accepted_activations()
131            .saturating_add(peer.operator.accepted_activations());
132        let frontier = [InputFrontier::default()];
133        let mut output = self
134            .operator
135            .process_with_frontiers(&[vec![batch.clone()]], &frontier)
136            .await?;
137        tokio::time::timeout(Duration::from_secs(5), async {
138            loop {
139                for (stage, batch) in peer.receiver.drain_checkpointed_staged() {
140                    peer.operator.stage_checkpointed_shuffle(
141                        &stage,
142                        crate::operator::RetainedBatch::from_received(batch),
143                        i64::MIN,
144                    )?;
145                }
146                output.extend(self.operator.process_with_frontiers(&[], &frontier).await?);
147                output.extend(peer.operator.process_with_frontiers(&[], &frontier).await?);
148                let processed = self
149                    .operator
150                    .accepted_activations()
151                    .saturating_add(peer.operator.accepted_activations())
152                    .saturating_sub(before);
153                if processed == batch.num_rows() as u64
154                    && !self.operator.checkpoint_drain_pending()
155                    && !peer.operator.checkpoint_drain_pending()
156                {
157                    break Ok(output);
158                }
159                tokio::task::yield_now().await;
160            }
161        })
162        .await
163        .map_err(|_| DbError::Config("two-owner benchmark drain timed out".into()))?
164    }
165}
166
167async fn cluster_operator(
168    descriptor: ProcessFunctionDescriptor,
169    handler: Arc<dyn NativeProcessFunction>,
170    node: u64,
171    owners: &[u64],
172    assignment: &CheckpointAssignmentFence,
173) -> Result<(ProcessFunctionOperator, ClusterShuffleConfig), DbError> {
174    let mut operator = ProcessFunctionOperator::new(descriptor, handler, VNODE_COUNT)?;
175    let boot = Uuid::from_u128(u128::from(node));
176    let registry = Arc::new(VnodeRegistry::single_owner(VNODE_COUNT, NodeId(node)));
177    registry.set_assignment_and_version(
178        Arc::from(owners.iter().copied().map(NodeId).collect::<Vec<_>>()),
179        7,
180    );
181    let sender = Arc::new(ShuffleSender::new(node, boot));
182    let receiver = Arc::new(
183        ShuffleReceiver::bind(node, std::net::SocketAddr::from(([127, 0, 0, 1], 0)), boot)
184            .await
185            .map_err(|error| DbError::Config(error.to_string()))?,
186    );
187    let deadline = Arc::new(LeaseDeadline::live_for(Duration::from_secs(300)));
188    sender
189        .bind_process_lease_deadline_pair(&receiver, Arc::clone(&deadline))
190        .map_err(|error| DbError::Config(error.to_string()))?;
191    sender
192        .install_assignment_fence(assignment, owners)
193        .map_err(|error| DbError::Config(error.to_string()))?;
194    receiver
195        .install_assignment_fence(assignment, owners)
196        .map_err(|error| DbError::Config(error.to_string()))?;
197    sender.set_recovery_gen(3);
198    receiver.set_recovery_gen(3);
199    operator.require_cluster_execution(
200        "activity",
201        tokio::runtime::Handle::current(),
202        NodeId(node),
203    )?;
204    operator.bind_startup_assignment(
205        assignment,
206        &owners
207            .iter()
208            .enumerate()
209            .filter_map(|(vnode, owner)| {
210                (*owner == node)
211                    .then_some(u32::try_from(vnode).expect("fixed benchmark vnode domain fits u32"))
212            })
213            .collect::<Vec<_>>(),
214    )?;
215    let scope = ClusterShuffleConfig {
216        registry,
217        sender,
218        receiver,
219        self_id: NodeId(node),
220        topology: None,
221    };
222    operator.bind_process_execution_authority(&scope, deadline)?;
223    Ok((operator, scope))
224}
laminar_db/process_function/benchmark.rs

laminar_db/process_function/
benchmark.rs
