ab_farmer_components/

plotting.rs

//! Plotting utilities
//!
//! This module contains functions and data structures that can be used for plotting purposes
//! (primarily with CPU).
//!
//! Plotted sectors can be written to plot and later [`read`](crate::reading) and/or
//! [`audited`](crate::auditing)/[`proven`](crate::proving) using other modules of this crate.

use crate::FarmerProtocolInfo;
use crate::sector::{
    FoundProofs, RawSector, RecordMetadata, SectorContentsMap, SectorMetadata,
    SectorMetadataChecksummed, sector_record_chunks_size, sector_size,
};
use crate::segment_reconstruction::recover_missing_piece;
use ab_core_primitives::hashes::Blake3Hash;
use ab_core_primitives::pieces::{Piece, PieceIndex, PieceOffset, Record, RecordChunk};
use ab_core_primitives::pos::PosSeed;
use ab_core_primitives::sectors::{SBucket, SectorId, SectorIndex};
use ab_core_primitives::segments::HistorySize;
use ab_core_primitives::solutions::ShardCommitmentHash;
use ab_data_retrieval::piece_getter::PieceGetter;
use ab_erasure_coding::ErasureCoding;
use ab_proof_of_space::{Table, TableGenerator};
use async_lock::{Mutex as AsyncMutex, Semaphore};
use backon::{ExponentialBuilder, Retryable};
use futures::stream::FuturesUnordered;
use futures::{StreamExt, select};
use parity_scale_codec::{Decode, Encode};
use parking_lot::Mutex;
use std::collections::HashMap;
use std::simd::Simd;
use std::sync::Arc;
use std::sync::atomic::{AtomicBool, Ordering};
use std::time::Duration;
use thiserror::Error;
use tracing::{debug, trace, warn};

const RECONSTRUCTION_CONCURRENCY_LIMIT: usize = 1;

/// Information about sector that was plotted
#[derive(Debug, Clone, Encode, Decode)]
pub struct PlottedSector {
    /// Sector ID
    pub sector_id: SectorId,
    /// Sector index
    pub sector_index: SectorIndex,
    /// Sector metadata
    pub sector_metadata: SectorMetadataChecksummed,
    /// Indexes of pieces that were plotted
    pub piece_indexes: Vec<PieceIndex>,
}

/// Plotting status
#[derive(Debug, Error)]
pub enum PlottingError {
    /// Records encoder error
    #[error("Records encoder error: {error}")]
    RecordsEncoderError {
        /// Lower-level error
        error: anyhow::Error,
    },
    /// Bad sector output size
    #[error("Bad sector output size: provided {provided}, expected {expected}")]
    BadSectorOutputSize {
        /// Actual size
        provided: usize,
        /// Expected size
        expected: usize,
    },
    /// Can't recover missing piece
    #[error("Can't recover missing piece {piece_index}: {error}")]
    PieceRecoveryFailed {
        /// Piece index
        piece_index: PieceIndex,
        /// Lower-level error
        error: anyhow::Error,
    },
    /// Failed to retrieve piece
    #[error("Failed to retrieve pieces: {error}")]
    FailedToRetrievePieces {
        /// Lower-level error
        error: anyhow::Error,
    },
    /// Abort early
    #[error("Abort early")]
    AbortEarly,
}

