Skip to main content

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