Skip to main content

subspace_farmer/cluster/
plotter.rs

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