Skip to main content

mz_compute_client/
service.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10//! Compute layer client and server.
11
12use std::collections::{BTreeMap, BTreeSet};
13use std::mem;
14
15use async_trait::async_trait;
16use bytesize::ByteSize;
17use differential_dataflow::lattice::Lattice;
18use mz_expr::row::RowCollection;
19use mz_ore::cast::CastInto;
20use mz_ore::soft_panic_or_log;
21use mz_ore::tracing::OpenTelemetryContext;
22use mz_repr::{GlobalId, Timestamp, UpdateCollection};
23use mz_service::client::{GenericClient, Partitionable, PartitionedState};
24use timely::PartialOrder;
25use timely::progress::frontier::{Antichain, MutableAntichain};
26use uuid::Uuid;
27
28use crate::protocol::command::ComputeCommand;
29use crate::protocol::response::{
30    ComputeResponse, CopyToResponse, FrontiersResponse, PeekError, PeekResponse,
31    StashedPeekResponse, SubscribeBatch, SubscribeResponse,
32};
33
34/// A client to a compute server.
35pub trait ComputeClient: GenericClient<ComputeCommand, ComputeResponse> {}
36
37impl<C> ComputeClient for C where C: GenericClient<ComputeCommand, ComputeResponse> {}
38
39#[async_trait]
40impl GenericClient<ComputeCommand, ComputeResponse> for Box<dyn ComputeClient> {
41    async fn send(&mut self, cmd: ComputeCommand) -> Result<(), anyhow::Error> {
42        (**self).send(cmd).await
43    }
44
45    /// # Cancel safety
46    ///
47    /// This method is cancel safe. If `recv` is used as the event in a [`tokio::select!`]
48    /// statement and some other branch completes first, it is guaranteed that no messages were
49    /// received by this client.
50    async fn recv(&mut self) -> Result<Option<ComputeResponse>, anyhow::Error> {
51        // `GenericClient::recv` is required to be cancel safe.
52        (**self).recv().await
53    }
54}
55
56/// Maintained state for partitioned compute clients.
57///
58/// This helper type unifies the responses of multiple partitioned workers in order to present as a
59/// single worker:
60///
61///   * It emits `Frontiers` responses reporting the minimum/meet of frontiers reported by the
62///     individual workers.
63///   * It emits `PeekResponse`s and `SubscribeResponse`s reporting the union of the responses
64///     received from the workers.
65///
66/// In the compute communication stack, this client is instantiated several times:
67///
68///   * One instance on the controller side, dispatching between cluster processes.
69///   * One instance in each cluster process, dispatching between timely worker threads.
70///
71/// Note that because compute commands, except `Hello` and `UpdateConfiguration`, are only
72/// sent to the first process, the cluster-side instances of `PartitionedComputeState` are not
73/// guaranteed to see all compute commands. Or more specifically: The instance running inside
74/// process 0 sees all commands, whereas the instances running inside the other processes only see
75/// `Hello` and `UpdateConfiguration`. The `PartitionedComputeState` implementation must be
76/// able to cope with this limited visibility. It does so by performing most of its state management
77/// based on observed compute responses rather than commands.
78#[derive(Debug)]
79pub struct PartitionedComputeState {
80    /// Number of partitions the state machine represents.
81    parts: usize,
82    /// The maximum result size this state machine can return.
83    ///
84    /// This is updated upon receiving [`ComputeCommand::UpdateConfiguration`]s.
85    max_result_size: u64,
86    /// Tracked frontiers for indexes and sinks.
87    ///
88    /// Frontier tracking for a collection is initialized when the first `Frontiers` response
89    /// for that collection is received. Frontier tracking is ceased when all shards have reported
90    /// advancement to the empty frontier for all frontier kinds.
91    ///
92    /// The compute protocol requires that shards always emit `Frontiers` responses reporting empty
93    /// frontiers for all frontier kinds when a collection is dropped. It further requires that no
94    /// further `Frontier` responses are emitted for a collection after the empty frontiers were
95    /// reported. These properties ensure that a) we always cease frontier tracking for collections
96    /// that have been dropped and b) frontier tracking for a collection is not re-initialized
97    /// after it was ceased.
98    frontiers: BTreeMap<GlobalId, TrackedFrontiers>,
99    /// For each in-progress peek the response data received so far, and the set of shards that
100    /// provided responses already.
101    ///
102    /// Tracking of responses for a peek is initialized when the first `PeekResponse` for that peek
103    /// is received. Once all shards have provided a `PeekResponse`, a unified peek response is
104    /// emitted and the peek tracking state is dropped again.
105    ///
106    /// The compute protocol requires that exactly one response is emitted for each peek. This
107    /// property ensures that a) we can eventually drop the tracking state maintained for a peek
108    /// and b) we won't re-initialize tracking for a peek we have already served.
109    peek_responses: BTreeMap<Uuid, PendingPeek>,
110    /// For each in-progress copy-to the response data received so far, and the set of shards that
111    /// provided responses already.
112    ///
113    /// Tracking of responses for a COPY TO is initialized when the first `CopyResponse` for that command
114    /// is received. Once all shards have provided a `CopyResponse`, a unified copy response is
115    /// emitted and the copy_to tracking state is dropped again.
116    ///
117    /// The compute protocol requires that exactly one response is emitted for each COPY TO command. This
118    /// property ensures that a) we can eventually drop the tracking state maintained for a copy
119    /// and b) we won't re-initialize tracking for a copy we have already served.
120    copy_to_responses: BTreeMap<GlobalId, (CopyToResponse, BTreeSet<usize>)>,
121    /// Tracks in-progress `SUBSCRIBE`s, and the stashed rows we are holding back until their
122    /// timestamps are complete.
123    ///
124    /// The updates may be `Err` if any of the batches have reported an error, in which case the
125    /// subscribe is permanently borked.
126    ///
127    /// Tracking of a subscribe is initialized when the first `SubscribeResponse` for that
128    /// subscribe is received. Once all shards have emitted an "end-of-subscribe" response the
129    /// subscribe tracking state is dropped again.
130    ///
131    /// The compute protocol requires that for a subscribe that shuts down an end-of-subscribe
132    /// response is emitted:
133    ///
134    ///   * Either a `Batch` response reporting advancement to the empty frontier...
135    ///   * ... or a `DroppedAt` response reporting that the subscribe was dropped before
136    ///     completing.
137    ///
138    /// The compute protocol further requires that no further `SubscribeResponse`s are emitted for
139    /// a subscribe after an end-of-subscribe was reported.
140    ///
141    /// These two properties ensure that a) once a subscribe has shut down, we can eventually drop
142    /// the tracking state maintained for it and b) we won't re-initialize tracking for a subscribe
143    /// we have already dropped.
144    pending_subscribes: BTreeMap<GlobalId, PendingSubscribe>,
145}
146
147impl Partitionable<ComputeCommand, ComputeResponse> for (ComputeCommand, ComputeResponse) {
148    type PartitionedState = PartitionedComputeState;
149
150    fn new(parts: usize) -> PartitionedComputeState {
151        PartitionedComputeState {
152            parts,
153            max_result_size: u64::MAX,
154            frontiers: BTreeMap::new(),
155            peek_responses: BTreeMap::new(),
156            pending_subscribes: BTreeMap::new(),
157            copy_to_responses: BTreeMap::new(),
158        }
159    }
160}
161
162impl PartitionedComputeState {
163    /// Observes commands that move past.
164    pub fn observe_command(&mut self, command: &ComputeCommand) {
165        match command {
166            ComputeCommand::UpdateConfiguration(config) => {
167                if let Some(max_result_size) = config.max_result_size {
168                    self.max_result_size = max_result_size;
169                }
170            }
171            _ => {
172                // We are not guaranteed to observe other compute commands. We
173                // must therefore not add any logic here that relies on doing so.
174            }
175        }
176    }
177
178    /// Absorb a [`ComputeResponse::Frontiers`].
179    fn absorb_frontiers(
180        &mut self,
181        shard_id: usize,
182        collection_id: GlobalId,
183        frontiers: FrontiersResponse,
184    ) -> Option<ComputeResponse> {
185        let tracked = self
186            .frontiers
187            .entry(collection_id)
188            .or_insert_with(|| TrackedFrontiers::new(self.parts));
189
190        let write_frontier = frontiers
191            .write_frontier
192            .and_then(|f| tracked.update_write_frontier(shard_id, &f));
193        let input_frontier = frontiers
194            .input_frontier
195            .and_then(|f| tracked.update_input_frontier(shard_id, &f));
196        let output_frontier = frontiers
197            .output_frontier
198            .and_then(|f| tracked.update_output_frontier(shard_id, &f));
199
200        let frontiers = FrontiersResponse {
201            write_frontier,
202            input_frontier,
203            output_frontier,
204        };
205        let result = frontiers
206            .has_updates()
207            .then_some(ComputeResponse::Frontiers(collection_id, frontiers));
208
209        if tracked.all_empty() {
210            // All shards have reported advancement to the empty frontier, so we do not
211            // expect further updates for this collection.
212            self.frontiers.remove(&collection_id);
213        }
214
215        result
216    }
217
218    /// Absorb a [`ComputeResponse::PeekResponse`].
219    fn absorb_peek_response(
220        &mut self,
221        shard_id: usize,
222        uuid: Uuid,
223        response: PeekResponse,
224        otel_ctx: OpenTelemetryContext,
225    ) -> Option<ComputeResponse> {
226        let pending = self
227            .peek_responses
228            .entry(uuid)
229            .or_insert_with(PendingPeek::new);
230        pending.absorb(shard_id, response, self.max_result_size);
231
232        if pending.ready_shards.len() == self.parts {
233            let response = self.peek_responses.remove(&uuid).unwrap().response;
234            Some(ComputeResponse::PeekResponse(uuid, response, otel_ctx))
235        } else {
236            None
237        }
238    }
239
240    /// Absorb a [`ComputeResponse::CopyToResponse`].
241    fn absorb_copy_to_response(
242        &mut self,
243        shard_id: usize,
244        copyto_id: GlobalId,
245        response: CopyToResponse,
246    ) -> Option<ComputeResponse> {
247        use CopyToResponse::*;
248
249        let (merged, ready_shards) = self
250            .copy_to_responses
251            .entry(copyto_id)
252            .or_insert((CopyToResponse::RowCount(0), BTreeSet::new()));
253
254        let first = ready_shards.insert(shard_id);
255        assert!(first, "duplicate copy-to response");
256
257        let resp1 = mem::replace(merged, Dropped);
258        *merged = match (resp1, response) {
259            (Dropped, _) | (_, Dropped) => Dropped,
260            (Error(e), _) | (_, Error(e)) => Error(e),
261            (RowCount(r1), RowCount(r2)) => RowCount(r1 + r2),
262        };
263
264        if ready_shards.len() == self.parts {
265            let (response, _) = self.copy_to_responses.remove(&copyto_id).unwrap();
266            Some(ComputeResponse::CopyToResponse(copyto_id, response))
267        } else {
268            None
269        }
270    }
271
272    /// Absorb a [`ComputeResponse::SubscribeResponse`].
273    fn absorb_subscribe_response(
274        &mut self,
275        subscribe_id: GlobalId,
276        response: SubscribeResponse,
277    ) -> Option<ComputeResponse> {
278        let tracked = self
279            .pending_subscribes
280            .entry(subscribe_id)
281            .or_insert_with(|| PendingSubscribe::new(self.parts));
282
283        let emit_response = match response {
284            SubscribeResponse::Batch(batch) => {
285                let frontiers = &mut tracked.frontiers;
286                let old_frontier = frontiers.frontier().to_owned();
287                frontiers.update_iter(batch.lower.into_iter().map(|t| (t, -1)));
288                frontiers.update_iter(batch.upper.into_iter().map(|t| (t, 1)));
289                let new_frontier = frontiers.frontier().to_owned();
290
291                tracked.stash(batch.updates, self.max_result_size);
292
293                // If the frontier has advanced, it is time to announce subscribe progress. Unless
294                // we have already announced that the subscribe has been dropped, in which case we
295                // must keep quiet.
296                if old_frontier != new_frontier && !tracked.dropped {
297                    let updates = match &mut tracked.stashed_updates {
298                        Ok(stashed_updates) => {
299                            // Split each collection along the frontier, passing the prefix along.
300                            let mut ship = vec![];
301                            let mut keep = vec![];
302                            for collection in stashed_updates.drain(..) {
303                                let partition_point = collection
304                                    .times()
305                                    .partition_point(|t| !new_frontier.less_equal(t));
306                                let (ship_coll, keep_coll) = collection.split_at(partition_point);
307                                if ship_coll.len() > 0 {
308                                    ship.push(ship_coll);
309                                }
310                                if keep_coll.len() > 0 {
311                                    keep.push(keep_coll);
312                                }
313                            }
314                            tracked.stashed_result_size = keep.iter().map(|c| c.byte_len()).sum();
315                            tracked.stashed_updates = Ok(keep);
316                            Ok(ship)
317                        }
318                        Err(text) => Err(text.clone()),
319                    };
320                    Some(ComputeResponse::SubscribeResponse(
321                        subscribe_id,
322                        SubscribeResponse::Batch(SubscribeBatch {
323                            lower: old_frontier,
324                            upper: new_frontier,
325                            updates,
326                        }),
327                    ))
328                } else {
329                    None
330                }
331            }
332            SubscribeResponse::DroppedAt(frontier) => {
333                tracked
334                    .frontiers
335                    .update_iter(frontier.iter().map(|t| (t.clone(), -1)));
336
337                if tracked.dropped {
338                    None
339                } else {
340                    tracked.dropped = true;
341                    Some(ComputeResponse::SubscribeResponse(
342                        subscribe_id,
343                        SubscribeResponse::DroppedAt(frontier),
344                    ))
345                }
346            }
347        };
348
349        if tracked.frontiers.frontier().is_empty() {
350            // All shards have reported advancement to the empty frontier or dropping, so
351            // we do not expect further updates for this subscribe.
352            self.pending_subscribes.remove(&subscribe_id);
353        }
354
355        emit_response
356    }
357}
358
359impl PartitionedState<ComputeCommand, ComputeResponse> for PartitionedComputeState {
360    fn split_command(&mut self, command: ComputeCommand) -> Vec<Option<ComputeCommand>> {
361        self.observe_command(&command);
362
363        // As specified by the compute protocol:
364        //  * Forward `Hello` and `UpdateConfiguration` commands to all shards.
365        //  * Forward all other commands to the first shard only.
366        match command {
367            command @ ComputeCommand::Hello { .. }
368            | command @ ComputeCommand::UpdateConfiguration(_) => {
369                vec![Some(command); self.parts]
370            }
371            command => {
372                let mut r = vec![None; self.parts];
373                r[0] = Some(command);
374                r
375            }
376        }
377    }
378
379    fn absorb_response(
380        &mut self,
381        shard_id: usize,
382        message: ComputeResponse,
383    ) -> Option<Result<ComputeResponse, anyhow::Error>> {
384        let response = match message {
385            ComputeResponse::Frontiers(id, frontiers) => {
386                self.absorb_frontiers(shard_id, id, frontiers)
387            }
388            ComputeResponse::PeekResponse(uuid, response, otel_ctx) => {
389                self.absorb_peek_response(shard_id, uuid, response, otel_ctx)
390            }
391            ComputeResponse::SubscribeResponse(id, response) => {
392                self.absorb_subscribe_response(id, response)
393            }
394            ComputeResponse::CopyToResponse(id, response) => {
395                self.absorb_copy_to_response(shard_id, id, response)
396            }
397            response @ ComputeResponse::Status(_) => {
398                // Pass through status responses.
399                Some(response)
400            }
401        };
402
403        response.map(Ok)
404    }
405}
406
407/// Tracked frontiers for an index or a sink collection.
408///
409/// Each frontier is maintained both as a `MutableAntichain` across all partitions and individually
410/// for each partition.
411#[derive(Debug)]
412struct TrackedFrontiers {
413    /// The tracked write frontier.
414    write_frontier: (MutableAntichain<Timestamp>, Vec<Antichain<Timestamp>>),
415    /// The tracked input frontier.
416    input_frontier: (MutableAntichain<Timestamp>, Vec<Antichain<Timestamp>>),
417    /// The tracked output frontier.
418    output_frontier: (MutableAntichain<Timestamp>, Vec<Antichain<Timestamp>>),
419}
420
421impl TrackedFrontiers {
422    /// Initializes frontier tracking state for a new collection.
423    fn new(parts: usize) -> Self {
424        // TODO(benesch): fix this dangerous use of `as`.
425        #[allow(clippy::as_conversions)]
426        let parts_diff = parts as i64;
427
428        let mut frontier = MutableAntichain::new();
429        frontier.update_iter([(Timestamp::MIN, parts_diff)]);
430        let part_frontiers = vec![Antichain::from_elem(Timestamp::MIN); parts];
431        let frontier_entry = (frontier, part_frontiers);
432
433        Self {
434            write_frontier: frontier_entry.clone(),
435            input_frontier: frontier_entry.clone(),
436            output_frontier: frontier_entry,
437        }
438    }
439
440    /// Returns whether all tracked frontiers have advanced to the empty frontier.
441    fn all_empty(&self) -> bool {
442        self.write_frontier.0.frontier().is_empty()
443            && self.input_frontier.0.frontier().is_empty()
444            && self.output_frontier.0.frontier().is_empty()
445    }
446
447    /// Updates write frontier tracking with a new shard frontier.
448    ///
449    /// If this causes the global write frontier to advance, the advanced frontier is returned.
450    fn update_write_frontier(
451        &mut self,
452        shard_id: usize,
453        new_shard_frontier: &Antichain<Timestamp>,
454    ) -> Option<Antichain<Timestamp>> {
455        Self::update_frontier(&mut self.write_frontier, shard_id, new_shard_frontier)
456    }
457
458    /// Updates input frontier tracking with a new shard frontier.
459    ///
460    /// If this causes the global input frontier to advance, the advanced frontier is returned.
461    fn update_input_frontier(
462        &mut self,
463        shard_id: usize,
464        new_shard_frontier: &Antichain<Timestamp>,
465    ) -> Option<Antichain<Timestamp>> {
466        Self::update_frontier(&mut self.input_frontier, shard_id, new_shard_frontier)
467    }
468
469    /// Updates output frontier tracking with a new shard frontier.
470    ///
471    /// If this causes the global output frontier to advance, the advanced frontier is returned.
472    fn update_output_frontier(
473        &mut self,
474        shard_id: usize,
475        new_shard_frontier: &Antichain<Timestamp>,
476    ) -> Option<Antichain<Timestamp>> {
477        Self::update_frontier(&mut self.output_frontier, shard_id, new_shard_frontier)
478    }
479
480    /// Updates the provided frontier entry with a new shard frontier.
481    fn update_frontier(
482        entry: &mut (MutableAntichain<Timestamp>, Vec<Antichain<Timestamp>>),
483        shard_id: usize,
484        new_shard_frontier: &Antichain<Timestamp>,
485    ) -> Option<Antichain<Timestamp>> {
486        let (frontier, shard_frontiers) = entry;
487
488        let old_frontier = frontier.frontier().to_owned();
489        let shard_frontier = &mut shard_frontiers[shard_id];
490        frontier.update_iter(shard_frontier.iter().map(|t| (t.clone(), -1)));
491        shard_frontier.join_assign(new_shard_frontier);
492        frontier.update_iter(shard_frontier.iter().map(|t| (t.clone(), 1)));
493
494        let new_frontier = frontier.frontier();
495
496        if PartialOrder::less_than(&old_frontier.borrow(), &new_frontier) {
497            Some(new_frontier.to_owned())
498        } else {
499            None
500        }
501    }
502}
503
504#[derive(Debug)]
505struct PendingSubscribe {
506    /// The subscribe frontiers of the partitioned shards.
507    frontiers: MutableAntichain<Timestamp>,
508    /// The updates we are holding back until their timestamps are complete.
509    stashed_updates: Result<Vec<UpdateCollection>, String>,
510    /// The row size of stashed updates, for `max_result_size` checking.
511    stashed_result_size: usize,
512    /// Whether we have already emitted a `DroppedAt` response for this subscribe.
513    ///
514    /// This field is used to ensure we emit such a response only once.
515    dropped: bool,
516}
517
518impl PendingSubscribe {
519    fn new(parts: usize) -> Self {
520        let mut frontiers = MutableAntichain::new();
521        // TODO(benesch): fix this dangerous use of `as`.
522        #[allow(clippy::as_conversions)]
523        frontiers.update_iter([(Timestamp::MIN, parts as i64)]);
524
525        Self {
526            frontiers,
527            stashed_updates: Ok(Vec::new()),
528            stashed_result_size: 0,
529            dropped: false,
530        }
531    }
532
533    /// Stash a new batch of updates.
534    ///
535    /// This also implements the short-circuit behavior of error responses, and performs
536    /// `max_result_size` checking.
537    fn stash(&mut self, new_updates: Result<Vec<UpdateCollection>, String>, max_result_size: u64) {
538        match (&mut self.stashed_updates, new_updates) {
539            (Err(_), _) => {
540                // Subscribe is borked; nothing to do.
541                // TODO: Consider refreshing error?
542            }
543            (_, Err(text)) => {
544                self.stashed_updates = Err(text);
545            }
546            (Ok(stashed), Ok(new)) => {
547                let new_size: usize = new.iter().map(|coll| coll.byte_len()).sum();
548                self.stashed_result_size += new_size;
549
550                if self.stashed_result_size > max_result_size.cast_into() {
551                    self.stashed_updates = Err(format!(
552                        "total result exceeds max size of {}",
553                        ByteSize::b(max_result_size)
554                    ));
555                } else {
556                    stashed.extend(new);
557                }
558            }
559        }
560    }
561}
562
563/// Accumulates the per-worker responses to one peek into the single response the controller
564/// hands upwards.
565#[derive(Debug)]
566struct PendingPeek {
567    /// The responses merged so far.
568    response: PeekResponse,
569    /// Inline result bytes seen so far, across all shards.
570    ///
571    /// Tracked separately from `response` because a worker's rows are dropped as soon as any
572    /// worker reports an error. Without this the aggregate size check would depend on the order
573    /// the responses happen to arrive in.
574    inline_byte_len: usize,
575    /// The shards that have provided responses.
576    ready_shards: BTreeSet<usize>,
577}
578
579impl PendingPeek {
580    fn new() -> Self {
581        Self {
582            response: PeekResponse::Rows(vec![RowCollection::default()]),
583            inline_byte_len: 0,
584            ready_shards: BTreeSet::new(),
585        }
586    }
587
588    fn absorb(&mut self, shard_id: usize, response: PeekResponse, max_result_size: u64) {
589        let first = self.ready_shards.insert(shard_id);
590        assert!(first, "duplicate peek response");
591
592        self.inline_byte_len = self
593            .inline_byte_len
594            .saturating_add(response.inline_byte_len());
595        let current = mem::replace(&mut self.response, PeekResponse::Canceled);
596        self.response = merge_peek_responses(current, response);
597
598        // Merging eagerly is what keeps the controller's memory bounded, so the size check has to
599        // happen on every response rather than once at the end.
600        if self.inline_byte_len > max_result_size.cast_into() {
601            let error = PeekError::ResultExceedsMaxSize {
602                max_result_size: max_result_size.cast_into(),
603            };
604            let current = mem::replace(&mut self.response, PeekResponse::Canceled);
605            self.response = merge_peek_responses(current, PeekResponse::Error(error));
606        }
607    }
608}
609
610/// Merge two [`PeekResponse`]s.
611fn merge_peek_responses(resp1: PeekResponse, resp2: PeekResponse) -> PeekResponse {
612    use PeekResponse::*;
613
614    // Cancelations and errors short-circuit. Cancelations take precedence over errors.
615    let (resp1, resp2) = match (resp1, resp2) {
616        (Canceled, _) | (_, Canceled) => return Canceled,
617        (Error(e1), Error(e2)) => return Error(merge_peek_errors(e1, e2)),
618        (Error(e), _) | (_, Error(e)) => return Error(e),
619        resps => resps,
620    };
621
622    match (resp1, resp2) {
623        (Rows(mut rows1), Rows(rows2)) => {
624            rows1.extend(rows2);
625            Rows(rows1)
626        }
627        (Rows(rows), Stashed(mut stashed)) | (Stashed(mut stashed), Rows(rows)) => {
628            stashed.inline_rows.extend(rows);
629            Stashed(stashed)
630        }
631        (Stashed(stashed1), Stashed(stashed2)) => {
632            // Deconstruct so we don't miss adding new fields. We need to be careful about
633            // merging everything!
634            let StashedPeekResponse {
635                num_rows_batches: num_rows_batches1,
636                encoded_size_bytes: encoded_size_bytes1,
637                relation_desc: relation_desc1,
638                shard_id: shard_id1,
639                batches: mut batches1,
640                inline_rows: mut inline_rows1,
641            } = *stashed1;
642            let StashedPeekResponse {
643                num_rows_batches: num_rows_batches2,
644                encoded_size_bytes: encoded_size_bytes2,
645                relation_desc: relation_desc2,
646                shard_id: shard_id2,
647                batches: mut batches2,
648                inline_rows: inline_rows2,
649            } = *stashed2;
650
651            if shard_id1 != shard_id2 {
652                soft_panic_or_log!(
653                    "shard IDs of stashed responses do not match: \
654                             {shard_id1} != {shard_id2}"
655                );
656                return Error(PeekError::unstructured("internal error"));
657            }
658            if relation_desc1 != relation_desc2 {
659                soft_panic_or_log!(
660                    "relation descs of stashed responses do not match: \
661                             {relation_desc1:?} != {relation_desc2:?}"
662                );
663                return Error(PeekError::unstructured("internal error"));
664            }
665
666            batches1.append(&mut batches2);
667            inline_rows1.extend(inline_rows2);
668
669            Stashed(Box::new(StashedPeekResponse {
670                num_rows_batches: num_rows_batches1 + num_rows_batches2,
671                encoded_size_bytes: encoded_size_bytes1 + encoded_size_bytes2,
672                relation_desc: relation_desc1,
673                shard_id: shard_id1,
674                batches: batches1,
675                inline_rows: inline_rows1,
676            }))
677        }
678        _ => unreachable!("handled above"),
679    }
680}
681
682/// Merge two [`PeekError`]s into the one we report.
683///
684/// A row-iteration-limit error only wins against another one, and then the smaller limit wins so
685/// the choice does not depend on which worker answered first. Any other error outranks it: a
686/// query that both overran the limit and failed for a real reason should report the real reason.
687fn merge_peek_errors(error1: PeekError, error2: PeekError) -> PeekError {
688    match (error1, error2) {
689        (
690            PeekError::RowIterationLimitExceeded { limit: limit1 },
691            PeekError::RowIterationLimitExceeded { limit: limit2 },
692        ) => PeekError::RowIterationLimitExceeded {
693            limit: limit1.min(limit2),
694        },
695        (PeekError::RowIterationLimitExceeded { .. }, error)
696        | (error, PeekError::RowIterationLimitExceeded { .. }) => error,
697        (error, _) => error,
698    }
699}
700
701#[cfg(test)]
702mod tests;