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_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);
45/// Intervals between pings from plotter server to client
46const PING_INTERVAL: Duration = Duration::from_secs(10);
47/// Timeout after which plotter that doesn't send pings is assumed to be down
48const PING_TIMEOUT: Duration = Duration::from_mins(1);
49
50/// Type alias used for event handlers
51pub 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/// An ephemeral identifier for a plotter
55#[derive(Debug, Copy, Clone, Ord, PartialOrd, Eq, PartialEq, Hash, Display)]
56pub enum ClusterPlotterId {
57    /// Plotter ID
58    Ulid(Ulid),
59}
60
61#[expect(
62    clippy::new_without_default,
63    reason = "Default has different semantics"
64)]
65impl ClusterPlotterId {
66    /// Creates new ID
67    pub fn new() -> Self {
68        Self::Ulid(Ulid::generate())
69    }
70}
71
72/// Request for free plotter instance
73#[derive(Debug, Clone, Encode, Decode)]
74struct ClusterPlotterFreeInstanceRequest;
75
76impl GenericRequest for ClusterPlotterFreeInstanceRequest {
77    const SUBJECT: &'static str = "ab.plotter.free-instance";
78    /// Might be `None` if instance had to respond, but turned out it was fully occupied already
79    type Response = Option<String>;
80}
81
82#[derive(Debug, Encode, Decode)]
83enum ClusterSectorPlottingProgress {
84    /// Plotter is already fully occupied with other work
85    Occupied,
86    /// Periodic ping indicating plotter is still busy
87    Ping,
88    /// Downloading sector pieces
89    Downloading,
90    /// Downloaded sector pieces
91    Downloaded(Duration),
92    /// Encoding sector pieces
93    Encoding,
94    /// Encoded sector pieces
95    Encoded(Duration),
96    /// Finished plotting, followed by a series of sector chunks
97    Finished {
98        /// Information about plotted sector
99        plotted_sector: PlottedSector,
100        /// How much time it took to plot a sector
101        time: Duration,
102    },
103    /// Sector chunk after finished plotting
104    SectorChunk(Result<Bytes, String>),
105    /// Plotting failed
106    Error {
107        /// Error message
108        error: String,
109    },
110}
111
112/// Request to plot sector from plotter
113#[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
127/// Size of sector chunks sent from plotter to farmer.
128///
129/// It is a multiple of disk page size, such that farmer can write chunks to disk without
130/// read-modify-write.
131fn 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/// Cluster plotter
141#[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        // Done outside the future below as a backpressure, ensuring that it is not possible to
183        // schedule unbounded number of plotting tasks
184        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    /// Create a new instance
260    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        // Basically runs plotting tasks in the background and allows to abort on drop
271        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                // Just so that `FuturesUnordered` will never end
276                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                            // Nothing to do
289                        }
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    /// Subscribe to plotting progress notifications
307    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        // Try to get plotter instance here first as a backpressure measure
339        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                // Take free instance that was found earlier if available or try to find a new one
357                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                // Allow to buffer up to the whole sector in memory to not block plotter on the
404                // other side
405                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                                    // Nothing to do
435                                }
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, &sector_index, &progress);
474        }
475    }
476}
477
478// Try to get free plotter instance and return `None` if it is not possible
479async 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            // Expected
608        }
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                // Use local time instead of reported by remote plotter
666                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        // This variant must be sent after Finished and it handled above
679        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    /// Returns `true` on success and `false` if progress receiver channel is gone
710    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
734/// Create plotter service that will be processing incoming requests.
735///
736/// Implementation is using concurrency with multiple tokio tasks, but can be started multiple times
737/// per controller instance in order to parallelize more work across threads if needed.
738pub 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                    // Drive above future and stream back any pieces that were downloaded so far
826                    stream::poll_fn(move |cx| {
827                        if !fut.is_terminated() {
828                            // Result doesn't matter, we'll need to poll stream below anyway
829                            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                        // Exit will be done by the stream above
837                        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    // Wrapper future just for instrumentation below
862    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                // 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    sector_chunk_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(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}