/// Options for plotting a sector.
///
/// Sector output and sector metadata output should be either empty (in which case they'll be
/// resized to correct size automatically) or correctly sized from the beginning or else error will
/// be returned.
#[derive(Debug)]
pub struct PlotSectorOptions<'a, RE, PG> {
    /// Public key corresponding to sector
    pub public_key_hash: &'a Blake3Hash,
    /// Root of the Merkle Tree of shard commitments
    pub shard_commitments_root: &'a ShardCommitmentHash,
    /// Sector index
    pub sector_index: SectorIndex,
    /// Getter for pieces of archival history
    pub piece_getter: &'a PG,
    /// Farmer protocol info
    pub farmer_protocol_info: FarmerProtocolInfo,
    /// Erasure coding instance
    pub erasure_coding: &'a ErasureCoding,
    /// How many pieces should sector contain
    pub pieces_in_sector: u16,
    /// Where plotted sector should be written, vector must either be empty (in which case it'll be
    /// resized to correct size automatically) or correctly sized from the beginning
    pub sector_output: &'a mut Vec<u8>,
    /// Semaphore for part of the plotting when farmer downloads new sector, allows to limit memory
    /// usage of the plotting process, permit will be held until the end of the plotting process
    pub downloading_semaphore: Option<Arc<Semaphore>>,
    /// Semaphore for part of the plotting when farmer encodes downloaded sector, should typically
    /// allow one permit at a time for efficient CPU utilization
    pub encoding_semaphore: Option<&'a Semaphore>,
    /// Proof of space table generators
    pub records_encoder: &'a mut RE,
    /// Whether encoding should be aborted early
    pub abort_early: &'a AtomicBool,
}

/// Plot a single sector.
///
/// This is a convenient wrapper around [`download_sector`] and [`encode_sector`] functions.
///
/// NOTE: Even though this function is async, it has blocking code inside and must be running in a
/// separate thread in order to prevent blocking an executor.
pub async fn plot_sector<RE, PG>(
    options: PlotSectorOptions<'_, RE, PG>,
) -> Result<PlottedSector, PlottingError>
where
    RE: RecordsEncoder,
    PG: PieceGetter + Send + Sync,
{
    let PlotSectorOptions {
        public_key_hash,
        shard_commitments_root,
        sector_index,
        piece_getter,
        farmer_protocol_info,
        erasure_coding,
        pieces_in_sector,
        sector_output,
        downloading_semaphore,
        encoding_semaphore,
        records_encoder,
        abort_early,
    } = options;

    let _downloading_permit = match downloading_semaphore {
        Some(downloading_semaphore) => Some(downloading_semaphore.acquire_arc().await),
        None => None,
    };

    let download_sector_fut = download_sector(DownloadSectorOptions {
        public_key_hash,
        shard_commitments_root,
        sector_index,
        piece_getter,
        farmer_protocol_info,
        erasure_coding,
        pieces_in_sector,
    });

    let _encoding_permit = match encoding_semaphore {
        Some(encoding_semaphore) => Some(encoding_semaphore.acquire().await),
        None => None,
    };

    let encoded_sector = encode_sector(
        download_sector_fut.await?,
        EncodeSectorOptions::<RE> {
            sector_index,
            records_encoder,
            abort_early,
        },
    )?;

    if abort_early.load(Ordering::Acquire) {
        return Err(PlottingError::AbortEarly);
    }

    write_sector(&encoded_sector, sector_output)?;

    Ok(encoded_sector.plotted_sector)
}

/// Opaque sector downloading result and ready for writing
#[derive(Debug)]
pub struct DownloadedSector {
    sector_id: SectorId,
    piece_indices: Vec<PieceIndex>,
    raw_sector: RawSector,
    history_size: HistorySize,
}

/// Options for sector downloading
#[derive(Debug)]
pub struct DownloadSectorOptions<'a, PG> {
    /// Public key corresponding to sector
    pub public_key_hash: &'a Blake3Hash,
    /// Root of the Merkle Tree of shard commitments
    pub shard_commitments_root: &'a ShardCommitmentHash,
    /// Sector index
    pub sector_index: SectorIndex,
    /// Getter for pieces of archival history
    pub piece_getter: &'a PG,
    /// Farmer protocol info
    pub farmer_protocol_info: FarmerProtocolInfo,
    /// Erasure coding instance
    pub erasure_coding: &'a ErasureCoding,
    /// How many pieces should sector contain
    pub pieces_in_sector: u16,
}

