gled-core 2.31.1

//! Copy a output buffer to the cpu and return.

use gled_gpu::{
    constants::{OUTPUT_BUFFER_SIZE, UNIVERSE_BUFFER_SIZE, UNIVERSE_SLOTS},
    gpu::OutputContext,
};
use kanal::{Receiver, Sender, bounded};
use std::sync::Arc;
use tracing::warn;
use wgpu::{
    Buffer, BufferDescriptor, BufferUsages, CommandEncoder, Device, MapMode, PollType,
    SubmissionIndex,
};

/// Number of reusable CPU-readback staging buffers kept around. A frame
/// finding none free is not read back and never reaches the output. A buffer
/// stays taken from its frame's submission until its readback is delivered,
/// which while the GPU is the limit spans several frames: with three, a
/// laptop skipped 13 % of the frames of 4096 universes (`tests/rig_bench`),
/// with six none. Each costs `OUTPUT_BUFFER_SIZE`, 4 MB.
const STAGING_POOL_SIZE: usize = 6;
/// Read-back frames queued for the UI preview. The UI drains this every frame
/// and only uses the newest; frames are not copied while it is full.
const PREVIEW_CHANNEL_CAPACITY: usize = 2;

/// Raw output bytes of one rendered frame, before color correction.
pub(crate) type PreviewOutput = Arc<Vec<u8>>;

pub(crate) use gled_protocol::state::OutputConfig;

/// One rendered frame on its way to the output devices, with the config it
/// was rendered for. Cloning the config only clones its shared parts.
#[derive(Debug)]
pub(crate) struct OutputFrame {
    pub(crate) data: Vec<u8>,
    pub(crate) config: OutputConfig,
}

/// The frames on their way to the output thread, and the bell that wakes
/// it. Whatever hands the output thread work queues it first and rings the
/// bell after, so the thread can block on the bell alone (kanal's timed
/// receive spins instead of sleeping).
#[derive(Debug)]
pub(crate) struct OutputFrames {
    pub(crate) frames: Receiver<OutputFrame>,
    pub(crate) wake: Receiver<()>,
    /// For the others that hand the output thread work.
    pub(crate) bell: Sender<()>,
}

pub(crate) struct ExtractOutput {
    output_sender: Sender<OutputFrame>,
    bell: Sender<()>,
    preview_sender: Sender<PreviewOutput>,
    /// Free-list of reusable map-read staging buffers. Buffers are lazily
    /// created on the first `run` (the wgpu device does not exist yet when
    /// `ExtractOutput::new` is called in `main`) and recycled afterwards.
    free_buffers_sender: Sender<Buffer>,
    free_buffers_receiver: Receiver<Buffer>,
    /// Submission indices of frames whose readback should be polled. The
    /// dedicated poll thread waits for each one individually so every frame's
    /// `map_buffer_on_submit` callback (the readback + network send) fires as
    /// soon as *that* frame's GPU work finishes, independently of any later
    /// frame the render thread submits in the meantime.
    submit_sender: Sender<SubmissionIndex>,
    /// Held until the poll thread is started on the first `run`, then taken.
    submit_receiver: Option<Receiver<SubmissionIndex>>,
    initialized: bool,
    config: OutputConfig,
}

impl std::fmt::Debug for ExtractOutput {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("ExtractOutput").finish_non_exhaustive()
    }
}

impl ExtractOutput {
    /// Returns the extractor, the frames for the Art-Net/DMX output and the
    /// receiver for the UI preview.
    pub(crate) fn new() -> (Self, OutputFrames, Receiver<PreviewOutput>) {
        // The output channel is sized to the staging pool: at most
        // `STAGING_POOL_SIZE` readbacks can be in flight at once (each holds a
        // staging buffer until its map callback returns it), and when several
        // complete in the same device poll every one of them must be delivered
        // instead of being collapsed into a single frame. The render thread
        // still never blocks: it `try_send`s and drops on backpressure, so a
        // genuinely overloaded output consumer simply loses the oldest surplus.
        let (output_sender, output_receiver) = bounded(STAGING_POOL_SIZE);
        // One ring is enough: the output thread takes every queued frame
        // when it wakes.
        let (bell, wake) = bounded(1);
        let (preview_sender, preview_receiver) = bounded(PREVIEW_CHANNEL_CAPACITY);
        let (free_buffers_sender, free_buffers_receiver) = bounded(STAGING_POOL_SIZE);
        // One slot per in-flight readback is enough: a submission is only worth
        // polling while its staging buffer is still mapped, and there are at
        // most `STAGING_POOL_SIZE` of those. Sized a little larger so that short
        // bursts of fast render frames are never dropped here.
        let (submit_sender, submit_receiver) = bounded(STAGING_POOL_SIZE * 2);

        (
            Self {
                output_sender,
                bell: bell.clone(),
                preview_sender,
                free_buffers_sender,
                free_buffers_receiver,
                submit_sender,
                submit_receiver: Some(submit_receiver),
                initialized: false,
                config: OutputConfig::default(),
            },
            OutputFrames {
                frames: output_receiver,
                wake,
                bell,
            },
            preview_receiver,
        )
    }

