ab_farmer/cluster/

controller.rs

//! Farming cluster controller
//!
//! Controller is responsible for managing farming cluster.
//!
//! This module exposes some data structures for NATS communication, custom piece getter and node
//! client implementations designed to work with cluster controller and a service function to drive
//! the backend part of the controller.

pub mod caches;
pub mod farms;
mod stream_map;

use crate::cluster::cache::{ClusterCacheReadPieceRequest, ClusterCacheReadPiecesRequest};
use crate::cluster::nats_client::{
    GenericBroadcast, GenericNotification, GenericRequest, GenericStreamRequest, NatsClient,
};
use crate::farm::{PieceCacheId, PieceCacheOffset};
use crate::farmer_cache::FarmerCache;
use crate::node_client::NodeClient;
use ab_core_primitives::pieces::{Piece, PieceIndex};
use ab_core_primitives::segments::{
    SegmentIndex, SuperSegmentHeader, SuperSegmentIndex, SuperSegmentRoot,
};
use ab_data_retrieval::piece_getter::PieceGetter;
use ab_farmer_rpc_primitives::{
    BlockSealInfo, BlockSealResponse, FarmerAppInfo, FarmerShardMembershipInfo, SlotInfo,
    SolutionResponse,
};
use anyhow::anyhow;
use async_nats::HeaderValue;
use async_trait::async_trait;
use futures::channel::mpsc;
use futures::future::FusedFuture;
use futures::stream::FuturesUnordered;
use futures::{FutureExt, Stream, StreamExt, select, stream};
use parity_scale_codec::{Decode, Encode};
use parking_lot::Mutex;
use rand::prelude::*;
use std::collections::{HashMap, HashSet};
use std::pin::Pin;
use std::task::Poll;
use tracing::{debug, error, trace, warn};
use triomphe::Arc;

/// Special "cache group" that all controllers subscribe to and that can be used to query any cache
/// group. The cache group for each query is chosen at random.
const GLOBAL_CACHE_GROUP: &str = "_";

/// Broadcast sent by controllers requesting farmers to identify themselves
#[derive(Debug, Copy, Clone, Encode, Decode)]
pub struct ClusterControllerFarmerIdentifyBroadcast;

impl GenericBroadcast for ClusterControllerFarmerIdentifyBroadcast {
    const SUBJECT: &'static str = "ab.controller.farmer-identify";
}

/// Broadcast sent by controllers requesting caches in cache group to identify themselves
#[derive(Debug, Copy, Clone, Encode, Decode)]
pub struct ClusterControllerCacheIdentifyBroadcast;

impl GenericBroadcast for ClusterControllerCacheIdentifyBroadcast {
    /// `*` here stands for cache group
    const SUBJECT: &'static str = "ab.controller.*.cache-identify";
}

/// Broadcast with slot info sent by controllers
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerSlotInfoBroadcast {
    slot_info: SlotInfo,
    instance: String,
}

impl GenericBroadcast for ClusterControllerSlotInfoBroadcast {
    const SUBJECT: &'static str = "ab.controller.slot-info";

    fn deterministic_message_id(&self) -> Option<HeaderValue> {
        // TODO: Depending on answer in `https://github.com/nats-io/nats.docs/issues/663` this might
        //  be simplified to just a slot number
        Some(
            HeaderValue::try_from(format!("slot-info-{}", self.slot_info.slot))
                .expect("No `\r` or `\n`; qed"),
        )
    }
}

/// Broadcast with block sealing info by controllers
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerBlockSealingBroadcast {
    block_sealing_info: BlockSealInfo,
}

impl GenericBroadcast for ClusterControllerBlockSealingBroadcast {
    const SUBJECT: &'static str = "ab.controller.block-sealing-info";
}

/// Broadcast with new super segment headers by controllers
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerNewSuperSegmentHeaderBroadcast {
    new_super_segment_header: SuperSegmentHeader,
}