/// Download sector for plotting.
///
/// This will identify necessary pieces and download them using provided piece getter, after which
/// they can be encoded using [`encode_sector`] and written to the plot.
pub async fn download_sector<PG>(
    options: DownloadSectorOptions<'_, PG>,
) -> Result<DownloadedSector, PlottingError>
where
    PG: PieceGetter + Send + Sync,
{
    let DownloadSectorOptions {
        public_key_hash,
        shard_commitments_root,
        sector_index,
        piece_getter,
        farmer_protocol_info,
        erasure_coding,
        pieces_in_sector,
    } = options;

    let sector_id = SectorId::new(
        public_key_hash,
        shard_commitments_root,
        sector_index,
        farmer_protocol_info.history_size,
    );

    let piece_indices = (PieceOffset::ZERO..)
        .take(pieces_in_sector.into())
        .map(|piece_offset| {
            sector_id.derive_piece_index(
                piece_offset,
                farmer_protocol_info.history_size,
                farmer_protocol_info.max_pieces_in_sector,
                farmer_protocol_info.recent_segments,
                farmer_protocol_info.recent_history_fraction,
            )
        })
        .collect::<Vec<_>>();

    let raw_sector = {
        let mut raw_sector = RawSector::new(pieces_in_sector);
        let mut pieces_to_download =
            HashMap::<PieceIndex, Vec<_>>::with_capacity(usize::from(pieces_in_sector));
        for (piece_index, (record, metadata)) in piece_indices
            .iter()
            .copied()
            .zip(raw_sector.records.iter_mut().zip(&mut raw_sector.metadata))
        {
            pieces_to_download
                .entry(piece_index)
                .or_default()
                .push((record, metadata));
        }
        // This map will be mutated, removing piece indices we have already processed
        let pieces_to_download = AsyncMutex::new(pieces_to_download);

        (|| async {
            let mut pieces_to_download = pieces_to_download.lock().await;

            if let Err(error) =
                download_sector_internal(&mut pieces_to_download, piece_getter, erasure_coding)
                    .await
            {
                warn!(
                    %sector_index,
                    %error,
                    %pieces_in_sector,
                    remaining_pieces = %pieces_to_download.len(),
                    "Sector downloading attempt failed, will retry later"
                );

                return Err(error);
            }

            debug!(%sector_index, "Sector downloaded successfully");

            Ok(())
        })
        .retry(
            ExponentialBuilder::default()
                .with_min_delay(Duration::from_secs(15))
                .with_max_delay(Duration::from_mins(10))
                // Try until we get a valid piece
                .without_max_times(),
        )
        .await?;

        raw_sector
    };

    Ok(DownloadedSector {
        sector_id,
        piece_indices,
        raw_sector,
        history_size: farmer_protocol_info.history_size,
    })
}

/// Records encoder for plotting purposes
pub trait RecordsEncoder {
    /// Encode provided sector records
    fn encode_records(
        &mut self,
        sector_id: &SectorId,
        records: &mut [Record],
        abort_early: &AtomicBool,
    ) -> anyhow::Result<SectorContentsMap>;
}

/// CPU implementation of [`RecordsEncoder`]
#[derive(Debug)]
pub struct CpuRecordsEncoder<'a, PosTable>
where
    PosTable: Table,
{
    table_generators: &'a [PosTable::Generator],
    erasure_coding: &'a ErasureCoding,
    global_mutex: &'a AsyncMutex<()>,
}

impl<PosTable> RecordsEncoder for CpuRecordsEncoder<'_, PosTable>
where
    PosTable: Table,
{
    fn encode_records(
        &mut self,
        sector_id: &SectorId,
        records: &mut [Record],
        abort_early: &AtomicBool,
    ) -> anyhow::Result<SectorContentsMap> {
        if self.table_generators.is_empty() {
            return Err(anyhow::anyhow!("No table generators"));
        }

        let pieces_in_sector = records
            .len()
            .try_into()
            .map_err(|error| anyhow::anyhow!("Failed to convert pieces in sector: {error}"))?;
        let mut sector_contents_map = SectorContentsMap::new(pieces_in_sector);

        {
            let global_mutex = self.global_mutex;
            let erasure_coding = self.erasure_coding;

            let iter = Mutex::new(
                (PieceOffset::ZERO..)
                    .zip(records.iter_mut())
                    .zip(sector_contents_map.iter_record_chunks_used_mut()),
            );

            rayon::scope(|scope| {
                for table_generator in self.table_generators {
                    scope.spawn(|_scope| {
                        loop {
                            // Take mutex briefly to make sure encoding is allowed right now
                            global_mutex.lock_blocking();

                            // This instead of `while` above because otherwise mutex will be held
                            // for the duration of the loop and will limit concurrency to 1 record
                            let Some(((piece_offset, record), record_chunks_used)) =
                                iter.lock().next()
                            else {
                                return;
                            };
                            let pos_seed = sector_id.derive_evaluation_seed(piece_offset);

                            record_encoding::<PosTable>(
                                &pos_seed,
                                record,
                                record_chunks_used,
                                table_generator,
                                erasure_coding,
                            );

                            if abort_early.load(Ordering::Relaxed) {
                                return;
                            }
                        }
                    });
                }
            });
        }

        Ok(sector_contents_map)
    }
}

