1use 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
34pub 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 async fn recv(&mut self) -> Result<Option<ComputeResponse>, anyhow::Error> {
51 (**self).recv().await
53 }
54}
55
56#[derive(Debug)]
79pub struct PartitionedComputeState {
80 parts: usize,
82 max_result_size: u64,
86 frontiers: BTreeMap<GlobalId, TrackedFrontiers>,
99 peek_responses: BTreeMap<Uuid, PendingPeek>,
110 copy_to_responses: BTreeMap<GlobalId, (CopyToResponse, BTreeSet<usize>)>,
121 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 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 }
175 }
176 }
177
178 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 self.frontiers.remove(&collection_id);
213 }
214
215 result
216 }
217
218 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 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(©to_id).unwrap();
266 Some(ComputeResponse::CopyToResponse(copyto_id, response))
267 } else {
268 None
269 }
270 }
271
272 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 old_frontier != new_frontier && !tracked.dropped {
297 let updates = match &mut tracked.stashed_updates {
298 Ok(stashed_updates) => {
299 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 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 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 Some(response)
400 }
401 };
402
403 response.map(Ok)
404 }
405}
406
407#[derive(Debug)]
412struct TrackedFrontiers {
413 write_frontier: (MutableAntichain<Timestamp>, Vec<Antichain<Timestamp>>),
415 input_frontier: (MutableAntichain<Timestamp>, Vec<Antichain<Timestamp>>),
417 output_frontier: (MutableAntichain<Timestamp>, Vec<Antichain<Timestamp>>),
419}
420
421impl TrackedFrontiers {
422 fn new(parts: usize) -> Self {
424 #[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 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 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 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 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 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 frontiers: MutableAntichain<Timestamp>,
508 stashed_updates: Result<Vec<UpdateCollection>, String>,
510 stashed_result_size: usize,
512 dropped: bool,
516}
517
518impl PendingSubscribe {
519 fn new(parts: usize) -> Self {
520 let mut frontiers = MutableAntichain::new();
521 #[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 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 }
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#[derive(Debug)]
566struct PendingPeek {
567 response: PeekResponse,
569 inline_byte_len: usize,
575 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 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
610fn merge_peek_responses(resp1: PeekResponse, resp2: PeekResponse) -> PeekResponse {
612 use PeekResponse::*;
613
614 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 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
682fn 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;