impl GenericBroadcast for ClusterControllerNewSuperSegmentHeaderBroadcast {
    const SUBJECT: &'static str = "ab.controller.new-super-segment-header";

    fn deterministic_message_id(&self) -> Option<HeaderValue> {
        // TODO: Depending on answer in `https://github.com/nats-io/nats.docs/issues/663` this might
        //  be simplified to just a super segment index
        Some(
            HeaderValue::try_from(format!(
                "new-super-segment-{}",
                self.new_super_segment_header.index
            ))
            .expect("No `\r` or `\n`; qed"),
        )
    }
}

/// Notification messages with solution by farmers
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerSolutionNotification {
    solution_response: SolutionResponse,
}

impl GenericNotification for ClusterControllerSolutionNotification {
    const SUBJECT: &'static str = "ab.controller.*.solution";
}

/// Notification messages with block seal by farmers
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerBlockSealNotification {
    block_seal: BlockSealResponse,
}

impl GenericNotification for ClusterControllerBlockSealNotification {
    const SUBJECT: &'static str = "ab.controller.block-seal";
}

/// Request farmer app info from controller
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerFarmerAppInfoRequest;

impl GenericRequest for ClusterControllerFarmerAppInfoRequest {
    const SUBJECT: &'static str = "ab.controller.farmer-app-info";
    type Response = Result<FarmerAppInfo, String>;
}

/// Request super segment headers with specified super segment indices
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerSuperSegmentHeadersRequest {
    super_segment_indices: Vec<SuperSegmentIndex>,
}

impl GenericRequest for ClusterControllerSuperSegmentHeadersRequest {
    const SUBJECT: &'static str = "ab.controller.super-segment-headers";
    type Response = Vec<Option<SuperSegmentHeader>>;
}

/// Request super segment root for a specified segment index
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerSuperSegmentRootForSegmentIndexRequest {
    segment_index: SegmentIndex,
}

impl GenericRequest for ClusterControllerSuperSegmentRootForSegmentIndexRequest {
    const SUBJECT: &'static str = "ab.controller.super-segment-root-for-segment-index";
    type Response = Option<SuperSegmentRoot>;
}

/// Find piece with specified index in cache
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerFindPieceInCacheRequest {
    piece_index: PieceIndex,
}

impl GenericRequest for ClusterControllerFindPieceInCacheRequest {
    const SUBJECT: &'static str = "ab.controller.*.find-piece-in-cache";
    type Response = Option<(PieceCacheId, PieceCacheOffset)>;
}

/// Find pieces with specified indices in cache
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerFindPiecesInCacheRequest {
    piece_indices: Vec<PieceIndex>,
}

impl GenericStreamRequest for ClusterControllerFindPiecesInCacheRequest {
    const SUBJECT: &'static str = "ab.controller.*.find-pieces-in-cache";
    /// Only pieces that were found are returned
    type Response = (PieceIndex, PieceCacheId, PieceCacheOffset);
}

/// Request piece with specified index
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerPieceRequest {
    piece_index: PieceIndex,
}

impl GenericRequest for ClusterControllerPieceRequest {
    const SUBJECT: &'static str = "ab.controller.piece";
    type Response = Option<Piece>;
}

/// Request pieces with specified index
#[derive(Debug, Clone, Encode, Decode)]
struct ClusterControllerPiecesRequest {
    piece_indices: Vec<PieceIndex>,
}

impl GenericStreamRequest for ClusterControllerPiecesRequest {
    const SUBJECT: &'static str = "ab.controller.pieces";
    /// Only pieces that were found are returned
    type Response = (PieceIndex, Piece);
}

/// Cluster piece getter
#[derive(Debug, Clone)]
pub struct ClusterPieceGetter {
    nats_client: NatsClient,
    cache_group: String,
}