impl<'a, PosTable> CpuRecordsEncoder<'a, PosTable>
where
    PosTable: Table,
{
    /// Create a new instance
    pub fn new(
        table_generators: &'a [PosTable::Generator],
        erasure_coding: &'a ErasureCoding,
        global_mutex: &'a AsyncMutex<()>,
    ) -> Self {
        Self {
            table_generators,
            erasure_coding,
            global_mutex,
        }
    }
}

/// Options for encoding a sector.
///
/// Sector output and sector metadata output should be either empty (in which case they'll be
/// resized to correct size automatically) or correctly sized from the beginning or else error will
/// be returned.
#[derive(Debug)]
pub struct EncodeSectorOptions<'a, RE>
where
    RE: RecordsEncoder,
{
    /// Sector index
    pub sector_index: SectorIndex,
    /// Records encoding instance
    pub records_encoder: &'a mut RE,
    /// Whether encoding should be aborted early
    pub abort_early: &'a AtomicBool,
}

/// Mostly opaque sector encoding result ready for writing
#[derive(Debug)]
#[expect(
    clippy::partial_pub_fields,
    reason = "Intentionally hiding some fields such that they can only be used by internal APIs"
)]
pub struct EncodedSector {
    /// Information about sector that was plotted
    pub plotted_sector: PlottedSector,
    raw_sector: RawSector,
    sector_contents_map: SectorContentsMap,
}

/// Encode downloaded sector.
///
/// This function encodes downloaded sector records and returns sector encoding result that can be
/// written using [`write_sector`].
pub fn encode_sector<RE>(
    downloaded_sector: DownloadedSector,
    encoding_options: EncodeSectorOptions<'_, RE>,
) -> Result<EncodedSector, PlottingError>
where
    RE: RecordsEncoder,
{
    let DownloadedSector {
        sector_id,
        piece_indices,
        mut raw_sector,
        history_size,
    } = downloaded_sector;
    let EncodeSectorOptions {
        sector_index,
        records_encoder,
        abort_early,
    } = encoding_options;

    let pieces_in_sector = raw_sector.records.len().try_into().expect(
        "Raw sector can only be created in this crate and it is always done correctly; qed",
    );

    let sector_contents_map = records_encoder
        .encode_records(&sector_id, &mut raw_sector.records, abort_early)
        .map_err(|error| PlottingError::RecordsEncoderError { error })?;

    let sector_metadata = SectorMetadataChecksummed::from(SectorMetadata {
        sector_index,
        pieces_in_sector,
        s_bucket_sizes: sector_contents_map.s_bucket_sizes(),
        history_size,
    });

    Ok(EncodedSector {
        plotted_sector: PlottedSector {
            sector_id,
            sector_index,
            sector_metadata,
            piece_indexes: piece_indices,
        },
        raw_sector,
        sector_contents_map,
    })
}

