dbsp 0.317.0

use std::{
    borrow::Cow,
    panic::Location,
    sync::{
        Arc,
        atomic::{AtomicUsize, Ordering},
    },
};

use size_of::SizeOf;
use typedmap::TypedMapKey;

use crate::{
    Circuit, Error, NumEntries, Runtime, Scope, Stream,
    circuit::{
        GlobalNodeId, LocalStoreMarker,
        circuit_builder::StreamId,
        metadata::{
            ALLOCATED_MEMORY_BYTES, BatchSizeStats, INPUT_BATCHES_STATS, MEMORY_ALLOCATIONS_COUNT,
            MetaItem, OUTPUT_BATCHES_STATS, OperatorLocation, OperatorMeta, SHARED_MEMORY_BYTES,
            STATE_RECORDS_COUNT, USED_MEMORY_BYTES,
        },
        operator_traits::{Operator, OperatorName, UnaryOperator},
    },
    circuit_cache_key,
    trace::{Batch, BatchReader, Spine, Trace},
};

circuit_cache_key!(AccumulatorId<C, B: Batch>(StreamId => Accumulation<Stream<C, Option<Spine<B>>>>));
circuit_cache_key!(ShardedAccumulatorId<C, B: Batch>(StreamId => Stream<C, Option<Spine<B>>>));

/// A stream produced by accumulating batches into a spine.
///
/// This is the type returned by [Stream::dyn_accumulate] and similar methods.
#[derive(Clone, Debug)]
pub struct Accumulation<T> {
    /// The result stream.
    pub stream: T,

    /// An [EnableCount].
    ///
    /// During a transaction when the enable count is disabled, the accumulator
    /// will simply discard the input batches instead of accumulating them.
    /// This is useful to increase performance when an output stream currently
    /// has nothing listening to it.
    pub enable_count: EnableCount,
}

impl<T> Accumulation<T> {
    /// Enables the [EnableCount] and returns the stream.
    ///
    /// Because this discards the `EnableCount`, this is intended for situations
    /// where the accumulator should be enabled permanently.
    pub fn into_enabled_stream(self) -> T {
        self.enable_count.enable();
        self.stream
    }

    /// Returns the stream and the enable count inside this `Accumulation`.
    pub fn into_parts(self) -> (T, EnableCount) {
        (self.stream, self.enable_count)
    }

    /// Returns an equivalent `Accumulation` with `f` applied to the inner
    /// stream.
    pub fn map<F, R>(self, f: F) -> Accumulation<R>
    where
        F: FnOnce(T) -> R,
    {
        Accumulation {
            stream: f(self.stream),
            enable_count: self.enable_count,
        }
    }
}

/// `TypedMapKey` entry used to share `enable_count` across instances of the same accumulator in multiple workers.
#[derive(Hash, PartialEq, Eq)]
struct EnableCountId {
    id: usize,
}

impl EnableCountId {
    fn new(id: usize) -> Self {
        Self { id }
    }
}

/// Used to enable/disable an accumulator during a transaction.
///
/// Most accumulators (created with dyn_accumulate()) are always enabled.
/// One special case is when the accumulator is used as part of an output handle
/// to collect updates to the output stream within a transaction. In this case,
/// if there is no output connector attached to the stream, there is no need to
/// store the updates (which during backfill can amount to storing a complete copy
/// of the table or view).
///
/// This flag enables this optimization by keeping track of the number of consumers
/// of the accumulator's output. It is equal to the number of attached output connectors
/// plus the number of times the same accumulator was instantiated as part of a regular
/// (non-output) operator with dyn_accumulate().
#[derive(Clone, Debug, Default)]
pub struct EnableCount(Arc<AtomicUsize>);

impl EnableCount {
    /// Creates a new `EnableCount` that is initially disabled.
    pub fn new() -> Self {
        Self::default()
    }

    /// Returns true if this `EnableCount` is enabled.
    pub fn is_enabled(&self) -> bool {
        self.0.load(Ordering::Acquire) > 0
    }

    /// Enable the accumulator for this output stream.
    ///
    /// This may be paired with a later call to [EnableCount::disable], if the
    /// stream should eventually be disabled.
    pub fn enable(&self) {
        self.0.fetch_add(1, Ordering::AcqRel);
    }

    /// Disable the accumulator for this output stream.
    ///
    /// This must be paired with a prior call to [EnableCount::enable].
    pub fn disable(&self) {
        let old = self.0.fetch_add(1, Ordering::AcqRel);
        assert!(old > 0);
    }
}

impl TypedMapKey<LocalStoreMarker> for EnableCountId {
    type Value = EnableCount;
}

impl<C, B> Stream<C, B>
where
    C: Circuit,
    B: Batch,
{
    /// See [`Stream::accumulate`].
    pub fn dyn_accumulate(
        &self,
        factories: &B::Factories,
    ) -> Accumulation<Stream<C, Option<Spine<B>>>> {
        self.circuit()
            .cache_get_or_insert_with(AccumulatorId::new(self.stream_id()), || {
                let accumulator = Accumulator::<B>::new(factories, Location::caller());
                let enable_count = accumulator.enable_count.clone();

                let stream = self
                    .circuit()
                    .add_unary_operator(accumulator, &self.try_sharded_version());
                stream.mark_sharded_if(self);
                Accumulation {
                    stream,
                    enable_count,
                }
            })
            .clone()
    }
}