#[async_trait]
impl PieceGetter for ClusterPieceGetter {
    async fn get_piece(&self, piece_index: PieceIndex) -> anyhow::Result<Option<Piece>> {
        if let Some((piece_cache_id, piece_cache_offset)) = self
            .nats_client
            .request(
                &ClusterControllerFindPieceInCacheRequest { piece_index },
                Some(&self.cache_group),
            )
            .await?
        {
            trace!(
                %piece_index,
                %piece_cache_id,
                %piece_cache_offset,
                "Found piece in cache, retrieving"
            );

            match self
                .nats_client
                .request(
                    &ClusterCacheReadPieceRequest {
                        offset: piece_cache_offset,
                    },
                    Some(&piece_cache_id.to_string()),
                )
                .await
                .map_err(|error| error.to_string())
                .flatten()
            {
                Ok(Some((retrieved_piece_index, piece))) => {
                    if retrieved_piece_index == piece_index {
                        trace!(
                            %piece_index,
                            %piece_cache_id,
                            %piece_cache_offset,
                            "Retrieved piece from cache successfully"
                        );

                        return Ok(Some(piece));
                    }

                    trace!(
                        %piece_index,
                        %piece_cache_id,
                        %piece_cache_offset,
                        "Retrieving piece was replaced in cache during retrieval"
                    );
                }
                Ok(None) => {
                    trace!(
                        %piece_index,
                        %piece_cache_id,
                        %piece_cache_offset,
                        "Piece cache didn't have piece at offset"
                    );
                }
                Err(error) => {
                    debug!(
                        %piece_index,
                        %piece_cache_id,
                        %piece_cache_offset,
                        %error,
                        "Retrieving piece from cache failed"
                    );
                }
            }
        } else {
            trace!(%piece_index, "Piece not found in cache");
        }

        Ok(self
            .nats_client
            .request(&ClusterControllerPieceRequest { piece_index }, None)
            .await?)
    }