/// Write encoded sector into sector output
pub fn write_sector(
    encoded_sector: &EncodedSector,
    sector_output: &mut Vec<u8>,
) -> Result<(), PlottingError> {
    let EncodedSector {
        plotted_sector: _,
        raw_sector,
        sector_contents_map,
    } = encoded_sector;

    let pieces_in_sector = raw_sector.records.len().try_into().expect(
        "Raw sector can only be created in this crate and it is always done correctly; qed",
    );

    let sector_size = sector_size(pieces_in_sector);

    if !sector_output.is_empty() && sector_output.len() != sector_size {
        return Err(PlottingError::BadSectorOutputSize {
            provided: sector_output.len(),
            expected: sector_size,
        });
    }

    sector_output.resize(sector_size, 0);

    // Write sector to disk in as the following regions:
    // * sector contents map
    // * record chunks as s-buckets
    // * record metadata
    // * checksum
    {
        let (sector_contents_map_region, remaining_bytes) =
            sector_output.split_at_mut(SectorContentsMap::encoded_size(pieces_in_sector));
        // Slice remaining memory into belonging to s-buckets and metadata
        let (s_buckets_region, metadata_region) =
            remaining_bytes.split_at_mut(sector_record_chunks_size(pieces_in_sector));

        // Write sector contents map so we can decode it later
        sector_contents_map
            .encode_into(sector_contents_map_region)
            .expect("Chunked into correct size above; qed");

        let mut next_record_chunks_offset = vec![0usize; pieces_in_sector.into()];
        // Write record chunks, one s-bucket at a time
        for (piece_offset, output) in (SBucket::ZERO..=SBucket::MAX)
            .flat_map(|s_bucket| {
                sector_contents_map
                    .iter_s_bucket_piece_offsets(s_bucket)
                    .expect("S-bucket guaranteed to be in range; qed")
            })
            .zip(s_buckets_region.as_chunks_mut::<{ RecordChunk::SIZE }>().0)
        {
            let next_record_chunks_offset =
                &mut next_record_chunks_offset[usize::from(piece_offset)];

            let chunk_position = *next_record_chunks_offset;
            *next_record_chunks_offset += 1;
            output.copy_from_slice(&raw_sector.records[usize::from(piece_offset)][chunk_position]);
        }

        let metadata_chunks = metadata_region
            .as_chunks_mut::<{ RecordMetadata::encoded_size() }>()
            .0;
        for (record_metadata, output) in raw_sector.metadata.iter().zip(metadata_chunks) {
            record_metadata.encode_to(&mut output.as_mut_slice());
        }

        // It would be more efficient to not re-read the whole sector again, but it makes above code
        // significantly more convoluted and most likely not worth it
        let (sector_contents, sector_checksum) =
            sector_output.split_at_mut(sector_size - Blake3Hash::SIZE);
        sector_checksum.copy_from_slice(
            {
                let mut hasher = blake3::Hasher::new();
                hasher.update_rayon(sector_contents);
                hasher.finalize()
            }
            .as_bytes(),
        );
    }

    Ok(())
}

fn record_encoding<PosTable>(
    pos_seed: &PosSeed,
    record: &mut Record,
    record_chunks_used: &mut FoundProofs,
    table_generator: &PosTable::Generator,
    erasure_coding: &ErasureCoding,
) where
    PosTable: Table,
{
    let pos_proofs = table_generator.create_proofs_parallel(pos_seed);

    let mut parity_record_chunks = Record::new_boxed();

    // Erasure code source record chunks
    erasure_coding
        .extend(record.iter(), parity_record_chunks.iter_mut())
        .expect("Statically guaranteed valid inputs; qed");

    *record_chunks_used = pos_proofs.found_proofs;

    // TODO: This can probably be optimized by using SIMD
    let mut num_found_proofs = 0_usize;
    for (s_buckets, found_proofs) in (0..Record::NUM_S_BUCKETS)
        .array_chunks::<{ u8::BITS as usize }>()
        .zip(pos_proofs.found_proofs)
    {
        for (proof_offset, s_bucket) in s_buckets.into_iter().enumerate() {
            if (found_proofs & (1 << proof_offset)) != 0 {
                let record_chunk = if s_bucket < Record::NUM_CHUNKS {
                    record[s_bucket]
                } else {
                    parity_record_chunks[s_bucket - Record::NUM_CHUNKS]
                };
                // TODO: SIMD hashing
                record[num_found_proofs] = (Simd::from(record_chunk)
                    ^ Simd::from(*pos_proofs.proofs[num_found_proofs].hash()))
                .to_array();
                num_found_proofs += 1;
            }
        }
    }
}