pub struct Accumulator<B>
where
    B: Batch,
{
    factories: B::Factories,
    name: OperatorName,
    state: Spine<B>,
    flush: bool,
    location: &'static Location<'static>,

    // Input batch sizes.
    input_batch_stats: BatchSizeStats,

    // Output batch sizes.
    output_batch_stats: BatchSizeStats,

    /// Used to enable/disable an accumulator during a transaction.
    enable_count: EnableCount,

    /// Whether the accumulator is enabled during the current transaction.
    ///
    /// An output connector can be attached in the middle of a transaction; however if the
    /// accumulator was disabled at the start of the transaction, it shouldn't produce
    /// partial outputs. This flag remembers the status of the accumulator at the start of the
    /// transaction.
    enabled_during_current_transaction: Option<bool>,
}

impl<B> Accumulator<B>
where
    B: Batch,
{
    pub fn new(factories: &B::Factories, location: &'static Location<'static>) -> Self {
        let enable_count = match Runtime::runtime() {
            None => EnableCount::default(),
            Some(runtime) => {
                let accumulator_id = runtime.sequence_next();
                runtime
                    .local_store()
                    .entry(EnableCountId::new(accumulator_id))
                    .or_insert_with(EnableCount::default)
                    .value()
                    .clone()
            }
        };

        let name = OperatorName::new("Accumulator");
        Self {
            factories: factories.clone(),
            state: Spine::new(factories, name.get()),
            name,
            flush: false,
            location,
            input_batch_stats: BatchSizeStats::new(),
            output_batch_stats: BatchSizeStats::new(),
            enable_count,
            enabled_during_current_transaction: None,
        }
    }
}

impl<B> Operator for Accumulator<B>
where
    B: Batch,
{
    fn name(&self) -> std::borrow::Cow<'static, str> {
        Cow::Borrowed("Accumulator")
    }

    fn init(&mut self, global_id: &GlobalNodeId) {
        self.name.init(global_id);
        self.state.set_name(self.name.get());
    }

    fn location(&self) -> OperatorLocation {
        Some(self.location)
    }

    fn metadata(&self, meta: &mut OperatorMeta) {
        let total_size = self.state.num_entries_deep();

        let bytes = self.state.size_of();

        meta.extend(metadata! {
            STATE_RECORDS_COUNT => MetaItem::Count(total_size),
            ALLOCATED_MEMORY_BYTES => MetaItem::bytes(bytes.total_bytes()),
            USED_MEMORY_BYTES => MetaItem::bytes(bytes.used_bytes()),
            MEMORY_ALLOCATIONS_COUNT => MetaItem::Count(bytes.distinct_allocations()),
            SHARED_MEMORY_BYTES => MetaItem::bytes(bytes.shared_bytes()),
            INPUT_BATCHES_STATS => self.input_batch_stats.metadata(),
            OUTPUT_BATCHES_STATS => self.output_batch_stats.metadata(),

        });

        self.state.metadata(meta);
    }

    fn clock_start(&mut self, _scope: Scope) {
        debug_assert!(self.state.is_empty());
    }

    fn clock_end(&mut self, _scope: Scope) {
        debug_assert!(self.state.is_empty());
    }

    fn fixedpoint(&self, _scope: Scope) -> bool {
        self.state.is_empty()
    }

    /// Clear the operator's state.
    fn clear_state(&mut self) -> Result<(), Error> {
        self.state = Spine::new(&self.factories, self.name.get());
        self.flush = false;
        Ok(())
    }

    fn flush(&mut self) {
        self.flush = true;
    }

    fn is_flush_complete(&self) -> bool {
        !self.flush
    }
}

impl<B> UnaryOperator<B, Option<Spine<B>>> for Accumulator<B>
where
    B: Batch,
{
    async fn eval(&mut self, batch: &B) -> Option<Spine<B>> {
        self.eval_owned(batch.clone()).await
    }

    async fn eval_owned(&mut self, batch: B) -> Option<Spine<B>> {
        // We don't have a start-of-transaction signal, so we sample enable_count when
        // we get the first non-empty batch.  This batch should belong to the next transaction
        // after the last one that was flushed, since the accumulator should not receive any
        // non-empty batches from the previous transaction at that point (in the top-level circuit).
        // This may not be the first batch in the transaction, but it's ok to admit some empty batches.
        let len = batch.len();

        if len > 0 {
            if self.enabled_during_current_transaction.is_none() {
                self.enabled_during_current_transaction = Some(self.enable_count.is_enabled());
            }

            if self.enabled_during_current_transaction == Some(true) {
                self.input_batch_stats.add_batch(len);
                self.state.insert(batch).await;
            }
        }

        if self.flush {
            self.flush = false;
            self.enabled_during_current_transaction = None;

            let mut spine = Spine::<B>::new(&self.factories, self.name.get());
            std::mem::swap(&mut self.state, &mut spine);

            self.output_batch_stats.add_batch(spine.len());
            Some(spine)
        } else {
            None
        }
    }
}