    async fn get_pieces<'a>(
        &'a self,
        piece_indices: Vec<PieceIndex>,
    ) -> anyhow::Result<
        Box<dyn Stream<Item = (PieceIndex, anyhow::Result<Option<Piece>>)> + Send + Unpin + 'a>,
    > {
        let (tx, mut rx) = mpsc::unbounded();

        let piece_indices_to_get =
            Mutex::new(piece_indices.iter().copied().collect::<HashSet<_>>());

        let mut cached_pieces_by_cache_id = HashMap::<PieceCacheId, Vec<PieceCacheOffset>>::new();

        {
            let mut cached_pieces = self
                .nats_client
                .stream_request(
                    &ClusterControllerFindPiecesInCacheRequest { piece_indices },
                    Some(&self.cache_group),
                )
                .await?;

            while let Some((_piece_index, piece_cache_id, piece_cache_offset)) =
                cached_pieces.next().await
            {
                cached_pieces_by_cache_id
                    .entry(piece_cache_id)
                    .or_default()
                    .push(piece_cache_offset);
            }
        }

        let fut = async move {
            let tx = &tx;

            cached_pieces_by_cache_id
                .into_iter()
                .map(|(piece_cache_id, offsets)| {
                    let piece_indices_to_get = &piece_indices_to_get;

                    async move {
                        let mut pieces_stream = match self
                            .nats_client
                            .stream_request(
                                &ClusterCacheReadPiecesRequest { offsets },
                                Some(&piece_cache_id.to_string()),
                            )
                            .await
                        {
                            Ok(pieces) => pieces,
                            Err(error) => {
                                warn!(
                                    %error,
                                    %piece_cache_id,
                                    "Failed to request pieces from cache"
                                );

                                return;
                            }
                        };

                        while let Some(piece_result) = pieces_stream.next().await {
                            let (piece_offset, maybe_piece) = match piece_result {
                                Ok(result) => result,
                                Err(error) => {
                                    warn!(%error, "Failed to get piece from cache");
                                    continue;
                                }
                            };

                            if let Some((piece_index, piece)) = maybe_piece {
                                piece_indices_to_get.lock().remove(&piece_index);

                                tx.unbounded_send((piece_index, Ok(Some(piece)))).expect(
                                    "This future isn't polled after receiver is dropped; qed",
                                );
                            } else {
                                warn!(
                                    %piece_cache_id,
                                    %piece_offset,
                                    "Failed to get piece from cache, it was missing or already gone"
                                );
                            }
                        }
                    }
                })
                .collect::<FuturesUnordered<_>>()
                // Simply drain everything
                .for_each(|()| async {})
                .await;

            let mut piece_indices_to_get = piece_indices_to_get.into_inner();
            if piece_indices_to_get.is_empty() {
                return;
            }

            let mut pieces_from_controller = match self
                .nats_client
                .stream_request(
                    &ClusterControllerPiecesRequest {
                        piece_indices: piece_indices_to_get.iter().copied().collect(),
                    },
                    None,
                )
                .await
            {
                Ok(pieces_from_controller) => pieces_from_controller,
                Err(error) => {
                    error!(%error, "Failed to get pieces from controller");

                    for piece_index in piece_indices_to_get {
                        tx.unbounded_send((
                            piece_index,
                            Err(anyhow::anyhow!("Failed to get piece from controller")),
                        ))
                        .expect("This future isn't polled after receiver is dropped; qed");
                    }
                    return;
                }
            };

            while let Some((piece_index, piece)) = pieces_from_controller.next().await {
                piece_indices_to_get.remove(&piece_index);
                tx.unbounded_send((piece_index, Ok(Some(piece))))
                    .expect("This future isn't polled after receiver is dropped; qed");
            }

            for piece_index in piece_indices_to_get {
                tx.unbounded_send((piece_index, Err(anyhow::anyhow!("Failed to get piece"))))
                    .expect("This future isn't polled after receiver is dropped; qed");
            }
        };
        let mut fut = Box::pin(fut.fuse());

        // Drive above future and stream back any pieces that were downloaded so far
        Ok(Box::new(stream::poll_fn(move |cx| {
            if !fut.is_terminated() {
                // Result doesn't matter, we'll need to poll stream below anyway
                let _: Poll<()> = fut.poll_unpin(cx);
            }

            if let Poll::Ready(maybe_result) = rx.poll_next_unpin(cx) {
                return Poll::Ready(maybe_result);
            }

            // Exit will be done by the stream above
            Poll::Pending
        })))
    }
}

impl ClusterPieceGetter {
    /// Create a new instance
    #[inline]
    pub fn new(nats_client: NatsClient, cache_group: Option<String>) -> Self {
        Self {
            nats_client,
            cache_group: cache_group.unwrap_or_else(|| GLOBAL_CACHE_GROUP.to_string()),
        }
    }
}

/// [`NodeClient`] used in cluster environment that connects to node through a controller instead
/// of to the node directly
#[derive(Debug, Clone)]
pub struct ClusterNodeClient {
    nats_client: NatsClient,
    // Store last slot info instance that can be used to send solution response to (some instances
    // may be not synced and not able to receive solution responses)
    last_slot_info_instance: Arc<Mutex<String>>,
}

impl ClusterNodeClient {
    /// Create a new instance
    pub fn new(nats_client: NatsClient) -> Self {
        Self {
            nats_client,
            last_slot_info_instance: Arc::default(),
        }
    }
}

#[async_trait]
impl NodeClient for ClusterNodeClient {
    async fn farmer_app_info(&self) -> anyhow::Result<FarmerAppInfo> {
        Ok(self
            .nats_client
            .request(&ClusterControllerFarmerAppInfoRequest, None)
            .await?
            .map_err(anyhow::Error::msg)?)
    }