async fn download_sector_internal<PG>(
    pieces_to_download: &mut HashMap<PieceIndex, Vec<(&mut Record, &mut RecordMetadata)>>,
    piece_getter: &PG,
    erasure_coding: &ErasureCoding,
) -> Result<(), PlottingError>
where
    PG: PieceGetter + Send + Sync,
{
    // TODO: Make configurable, likely allowing user to specify RAM usage expectations and inferring
    //  concurrency from there
    let recovery_semaphore = &Semaphore::new(RECONSTRUCTION_CONCURRENCY_LIMIT);

    // Allocate to decouple lifetime from `pieces_to_download` that will be modified below
    let piece_indices = pieces_to_download.keys().copied().collect::<Vec<_>>();
    let mut downloaded_pieces = piece_getter
        .get_pieces(piece_indices)
        .await
        .map_err(|error| PlottingError::FailedToRetrievePieces { error })?
        .fuse();
    let mut reconstructed_pieces = FuturesUnordered::new();

    let mut final_result = Ok(());

    loop {
        let (piece_index, result) = select! {
            (piece_index, result) = downloaded_pieces.select_next_some() => {
                match result {
                    Ok(Some(piece)) => (piece_index, Ok(piece)),
                    Ok(None) => {
                        trace!(%piece_index, "Piece was not found, trying reconstruction");

                        reconstructed_pieces.push(reconstruct_piece(
                            piece_index,
                            recovery_semaphore,
                            piece_getter,
                            erasure_coding,
                        ));
                        continue;
                    }
                    Err(error) => {
                        trace!(
                            %error,
                            %piece_index,
                            "Failed to download piece, trying reconstruction"
                        );

                        reconstructed_pieces.push(reconstruct_piece(
                            piece_index,
                            recovery_semaphore,
                            piece_getter,
                            erasure_coding,
                        ));
                        continue;
                    }
                }
            },
            (piece_index, result) = reconstructed_pieces.select_next_some() => {
                (piece_index, result)
            },
            complete => {
                break;
            }
        };

        match result {
            Ok(piece) => {
                process_piece(piece_index, piece, pieces_to_download);
            }
            Err(error) => {
                trace!(%error, %piece_index, "Failed to download piece");

                if final_result.is_ok() {
                    final_result = Err(error);
                }
            }
        }
    }

    if final_result.is_ok() && !pieces_to_download.is_empty() {
        return Err(PlottingError::FailedToRetrievePieces {
            error: anyhow::anyhow!(
                "Successful result, but not all pieces were downloaded, this is likely a piece \
                getter implementation bug"
            ),
        });
    }

    final_result
}

async fn reconstruct_piece<PG>(
    piece_index: PieceIndex,
    recovery_semaphore: &Semaphore,
    piece_getter: &PG,
    erasure_coding: &ErasureCoding,
) -> (PieceIndex, Result<Piece, PlottingError>)
where
    PG: PieceGetter + Send + Sync,
{
    let _permit = recovery_semaphore.acquire().await;
    let recovered_piece_fut =
        recover_missing_piece(piece_getter, erasure_coding.clone(), piece_index);

    (
        piece_index,
        recovered_piece_fut
            .await
            .map_err(|error| PlottingError::PieceRecoveryFailed {
                piece_index,
                error: error.into(),
            }),
    )
}

fn process_piece(
    piece_index: PieceIndex,
    piece: Piece,
    pieces_to_download: &mut HashMap<PieceIndex, Vec<(&mut Record, &mut RecordMetadata)>>,
) {
    for (record, metadata) in pieces_to_download.remove(&piece_index).unwrap_or_default() {
        *metadata = RecordMetadata {
            piece_header: piece.header,
            piece_checksum: blake3::hash(piece.as_ref()).into(),
        };
        // Fancy way to insert value to avoid going through stack (if naive dereferencing is used)
        // and potentially causing stack overflow as the result
        record.copy_from_slice(&*piece.record);
    }
}