    /// Lazily initialise GPU-side resources on the output device: the reusable
    /// staging buffers and the dedicated poll thread that delivers each
    /// finished frame's readback independently.
    fn ensure_initialized(&mut self, device: &Device) {
        if self.initialized {
            return;
        }
        for index in 0..STAGING_POOL_SIZE {
            let buffer = device.create_buffer(&BufferDescriptor {
                size: OUTPUT_BUFFER_SIZE,
                usage: BufferUsages::MAP_READ | BufferUsages::COPY_DST,
                label: Some(&format!("ExtractOutput staging buffer {index}")),
                mapped_at_creation: false,
            });
            self.free_buffers_sender
                .try_send(buffer)
                .expect("staging pool channel has capacity for the pool");
        }

        // Spawn a thread that polls the device once per submitted frame. By
        // waiting for each submission *index* in turn it lets every frame's
        // `map_buffer_on_submit` callback run the instant that frame is done on
        // the GPU – so frame A is read back and sent before frame B finishes,
        // instead of both being delivered together at the next present poll.
        if let Some(submit_receiver) = self.submit_receiver.take() {
            let poll_device = device.clone();
            std::thread::Builder::new()
                .name("gled:output:poll".to_owned())
                .spawn(move || {
                    #[cfg(feature = "profiling")]
                    profiling::register_thread!("output:poll");

                    while let Ok(index) = submit_receiver.recv() {
                        // Blocks only until this specific submission completes;
                        // the readback callback (and network send) runs here.
                        let _ = poll_device.poll(PollType::Wait {
                            submission_index: Some(index),
                            timeout: None,
                        });
                    }
                })
                .expect("Could not spawn output poll thread");
        }

        self.initialized = true;
    }

    /// Notify the poll thread that a frame has been submitted so its readback is
    /// delivered as soon as the GPU finishes it. Called by `Project::render`
    /// right after `queue.submit`. Non-blocking: if the poll thread is briefly
    /// behind, the frame is still delivered by the next poll.
    pub(crate) fn notify_submitted(&self, index: SubmissionIndex) {
        let _ = self.submit_sender.try_send(index);
    }

    /// Takes the config the next frames are rendered for. A patch with more
    /// universes than the output buffer has slots is reported once, when its
    /// universes arrive, not with every frame.
    pub(crate) fn set_config(&mut self, config: OutputConfig) {
        let universes = &config.universes;
        if !Arc::ptr_eq(universes, &self.config.universes)
            && universes != &self.config.universes
            && let Some(first_dark) = universes.iter().nth(UNIVERSE_SLOTS as usize)
        {
            warn!(
                "The patch uses {} universes, gled renders {UNIVERSE_SLOTS} at once: \
                 universe {first_dark} and the ones above it get no output",
                universes.len()
            );
        }
        self.config = config;
    }

    /// The output buffer slots the patch's universes fill.
    pub(crate) fn slots(&self) -> usize {
        self.config.universes.len().min(UNIVERSE_SLOTS as usize)
    }

    #[cfg_attr(feature = "profiling", profiling::function)]
    pub(crate) fn run(&mut self, context: &OutputContext, encoder: &mut CommandEncoder) {
        let active_len = self.slots() as u64 * UNIVERSE_BUFFER_SIZE;
        if active_len == 0 {
            return;
        }

        self.ensure_initialized(&context.gpu.device);

        // Grab a free staging buffer. If every buffer is still in flight the
        // CPU side has not kept up, so we simply skip this frame's readback –
        // the next rendered frame carries fresher data anyway.
        let Ok(Some(output_cpu)) = self.free_buffers_receiver.try_recv() else {
            return;
        };

        encoder.copy_buffer_to_buffer(&context.output_buffer, 0, &output_cpu, 0, active_len);

        let output_sender = self.output_sender.clone();
        let bell = self.bell.clone();
        let preview_sender = self.preview_sender.clone();
        let free_buffers_sender = self.free_buffers_sender.clone();
        let config = self.config.clone();
        let capturable = output_cpu.clone();
        encoder.map_buffer_on_submit(&output_cpu, MapMode::Read, ..active_len, move |_v| {
            let output_data = capturable
                .get_mapped_range(..active_len)
                .expect("Could not get mapped range of output buffer")
                .to_vec();
            capturable.unmap();
            if !preview_sender.is_full() {
                let _ = preview_sender.try_send(Arc::new(output_data.clone()));
            }
            // Non-blocking: keep only the latest frame, never stall the render
            // thread on a slow output consumer.
            let _ = output_sender.try_send(OutputFrame {
                data: output_data,
                config,
            });
            let _ = bell.try_send(());
            // Return the buffer to the pool for reuse.
            let _ = free_buffers_sender.try_send(capturable);
        });
    }
}