    async fn subscribe_slot_info(
        &self,
    ) -> anyhow::Result<Pin<Box<dyn Stream<Item = SlotInfo> + Send + 'static>>> {
        let subscription = self
            .nats_client
            .subscribe_to_broadcasts::<ClusterControllerSlotInfoBroadcast>(None, None)
            .await?
            .filter_map({
                let mut last_slot_number = None;
                let last_slot_info_instance = Arc::clone(&self.last_slot_info_instance);

                move |broadcast| {
                    let slot_info = broadcast.slot_info;

                    let maybe_slot_info = if let Some(last_slot_number) = last_slot_number
                        && last_slot_number >= slot_info.slot
                    {
                        None
                    } else {
                        last_slot_number.replace(slot_info.slot);
                        *last_slot_info_instance.lock() = broadcast.instance;

                        Some(slot_info)
                    };

                    async move { maybe_slot_info }
                }
            });

        Ok(Box::pin(subscription))
    }

    async fn submit_solution_response(
        &self,
        solution_response: SolutionResponse,
    ) -> anyhow::Result<()> {
        let last_slot_info_instance = self.last_slot_info_instance.lock().clone();
        Ok(self
            .nats_client
            .notification(
                &ClusterControllerSolutionNotification { solution_response },
                Some(&last_slot_info_instance),
            )
            .await?)
    }

    async fn subscribe_block_sealing(
        &self,
    ) -> anyhow::Result<Pin<Box<dyn Stream<Item = BlockSealInfo> + Send + 'static>>> {
        let subscription = self
            .nats_client
            .subscribe_to_broadcasts::<ClusterControllerBlockSealingBroadcast>(None, None)
            .await?
            .map(|broadcast| broadcast.block_sealing_info);

        Ok(Box::pin(subscription))
    }

    async fn submit_block_seal(&self, block_seal: BlockSealResponse) -> anyhow::Result<()> {
        Ok(self
            .nats_client
            .notification(&ClusterControllerBlockSealNotification { block_seal }, None)
            .await?)
    }

    async fn subscribe_new_super_segment_headers(
        &self,
    ) -> anyhow::Result<Pin<Box<dyn Stream<Item = SuperSegmentHeader> + Send + 'static>>> {
        let subscription = self
            .nats_client
            .subscribe_to_broadcasts::<ClusterControllerNewSuperSegmentHeaderBroadcast>(None, None)
            .await?
            .filter_map({
                let mut last_super_segment_index = None;

                move |broadcast| {
                    let new_super_segment_header = broadcast.new_super_segment_header;
                    let super_segment_index = new_super_segment_header.index.as_inner();

                    let maybe_super_segment_header = if let Some(last_super_segment_index) =
                        last_super_segment_index
                        && last_super_segment_index >= super_segment_index
                    {
                        None
                    } else {
                        last_super_segment_index.replace(super_segment_index);

                        Some(new_super_segment_header)
                    };

                    async move { maybe_super_segment_header }
                }
            });

        Ok(Box::pin(subscription))
    }

    async fn super_segment_headers(
        &self,
        super_segment_indices: Vec<SuperSegmentIndex>,
    ) -> anyhow::Result<Vec<Option<SuperSegmentHeader>>> {
        Ok(self
            .nats_client
            .request(
                &ClusterControllerSuperSegmentHeadersRequest {
                    super_segment_indices,
                },
                None,
            )
            .await?)
    }

    async fn super_segment_root_for_segment_index(
        &self,
        segment_index: SegmentIndex,
    ) -> anyhow::Result<Option<SuperSegmentRoot>> {
        Ok(self
            .nats_client
            .request(
                &ClusterControllerSuperSegmentRootForSegmentIndexRequest { segment_index },
                None,
            )
            .await?)
    }

    async fn piece(&self, piece_index: PieceIndex) -> anyhow::Result<Option<Piece>> {
        Ok(self
            .nats_client
            .request(&ClusterControllerPieceRequest { piece_index }, None)
            .await?)
    }

    async fn update_shard_membership_info(
        &self,
        _info: FarmerShardMembershipInfo,
    ) -> anyhow::Result<()> {
        // Controller aggregates these on its own
        Ok(())
    }
}

