1use crate::cluster::nats_client::{GenericRequest, GenericStreamRequest, NatsClient};
10use crate::plotter::{Plotter, SectorPlottingProgress};
11use crate::utils::AsyncJoinOnDrop;
12use ab_core_primitives::ed25519::Ed25519PublicKey;
13use ab_core_primitives::sectors::SectorIndex;
14use ab_core_primitives::solutions::ShardCommitmentHash;
15use ab_direct_io_file::AlignedPage;
16use ab_farmer_components::FarmerProtocolInfo;
17use ab_farmer_components::plotting::PlottedSector;
18use ab_farmer_components::sector::sector_size;
19use anyhow::anyhow;
20use async_nats::RequestErrorKind;
21use async_trait::async_trait;
22use backon::{BackoffBuilder, ExponentialBackoff, ExponentialBuilder};
23use bytes::Bytes;
24use derive_more::Display;
25use event_listener_primitives::{Bag, HandlerId};
26use futures::channel::mpsc;
27use futures::future::FusedFuture;
28use futures::stream::FuturesUnordered;
29use futures::{FutureExt, Sink, SinkExt, StreamExt, select, stream};
30use parity_scale_codec::{Decode, Encode};
31use std::error::Error;
32use std::future::pending;
33use std::num::NonZeroUsize;
34use std::pin::pin;
35use std::sync::Arc as StdArc;
36use std::task::Poll;
37use std::time::{Duration, Instant};
38use tokio::sync::{OwnedSemaphorePermit, Semaphore};
39use tokio::time::MissedTickBehavior;
40use tracing::{Instrument, debug, info, info_span, trace, warn};
41use triomphe::Arc;
42use ulid::Ulid;
43
44const FREE_CAPACITY_CHECK_INTERVAL: Duration = Duration::from_secs(1);
45const PING_INTERVAL: Duration = Duration::from_secs(10);
47const PING_TIMEOUT: Duration = Duration::from_mins(1);
49
50pub type HandlerFn3<A, B, C> = StdArc<dyn Fn(&A, &B, &C) + Send + Sync + 'static>;
52type Handler3<A, B, C> = Bag<HandlerFn3<A, B, C>, A, B, C>;
53
54#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Display)]
56pub enum ClusterPlotterId {
57 Ulid(Ulid),
59}
60
61#[expect(
62 clippy::new_without_default,
63 reason = "Default has different semantics"
64)]
65impl ClusterPlotterId {
66 pub fn new() -> Self {
68 Self::Ulid(Ulid::generate())
69 }
70}
71
72#[derive(Debug, Clone, Encode, Decode)]
74struct ClusterPlotterFreeInstanceRequest;
75
76impl GenericRequest for ClusterPlotterFreeInstanceRequest {
77 const SUBJECT: &'static str = "ab.plotter.free-instance";
78 type Response = Option<String>;
80}
81
82#[derive(Debug, Encode, Decode)]
83enum ClusterSectorPlottingProgress {
84 Occupied,
86 Ping,
88 Downloading,
90 Downloaded(Duration),
92 Encoding,
94 Encoded(Duration),
96 Finished {
98 plotted_sector: PlottedSector,
100 time: Duration,
102 },
103 SectorChunk(Result<Bytes, String>),
105 Error {
107 error: String,
109 },
110}
111
112#[derive(Debug, Clone, Encode, Decode)]
114struct ClusterPlotterPlotSectorRequest {
115 public_key: Ed25519PublicKey,
116 shard_commitments_root: ShardCommitmentHash,
117 sector_index: SectorIndex,
118 farmer_protocol_info: FarmerProtocolInfo,
119 pieces_in_sector: u16,
120}
121
122impl GenericStreamRequest for ClusterPlotterPlotSectorRequest {
123 const SUBJECT: &'static str = "ab.plotter.*.plot-sector";
124 type Response = ClusterSectorPlottingProgress;
125}
126
127fn sector_chunk_size(nats_client: &NatsClient) -> usize {
132 nats_client.approximate_max_message_size() / AlignedPage::SIZE * AlignedPage::SIZE
133}
134
135#[derive(Default, Debug)]
136struct Handlers {
137 plotting_progress: Handler3<Ed25519PublicKey, SectorIndex, SectorPlottingProgress>,
138}
139
140#[derive(Debug)]
142pub struct ClusterPlotter {
143 sector_encoding_semaphore: StdArc<Semaphore>,
144 retry_backoff_policy: ExponentialBuilder,
145 nats_client: NatsClient,
146 handlers: Arc<Handlers>,
147 tasks_sender: mpsc::Sender<AsyncJoinOnDrop<()>>,
148 _background_tasks: AsyncJoinOnDrop<()>,
149}
150
151impl Drop for ClusterPlotter {
152 #[inline]
153 fn drop(&mut self) {
154 self.tasks_sender.close_channel();
155 }
156}
157
158#[async_trait]
159impl Plotter for ClusterPlotter {
160 async fn has_free_capacity(&self) -> Result<bool, String> {
161 Ok(self.sector_encoding_semaphore.available_permits() > 0
162 && self
163 .nats_client
164 .request(&ClusterPlotterFreeInstanceRequest, None)
165 .await
166 .map_err(|error| error.to_string())?
167 .is_some())
168 }
169
170 async fn plot_sector(
171 &self,
172 public_key: Ed25519PublicKey,
173 shard_commitments_root: ShardCommitmentHash,
174 sector_index: SectorIndex,
175 farmer_protocol_info: FarmerProtocolInfo,
176 pieces_in_sector: u16,
177 _replotting: bool,
178 mut progress_sender: mpsc::Sender<SectorPlottingProgress>,
179 ) {
180 let start = Instant::now();
181
182 let sector_encoding_permit = match StdArc::clone(&self.sector_encoding_semaphore)
185 .acquire_owned()
186 .await
187 {
188 Ok(sector_encoding_permit) => sector_encoding_permit,
189 Err(error) => {
190 warn!(%error, "Failed to acquire sector encoding permit");
191
192 let progress_updater = ProgressUpdater {
193 public_key,
194 sector_index,
195 handlers: Arc::clone(&self.handlers),
196 };
197
198 progress_updater
199 .update_progress_and_events(
200 &mut progress_sender,
201 SectorPlottingProgress::Error {
202 error: format!("Failed to acquire sector encoding permit: {error}"),
203 },
204 )
205 .await;
206
207 return;
208 }
209 };
210
211 self.plot_sector_internal(
212 start,
213 sector_encoding_permit,
214 public_key,
215 shard_commitments_root,
216 sector_index,
217 farmer_protocol_info,
218 pieces_in_sector,
219 progress_sender,
220 )
221 .await;
222 }
223
224 async fn try_plot_sector(
225 &self,
226 public_key: Ed25519PublicKey,
227 shard_commitments_root: ShardCommitmentHash,
228 sector_index: SectorIndex,
229 farmer_protocol_info: FarmerProtocolInfo,
230 pieces_in_sector: u16,
231 _replotting: bool,
232 progress_sender: mpsc::Sender<SectorPlottingProgress>,
233 ) -> bool {
234 let start = Instant::now();
235
236 let Ok(sector_encoding_permit) =
237 StdArc::clone(&self.sector_encoding_semaphore).try_acquire_owned()
238 else {
239 return false;
240 };
241
242 self.plot_sector_internal(
243 start,
244 sector_encoding_permit,
245 public_key,
246 shard_commitments_root,
247 sector_index,
248 farmer_protocol_info,
249 pieces_in_sector,
250 progress_sender,
251 )
252 .await;
253
254 true
255 }
256}
257
258impl ClusterPlotter {
259 pub fn new(
261 nats_client: NatsClient,
262 sector_encoding_concurrency: NonZeroUsize,
263 retry_backoff_policy: ExponentialBuilder,
264 ) -> Self {
265 let sector_encoding_semaphore =
266 StdArc::new(Semaphore::new(sector_encoding_concurrency.get()));
267
268 let (tasks_sender, mut tasks_receiver) = mpsc::channel(1);
269
270 let background_tasks = AsyncJoinOnDrop::new(
272 tokio::spawn(async move {
273 let background_tasks = FuturesUnordered::new();
274 let mut background_tasks = pin!(background_tasks);
275 background_tasks.push(AsyncJoinOnDrop::new(tokio::spawn(pending::<()>()), true));
277
278 loop {
279 select! {
280 maybe_background_task = tasks_receiver.next().fuse() => {
281 let Some(background_task) = maybe_background_task else {
282 break;
283 };
284
285 background_tasks.push(background_task);
286 },
287 _ = background_tasks.select_next_some() => {
288 }
290 }
291 }
292 }),
293 true,
294 );
295
296 Self {
297 sector_encoding_semaphore,
298 retry_backoff_policy,
299 nats_client,
300 handlers: Arc::default(),
301 tasks_sender,
302 _background_tasks: background_tasks,
303 }
304 }
305
306 pub fn on_plotting_progress(
308 &self,
309 callback: HandlerFn3<Ed25519PublicKey, SectorIndex, SectorPlottingProgress>,
310 ) -> HandlerId {
311 self.handlers.plotting_progress.add(callback)
312 }
313
314 #[expect(clippy::too_many_arguments, reason = "Internal API")]
315 async fn plot_sector_internal<PS>(
316 &self,
317 start: Instant,
318 sector_encoding_permit: OwnedSemaphorePermit,
319 public_key: Ed25519PublicKey,
320 shard_commitments_root: ShardCommitmentHash,
321 sector_index: SectorIndex,
322 farmer_protocol_info: FarmerProtocolInfo,
323 pieces_in_sector: u16,
324 mut progress_sender: PS,
325 ) where
326 PS: Sink<SectorPlottingProgress, Error: Error> + Unpin + Send + 'static,
327 {
328 trace!("Starting plotting, getting plotting permit");
329
330 let progress_updater = ProgressUpdater {
331 public_key,
332 sector_index,
333 handlers: Arc::clone(&self.handlers),
334 };
335
336 let mut retry_backoff_policy = self.retry_backoff_policy.build();
337
338 let free_plotter_instance_fut = get_free_plotter_instance(
340 &self.nats_client,
341 &progress_updater,
342 &mut progress_sender,
343 &mut retry_backoff_policy,
344 );
345 let mut maybe_free_instance = free_plotter_instance_fut.await;
346 if maybe_free_instance.is_none() {
347 return;
348 }
349
350 trace!("Got plotting permit #1");
351
352 let nats_client = self.nats_client.clone();
353
354 let plotting_fut = async move {
355 'outer: loop {
356 let free_instance = if let Some(free_instance) = maybe_free_instance.take() {
358 free_instance
359 } else {
360 let free_plotter_instance_fut = get_free_plotter_instance(
361 &nats_client,
362 &progress_updater,
363 &mut progress_sender,
364 &mut retry_backoff_policy,
365 );
366 let Some(free_instance) = free_plotter_instance_fut.await else {
367 break;
368 };
369 trace!("Got plotting permit #2");
370 free_instance
371 };
372
373 let response_stream_result = nats_client
374 .stream_request(
375 &ClusterPlotterPlotSectorRequest {
376 public_key,
377 shard_commitments_root,
378 sector_index,
379 farmer_protocol_info,
380 pieces_in_sector,
381 },
382 Some(&free_instance),
383 )
384 .await;
385 trace!("Subscribed to plotting notifications");
386
387 let mut response_stream = match response_stream_result {
388 Ok(response_stream) => response_stream,
389 Err(error) => {
390 progress_updater
391 .update_progress_and_events(
392 &mut progress_sender,
393 SectorPlottingProgress::Error {
394 error: format!("Failed make stream request: {error}"),
395 },
396 )
397 .await;
398
399 break;
400 }
401 };
402
403 let (mut sector_sender, sector_receiver) = mpsc::channel(
406 sector_size(pieces_in_sector).div_ceil(sector_chunk_size(&nats_client)),
407 );
408 let mut maybe_sector_receiver = Some(sector_receiver);
409 loop {
410 match tokio::time::timeout(PING_TIMEOUT, response_stream.next()).await {
411 Ok(Some(response)) => {
412 match process_response_notification(
413 &start,
414 &free_instance,
415 &progress_updater,
416 &mut progress_sender,
417 &mut retry_backoff_policy,
418 response,
419 &mut sector_sender,
420 &mut maybe_sector_receiver,
421 )
422 .await
423 {
424 ResponseProcessingResult::Retry => {
425 debug!("Retrying");
426 continue 'outer;
427 }
428 ResponseProcessingResult::Abort => {
429 debug!("Aborting");
430 break 'outer;
431 }
432 ResponseProcessingResult::Continue => {
433 trace!("Continue");
434 }
436 }
437 }
438 Ok(None) => {
439 trace!("Plotting done");
440 break;
441 }
442 Err(_error) => {
443 progress_updater
444 .update_progress_and_events(
445 &mut progress_sender,
446 SectorPlottingProgress::Error {
447 error: "Timed out without ping from plotter".to_string(),
448 },
449 )
450 .await;
451 break;
452 }
453 }
454 }
455
456 break;
457 }
458
459 drop(sector_encoding_permit);
460 };
461
462 let plotting_task =
463 AsyncJoinOnDrop::new(tokio::spawn(plotting_fut.in_current_span()), true);
464 if let Err(error) = self.tasks_sender.clone().send(plotting_task).await {
465 warn!(%error, "Failed to send plotting task");
466
467 let progress = SectorPlottingProgress::Error {
468 error: format!("Failed to send plotting task: {error}"),
469 };
470
471 self.handlers
472 .plotting_progress
473 .call_simple(&public_key, §or_index, &progress);
474 }
475 }
476}
477
478async fn get_free_plotter_instance<PS>(
480 nats_client: &NatsClient,
481 progress_updater: &ProgressUpdater,
482 progress_sender: &mut PS,
483 retry_backoff_policy: &mut ExponentialBackoff,
484) -> Option<String>
485where
486 PS: Sink<SectorPlottingProgress, Error: Error> + Unpin + Send + 'static,
487{
488 loop {
489 match nats_client
490 .request(&ClusterPlotterFreeInstanceRequest, None)
491 .await
492 {
493 Ok(Some(free_instance)) => {
494 return Some(free_instance);
495 }
496 Ok(None) => {
497 let Some(delay) = retry_backoff_policy.next() else {
498 progress_updater
499 .update_progress_and_events(
500 progress_sender,
501 SectorPlottingProgress::Error {
502 error: "Instance was occupied, exiting #1".to_string(),
503 },
504 )
505 .await;
506 return None;
507 };
508
509 debug!("Instance was occupied, retrying #1");
510
511 tokio::time::sleep(delay).await;
512 }
513 Err(error) => match error.kind() {
514 RequestErrorKind::TimedOut => {
515 let Some(delay) = retry_backoff_policy.next() else {
516 progress_updater
517 .update_progress_and_events(
518 progress_sender,
519 SectorPlottingProgress::Error {
520 error: "Plotter request timed out, exiting".to_string(),
521 },
522 )
523 .await;
524 return None;
525 };
526
527 debug!("Plotter request timed out, retrying");
528
529 tokio::time::sleep(delay).await;
530 }
531 RequestErrorKind::NoResponders => {
532 let Some(delay) = retry_backoff_policy.next() else {
533 progress_updater
534 .update_progress_and_events(
535 progress_sender,
536 SectorPlottingProgress::Error {
537 error: "No plotters, exiting".to_string(),
538 },
539 )
540 .await;
541 return None;
542 };
543
544 debug!("No plotters, retrying");
545
546 tokio::time::sleep(delay).await;
547 }
548 RequestErrorKind::InvalidSubject
549 | RequestErrorKind::MaxPayloadExceeded
550 | RequestErrorKind::Other => {
551 progress_updater
552 .update_progress_and_events(
553 progress_sender,
554 SectorPlottingProgress::Error {
555 error: format!("Failed to get free plotter instance: {error}"),
556 },
557 )
558 .await;
559 return None;
560 }
561 },
562 }
563 }
564}
565
566enum ResponseProcessingResult {
567 Retry,
568 Abort,
569 Continue,
570}
571
572#[expect(clippy::too_many_arguments, reason = "Internal API")]
573async fn process_response_notification<PS>(
574 start: &Instant,
575 free_instance: &str,
576 progress_updater: &ProgressUpdater,
577 progress_sender: &mut PS,
578 retry_backoff_policy: &mut ExponentialBackoff,
579 response: ClusterSectorPlottingProgress,
580 sector_sender: &mut mpsc::Sender<Result<Bytes, String>>,
581 maybe_sector_receiver: &mut Option<mpsc::Receiver<Result<Bytes, String>>>,
582) -> ResponseProcessingResult
583where
584 PS: Sink<SectorPlottingProgress, Error: Error> + Unpin + Send + 'static,
585{
586 if matches!(response, ClusterSectorPlottingProgress::SectorChunk(_)) {
587 trace!("Processing plotting response notification (sector chunk)");
588 } else {
589 trace!(?response, "Processing plotting response notification");
590 }
591
592 match response {
593 ClusterSectorPlottingProgress::Occupied => {
594 debug!(%free_instance, "Instance was occupied, retrying #2");
595
596 return if let Some(delay) = retry_backoff_policy.next() {
597 debug!("Instance was occupied, retrying #2");
598
599 tokio::time::sleep(delay).await;
600 ResponseProcessingResult::Retry
601 } else {
602 debug!("Instance was occupied, exiting #2");
603 ResponseProcessingResult::Abort
604 };
605 }
606 ClusterSectorPlottingProgress::Ping => {
607 }
609 ClusterSectorPlottingProgress::Downloading => {
610 if !progress_updater
611 .update_progress_and_events(progress_sender, SectorPlottingProgress::Downloading)
612 .await
613 {
614 return ResponseProcessingResult::Abort;
615 }
616 }
617 ClusterSectorPlottingProgress::Downloaded(time) => {
618 if !progress_updater
619 .update_progress_and_events(
620 progress_sender,
621 SectorPlottingProgress::Downloaded(time),
622 )
623 .await
624 {
625 return ResponseProcessingResult::Abort;
626 }
627 }
628 ClusterSectorPlottingProgress::Encoding => {
629 if !progress_updater
630 .update_progress_and_events(progress_sender, SectorPlottingProgress::Encoding)
631 .await
632 {
633 return ResponseProcessingResult::Abort;
634 }
635 }
636 ClusterSectorPlottingProgress::Encoded(time) => {
637 if !progress_updater
638 .update_progress_and_events(progress_sender, SectorPlottingProgress::Encoded(time))
639 .await
640 {
641 return ResponseProcessingResult::Abort;
642 }
643 }
644 ClusterSectorPlottingProgress::Finished {
645 plotted_sector,
646 time: _,
647 } => {
648 let Some(sector_receiver) = maybe_sector_receiver.take() else {
649 debug!("Unexpected duplicated sector plotting progress Finished");
650
651 progress_updater
652 .update_progress_and_events(
653 progress_sender,
654 SectorPlottingProgress::Error {
655 error: "Unexpected duplicated sector plotting progress Finished"
656 .to_string(),
657 },
658 )
659 .await;
660 return ResponseProcessingResult::Abort;
661 };
662
663 let progress = SectorPlottingProgress::Finished {
664 plotted_sector,
665 time: start.elapsed(),
667 sector: Box::pin(sector_receiver),
668 };
669 if !progress_updater
670 .update_progress_and_events(progress_sender, progress)
671 .await
672 {
673 return ResponseProcessingResult::Abort;
674 }
675
676 return ResponseProcessingResult::Continue;
677 }
678 ClusterSectorPlottingProgress::SectorChunk(maybe_sector_chunk) => {
680 if let Err(error) = sector_sender.send(maybe_sector_chunk).await {
681 warn!(%error, "Failed to send sector chunk");
682 return ResponseProcessingResult::Abort;
683 }
684 return ResponseProcessingResult::Continue;
685 }
686 ClusterSectorPlottingProgress::Error { error } => {
687 if !progress_updater
688 .update_progress_and_events(
689 progress_sender,
690 SectorPlottingProgress::Error { error },
691 )
692 .await
693 {
694 return ResponseProcessingResult::Abort;
695 }
696 }
697 }
698
699 ResponseProcessingResult::Continue
700}
701
702struct ProgressUpdater {
703 public_key: Ed25519PublicKey,
704 sector_index: SectorIndex,
705 handlers: Arc<Handlers>,
706}
707
708impl ProgressUpdater {
709 async fn update_progress_and_events<PS>(
711 &self,
712 progress_sender: &mut PS,
713 progress: SectorPlottingProgress,
714 ) -> bool
715 where
716 PS: Sink<SectorPlottingProgress, Error: Error> + Unpin,
717 {
718 self.handlers.plotting_progress.call_simple(
719 &self.public_key,
720 &self.sector_index,
721 &progress,
722 );
723
724 if let Err(error) = progress_sender.send(progress).await {
725 warn!(%error, "Failed to send error progress update");
726
727 false
728 } else {
729 true
730 }
731 }
732}
733
734pub async fn plotter_service<P>(nats_client: &NatsClient, plotter: &P) -> anyhow::Result<()>
739where
740 P: Plotter + Sync,
741{
742 let plotter_id = ClusterPlotterId::new();
743
744 select! {
745 result = free_instance_responder(&plotter_id, nats_client, plotter).fuse() => {
746 result
747 }
748 result = plot_sector_responder(&plotter_id, nats_client, plotter).fuse() => {
749 result
750 }
751 }
752}
753
754async fn free_instance_responder<P>(
755 plotter_id: &ClusterPlotterId,
756 nats_client: &NatsClient,
757 plotter: &P,
758) -> anyhow::Result<()>
759where
760 P: Plotter + Sync,
761{
762 loop {
763 while !plotter.has_free_capacity().await.unwrap_or_default() {
764 tokio::time::sleep(FREE_CAPACITY_CHECK_INTERVAL).await;
765 }
766
767 let mut subscription = nats_client
768 .queue_subscribe(
769 ClusterPlotterFreeInstanceRequest::SUBJECT,
770 "ab.plotter".to_string(),
771 )
772 .await
773 .map_err(|error| anyhow!("Failed to subscribe to free instance requests: {error}"))?;
774 debug!(?subscription, "Free instance subscription");
775
776 while let Some(message) = subscription.next().await {
777 let Some(reply_subject) = message.reply else {
778 continue;
779 };
780
781 debug!(%reply_subject, "Free instance request");
782
783 let has_free_capacity = plotter.has_free_capacity().await.unwrap_or_default();
784 let response: <ClusterPlotterFreeInstanceRequest as GenericRequest>::Response =
785 has_free_capacity.then(|| plotter_id.to_string());
786
787 if let Err(error) = nats_client
788 .publish(reply_subject, response.encode().into())
789 .await
790 {
791 warn!(%error, "Failed to send free instance response");
792 }
793
794 if !has_free_capacity {
795 subscription.unsubscribe().await.map_err(|error| {
796 anyhow!("Failed to unsubscribe from free instance requests: {error}")
797 })?;
798 }
799 }
800 }
801}
802
803async fn plot_sector_responder<P>(
804 plotter_id: &ClusterPlotterId,
805 nats_client: &NatsClient,
806 plotter: &P,
807) -> anyhow::Result<()>
808where
809 P: Plotter + Sync,
810{
811 let plotter_id_string = plotter_id.to_string();
812
813 nats_client
814 .stream_request_responder(
815 Some(&plotter_id_string),
816 Some(plotter_id_string.clone()),
817 |request| async move {
818 let (progress_sender, mut progress_receiver) = mpsc::channel(10);
819
820 let fut =
821 process_plot_sector_request(nats_client, plotter, request, progress_sender);
822 let mut fut = Box::pin(fut.fuse());
823
824 Some(
825 stream::poll_fn(move |cx| {
827 if !fut.is_terminated() {
828 let _: Poll<()> = fut.poll_unpin(cx);
830 }
831
832 if let Poll::Ready(maybe_result) = progress_receiver.poll_next_unpin(cx) {
833 return Poll::Ready(maybe_result);
834 }
835
836 Poll::Pending
838 }),
839 )
840 },
841 )
842 .await
843}
844
845async fn process_plot_sector_request<P>(
846 nats_client: &NatsClient,
847 plotter: &P,
848 request: ClusterPlotterPlotSectorRequest,
849 mut response_proxy_sender: mpsc::Sender<ClusterSectorPlottingProgress>,
850) where
851 P: Plotter,
852{
853 let ClusterPlotterPlotSectorRequest {
854 public_key,
855 shard_commitments_root,
856 sector_index,
857 farmer_protocol_info,
858 pieces_in_sector,
859 } = request;
860
861 let inner_fut = async {
863 info!("Plot sector request");
864
865 let (progress_sender, mut progress_receiver) = mpsc::channel(1);
866
867 if !plotter
868 .try_plot_sector(
869 public_key,
870 shard_commitments_root,
871 sector_index,
872 farmer_protocol_info,
873 pieces_in_sector,
874 false,
875 progress_sender,
876 )
877 .await
878 {
879 debug!("Plotter is currently occupied and can't plot more sectors");
880
881 if let Err(error) = response_proxy_sender
882 .send(ClusterSectorPlottingProgress::Occupied)
883 .await
884 {
885 warn!(%error, "Failed to send plotting progress");
886 return;
887 }
888 return;
889 }
890
891 let progress_proxy_fut = {
892 let mut response_proxy_sender = response_proxy_sender.clone();
893 let sector_chunk_size = sector_chunk_size(nats_client);
894
895 async move {
896 while let Some(progress) = progress_receiver.next().await {
897 send_publish_progress(&mut response_proxy_sender, progress, sector_chunk_size)
898 .await;
899 }
900 }
901 };
902
903 let mut ping_interval = tokio::time::interval(PING_INTERVAL);
904 ping_interval.set_missed_tick_behavior(MissedTickBehavior::Delay);
905 let ping_fut = async {
906 loop {
907 ping_interval.tick().await;
908 if let Err(error) = response_proxy_sender
909 .send(ClusterSectorPlottingProgress::Ping)
910 .await
911 {
912 warn!(%error, "Failed to send plotting ping");
913 return;
914 }
915 }
916 };
917
918 select! {
919 () = progress_proxy_fut.fuse() => {
920 }
922 () = ping_fut.fuse() => {
923 unreachable!("Ping loop never ends");
924 }
925 }
926
927 info!("Finished plotting sector successfully");
928 };
929
930 inner_fut
931 .instrument(info_span!("", %public_key, %sector_index))
932 .await;
933}
934
935async fn send_publish_progress(
936 response_sender: &mut mpsc::Sender<ClusterSectorPlottingProgress>,
937 progress: SectorPlottingProgress,
938 sector_chunk_size: usize,
939) {
940 let cluster_progress = match progress {
942 SectorPlottingProgress::Downloading => ClusterSectorPlottingProgress::Downloading,
943 SectorPlottingProgress::Downloaded(time) => ClusterSectorPlottingProgress::Downloaded(time),
944 SectorPlottingProgress::Encoding => ClusterSectorPlottingProgress::Encoding,
945 SectorPlottingProgress::Encoded(time) => ClusterSectorPlottingProgress::Encoded(time),
946 SectorPlottingProgress::Finished {
947 plotted_sector,
948 time,
949 mut sector,
950 } => {
951 if let Err(error) = response_sender
952 .send(ClusterSectorPlottingProgress::Finished {
953 plotted_sector,
954 time,
955 })
956 .await
957 {
958 warn!(%error, "Failed to send plotting progress");
959 return;
960 }
961
962 while let Some(maybe_sector_chunk) = sector.next().await {
963 match maybe_sector_chunk {
964 Ok(sector_chunk) => {
965 for small_sector_chunk in sector_chunk.chunks(sector_chunk_size) {
967 if let Err(error) = response_sender
968 .send(ClusterSectorPlottingProgress::SectorChunk(Ok(
969 sector_chunk.slice_ref(small_sector_chunk)
970 )))
971 .await
972 {
973 warn!(%error, "Failed to send plotting progress");
974 return;
975 }
976 }
977 }
978 Err(error) => {
979 if let Err(error) = response_sender
980 .send(ClusterSectorPlottingProgress::SectorChunk(Err(error)))
981 .await
982 {
983 warn!(%error, "Failed to send plotting progress");
984 return;
985 }
986 }
987 }
988 }
989
990 return;
991 }
992 SectorPlottingProgress::Error { error } => ClusterSectorPlottingProgress::Error { error },
993 };
994
995 if let Err(error) = response_sender.send(cluster_progress).await {
996 warn!(%error, "Failed to send plotting progress");
997 }
998}