/// Create controller service that handles things like broadcasting information (for example slot
/// notifications) as well as responding to incoming requests (like piece requests).
///
/// Implementation is using concurrency with multiple tokio tasks, but can be started multiple times
/// per controller instance in order to parallelize more work across threads if needed.
pub async fn controller_service<NC, PG>(
    nats_client: &NatsClient,
    node_client: &NC,
    piece_getter: &PG,
    farmer_caches: &[(&str, &FarmerCache)],
    instance: &str,
    primary_instance: bool,
) -> anyhow::Result<()>
where
    NC: NodeClient,
    PG: PieceGetter + Sync,
{
    if primary_instance {
        select! {
            result = slot_info_broadcaster(nats_client, node_client, instance).fuse() => {
                result
            },
            result = block_sealing_broadcaster(nats_client, node_client, instance).fuse() => {
                result
            },
            result = new_super_segment_headers_broadcaster(nats_client, node_client, instance).fuse() => {
                result
            },
            result = solution_response_forwarder(nats_client, node_client, instance).fuse() => {
                result
            },
            result = block_seal_forwarder(nats_client, node_client, instance).fuse() => {
                result
            },
            result = farmer_app_info_responder(nats_client, node_client).fuse() => {
                result
            },
            result = super_segment_headers_responder(nats_client, node_client).fuse() => {
                result
            },
            result = super_segment_root_for_segment_index_responder(nats_client, node_client).fuse() => {
                result
            },
            result = find_piece_responder(nats_client, farmer_caches).fuse() => {
                result
            },
            result = find_pieces_responder(nats_client, farmer_caches).fuse() => {
                result
            },
            result = piece_responder(nats_client, piece_getter).fuse() => {
                result
            },
            result = pieces_responder(nats_client, piece_getter).fuse() => {
                result
            },
        }
    } else {
        select! {
            result = farmer_app_info_responder(nats_client, node_client).fuse() => {
                result
            },
            result = super_segment_headers_responder(nats_client, node_client).fuse() => {
                result
            },
            result = super_segment_root_for_segment_index_responder(nats_client, node_client).fuse() => {
                result
            },
            result = find_piece_responder(nats_client, farmer_caches).fuse() => {
                result
            },
            result = find_pieces_responder(nats_client, farmer_caches).fuse() => {
                result
            },
            result = piece_responder(nats_client, piece_getter).fuse() => {
                result
            },
            result = pieces_responder(nats_client, piece_getter).fuse() => {
                result
            },
        }
    }
}

async fn slot_info_broadcaster<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
    instance: &str,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    let mut slot_info_notifications = node_client
        .subscribe_slot_info()
        .await
        .map_err(|error| anyhow!("Failed to subscribe to slot info notifications: {error}"))?;

    while let Some(slot_info) = slot_info_notifications.next().await {
        debug!(?slot_info, "New slot");

        let slot = slot_info.slot;

        if let Err(error) = nats_client
            .broadcast(
                &ClusterControllerSlotInfoBroadcast {
                    slot_info,
                    instance: instance.to_string(),
                },
                instance,
            )
            .await
        {
            warn!(%slot, %error, "Failed to broadcast slot info");
        }
    }

    Ok(())
}

async fn block_sealing_broadcaster<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
    instance: &str,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    let mut block_sealing_subscription = node_client
        .subscribe_block_sealing()
        .await
        .map_err(|error| anyhow!("Failed to subscribe to block sealing notifications: {error}"))?;

    while let Some(block_sealing_info) = block_sealing_subscription.next().await {
        trace!(?block_sealing_info, "New block sealing notification");

        if let Err(error) = nats_client
            .broadcast(
                &ClusterControllerBlockSealingBroadcast { block_sealing_info },
                instance,
            )
            .await
        {
            warn!(%error, "Failed to broadcast block sealing info");
        }
    }

    Ok(())
}

async fn new_super_segment_headers_broadcaster<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
    instance: &str,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    let mut new_super_segments_notifications = node_client
        .subscribe_new_super_segment_headers()
        .await
        .map_err(|error| {
            anyhow!("Failed to subscribe to new super segment header notifications: {error}")
        })?;

    while let Some(new_super_segment_header) = new_super_segments_notifications.next().await {
        trace!(
            ?new_super_segment_header,
            "New archived new super segment header notification"
        );

        if let Err(error) = nats_client
            .broadcast(
                &ClusterControllerNewSuperSegmentHeaderBroadcast {
                    new_super_segment_header,
                },
                instance,
            )
            .await
        {
            warn!(%error, "Failed to broadcast new super segment header info");
        }
    }

    Ok(())
}

async fn solution_response_forwarder<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
    instance: &str,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    let mut subscription = nats_client
        .subscribe_to_notifications::<ClusterControllerSolutionNotification>(
            Some(instance),
            Some(instance.to_string()),
        )
        .await
        .map_err(|error| anyhow!("Failed to subscribe to solution notifications: {error}"))?;

    while let Some(notification) = subscription.next().await {
        debug!(?notification, "Solution notification");

        let slot = notification.solution_response.slot_number;
        let public_key_hash = notification.solution_response.solution.public_key_hash;
        let sector_index = notification.solution_response.solution.sector_index;

        if let Err(error) = node_client
            .submit_solution_response(notification.solution_response)
            .await
        {
            warn!(
                %error,
                %slot,
                %public_key_hash,
                %sector_index,
                "Failed to send solution response"
            );
        }
    }

    Ok(())
}

async fn block_seal_forwarder<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
    instance: &str,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    let mut subscription = nats_client
        .subscribe_to_notifications::<ClusterControllerBlockSealNotification>(
            None,
            Some(instance.to_string()),
        )
        .await
        .map_err(|error| anyhow!("Failed to subscribe to block seal notifications: {error}"))?;

    while let Some(notification) = subscription.next().await {
        debug!(?notification, "Block seal notification");

        if let Err(error) = node_client.submit_block_seal(notification.block_seal).await {
            warn!(%error, "Failed to send block seal");
        }
    }

    Ok(())
}

async fn farmer_app_info_responder<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    nats_client
        .request_responder(
            None,
            Some("ab.controller".to_string()),
            |_: ClusterControllerFarmerAppInfoRequest| async move {
                Some(
                    node_client
                        .farmer_app_info()
                        .await
                        .map_err(|error| error.to_string()),
                )
            },
        )
        .await
}

async fn super_segment_headers_responder<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    nats_client
        .request_responder(
            None,
            Some("ab.controller".to_string()),
            |ClusterControllerSuperSegmentHeadersRequest { super_segment_indices }| async move {
                node_client
                    .super_segment_headers(super_segment_indices.clone())
                    .await
                    .inspect_err(|error| {
                        warn!(%error, ?super_segment_indices, "Failed to get super segment headers");
                    })
                    .ok()
            },
        )
        .await
}

async fn super_segment_root_for_segment_index_responder<NC>(
    nats_client: &NatsClient,
    node_client: &NC,
) -> anyhow::Result<()>
where
    NC: NodeClient,
{
    nats_client
        .request_responder(
            None,
            Some("ab.controller".to_string()),
            |ClusterControllerSuperSegmentRootForSegmentIndexRequest { segment_index }| async move {
                node_client
                    .super_segment_root_for_segment_index(segment_index)
                    .await
                    .inspect_err(|error| {
                        warn!(
                            %error,
                            %segment_index,
                            "Failed to get super segment root for segment index"
                        );
                    })
                    .ok()
            },
        )
        .await
}

async fn find_piece_responder(
    nats_client: &NatsClient,
    farmer_caches: &[(&str, &FarmerCache)],
) -> anyhow::Result<()> {
    futures::future::try_join(
        farmer_caches
            .iter()
            .map(|(cache_group, farmer_cache)| {
                nats_client.request_responder(
                    Some(cache_group),
                    Some("ab.controller".to_string()),
                    move |ClusterControllerFindPieceInCacheRequest { piece_index }| async move {
                        Some(farmer_cache.find_piece(piece_index).await)
                    },
                )
            })
            .collect::<FuturesUnordered<_>>()
            .next()
            .map(|result| result.unwrap_or(Ok(()))),
        nats_client.request_responder(
            Some(GLOBAL_CACHE_GROUP),
            Some("ab.controller".to_string()),
            |ClusterControllerFindPieceInCacheRequest { piece_index }| async move {
                let (_cache_group, farmer_cache) = farmer_caches.iter().choose(&mut rand::rng())?;
                Some(farmer_cache.find_piece(piece_index).await)
            },
        ),
    )
    .await
    .map(|((), ())| ())
}

async fn find_pieces_responder(
    nats_client: &NatsClient,
    farmer_caches: &[(&str, &FarmerCache)],
) -> anyhow::Result<()> {
    futures::future::try_join(
        farmer_caches
            .iter()
            .map(|(cache_group, farmer_cache)| {
                nats_client.stream_request_responder(
                    Some(cache_group),
                    Some("ab.controller".to_string()),
                    move |ClusterControllerFindPiecesInCacheRequest { piece_indices }| async move {
                        Some(stream::iter(farmer_cache.find_pieces(piece_indices).await))
                    },
                )
            })
            .collect::<FuturesUnordered<_>>()
            .next()
            .map(|result| result.unwrap_or(Ok(()))),
        nats_client.stream_request_responder(
            Some(GLOBAL_CACHE_GROUP),
            Some("ab.controller".to_string()),
            |ClusterControllerFindPiecesInCacheRequest { piece_indices }| async move {
                let (_cache_group, farmer_cache) = farmer_caches.iter().choose(&mut rand::rng())?;
                Some(stream::iter(farmer_cache.find_pieces(piece_indices).await))
            },
        ),
    )
    .await
    .map(|((), ())| ())
}

async fn piece_responder<PG>(nats_client: &NatsClient, piece_getter: &PG) -> anyhow::Result<()>
where
    PG: PieceGetter + Sync,
{
    nats_client
        .request_responder(
            None,
            Some("ab.controller".to_string()),
            |ClusterControllerPieceRequest { piece_index }| async move {
                piece_getter
                    .get_piece(piece_index)
                    .await
                    .inspect_err(|error| warn!(%error, %piece_index, "Failed to get piece"))
                    .ok()
            },
        )
        .await
}

async fn pieces_responder<PG>(nats_client: &NatsClient, piece_getter: &PG) -> anyhow::Result<()>
where
    PG: PieceGetter + Sync,
{
    nats_client
        .stream_request_responder(
            None,
            Some("ab.controller".to_string()),
            |ClusterControllerPiecesRequest { piece_indices }| async move {
                piece_getter
                    .get_pieces(piece_indices)
                    .await
                    .map(|stream| {
                        Box::pin(stream.filter_map(
                            |(piece_index, maybe_piece_result)| async move {
                                match maybe_piece_result {
                                    Ok(Some(piece)) => Some((piece_index, piece)),
                                    Ok(None) => None,
                                    Err(error) => {
                                        warn!(%error, %piece_index, "Failed to get piece");
                                        None
                                    }
                                }
                            },
                        ))
                    })
                    .inspect_err(|error| warn!(%error, "Failed to get pieces"))
                    .ok()
            },
        )
        .await
}