acquire.hpp

/*
 * This file is part of libtcspc
 * Copyright 2019-2026 Board of Regents of the University of Wisconsin System
 * SPDX-License-Identifier: MIT
 */
 
#pragma once
 
#include "arg_wrappers.hpp"
#include "bucket.hpp"
#include "context.hpp"
#include "core.hpp"
#include "errors.hpp"
#include "introspect.hpp"
#include "processor.hpp"
 
#include <chrono>
#include <condition_variable>
#include <cstddef>
#include <exception>
#include <functional>
#include <memory>
#include <mutex>
#include <optional>
#include <span>
#include <stdexcept>
#include <utility>
 
namespace tcspc {

class acquire_access {
    std::function<void()> halt_fn;
 
  public:
    template <typename Func>
    explicit acquire_access(Func halt_func) : halt_fn(halt_func) {}

    void halt() { halt_fn(); }
};
 
namespace internal {
 
// The max sleep duration when a read doesn't fill a batch is chosen to be
// short enough that (1) its effect is unnoticeable in a live display of the
// data and (2) hardware buffers are unlikely to fill up if they started out
// empty, given the buffer capacity and maximum count rates of typical devices.
constexpr auto slow_acq_sleep = std::chrono::milliseconds(10);
 
template <typename T, typename Reader, typename Downstream> class acquire {
    static_assert(processor<Downstream, bucket<T>>);
 
    Reader reader;
    std::shared_ptr<bucket_source<T>> bsource;
    std::size_t bsize;
 
    std::mutex halt_mutex;
    std::condition_variable halt_cv;
    bool halted = false;
 
    Downstream downstream;
 
    // Cold data after downstream.
    access_tracker<acquire_access> trk;
 
    void halt() {
        {
            auto const lock = std::lock_guard(halt_mutex);
            halted = true;
        }
        halt_cv.notify_one();
    }
 
  public:
    explicit acquire(Reader reader,
                     std::shared_ptr<bucket_source<T>> buffer_provider,
                     arg::batch_size<std::size_t> batch_size,
                     access_tracker<acquire_access> tracker,
                     Downstream downstream)
        : reader(std::move(reader)), bsource(std::move(buffer_provider)),
          bsize(batch_size.value), downstream(std::move(downstream)),
          trk(std::move(tracker)) {
        if (not bsource)
            throw std::invalid_argument(
                "acquire buffer_provider must not be null");
        if (bsize == 0)
            throw std::invalid_argument("acquire batch size must be positive");
 
        trk.register_access_factory([](auto &tracker) {
            auto *self = LIBTCSPC_OBJECT_FROM_TRACKER(acquire, trk, tracker);
            return acquire_access([self] { self->halt(); });
        });
    }
 
    // Custom move because we have a mutex. Move only works when not running.
    ~acquire() = default;
 
    acquire(acquire const &) = delete;
    auto operator=(acquire const &) = delete;
 
    acquire(acquire &&other) noexcept
        : reader(std::move(other.reader)), bsource(std::move(other.bsource)),
          bsize(other.bsize), halted(other.halted),
          downstream(std::move(other.downstream)), trk(std::move(other.trk)) {}
 
    auto operator=(acquire &&rhs) noexcept -> acquire & {
        reader = std::move(rhs.reader);
        bsource = std::move(rhs.bsource);
        bsize = rhs.bsize;
        halted = rhs.halted;
        downstream = std::move(rhs.downstream);
        trk = std::move(rhs.trk);
        return *this;
    }
 
    [[nodiscard]] auto introspect_node() const -> processor_info {
        return processor_info(this, "acquire");
    }
 
    [[nodiscard]] auto introspect_graph() const -> processor_graph {
        return downstream.introspect_graph().push_entry_point(this);
    }
 
    void flush() {
        bucket<T> b;
        bool reached_end = false;
        {
            auto lock = std::unique_lock(halt_mutex);
            while (not halted) {
                lock.unlock();
                auto const start_time = std::chrono::steady_clock::now();
                if (b.empty())
                    b = bsource->bucket_of_size(bsize);
                std::optional<std::size_t> const read = reader(std::span(b));
                if (not read) {
                    reached_end = true;
                    break;
                }
                if (*read > 0) {
                    b.shrink(0, *read);
                    downstream.handle(std::move(b));
                    b = {};
                }
                lock.lock();
                if (*read < bsize) { // Not enough data to fill the batch.
                    halt_cv.wait_until(lock, start_time + slow_acq_sleep,
                                       [&] { return halted; });
                }
            }
        }
        if (reached_end)
            downstream.flush();
        else
            throw acquisition_halted();
    }
};
 
template <typename T, typename Reader, typename LiveDownstream,
          typename BatchDownstream>
class acquire_full_buckets {
    static_assert(processor<LiveDownstream, bucket<T const>>);
    static_assert(processor<BatchDownstream, bucket<T>>);
 
    Reader reader;
    std::shared_ptr<bucket_source<T>> bsource;
    std::size_t bsize;
 
    std::mutex halt_mutex;
    std::condition_variable halt_cv;
    bool halted = false;
 
    LiveDownstream live_downstream;
    BatchDownstream batch_downstream;
 
    // Cold data after downstream.
    access_tracker<acquire_access> trk;
 
    void halt() {
        {
            auto const lock = std::lock_guard(halt_mutex);
            halted = true;
        }
        halt_cv.notify_one();
    }
 
    // Mutates 'b' only when throwing.
    void emit_live(bucket<T> &b, std::size_t start, std::size_t count) {
        if (count > 0) {
            try {
                auto v = bsource->shared_view_of(b);
                v.shrink(start, count);
                live_downstream.handle(std::move(v));
            } catch (end_of_processing const &) {
                b.shrink(0, start + count);
                batch_downstream.handle(std::move(b));
                batch_downstream.flush();
                throw;
            }
        }
    }
 
    void emit_batch(bucket<T> &&b) {
        try {
            batch_downstream.handle(std::move(b));
        } catch (end_of_processing const &) {
            live_downstream.flush();
            throw;
        }
    }
 
    void flush_downstreams(bucket<T> &&b, std::size_t filled) {
        std::exception_ptr end;
        try {
            live_downstream.flush();
        } catch (end_of_processing const &) {
            end = std::current_exception();
        }
        if (not b.empty() && filled > 0) {
            b.shrink(0, filled);
            batch_downstream.handle(std::move(b));
        }
        batch_downstream.flush();
        if (end)
            std::rethrow_exception(end);
    }
 
  public:
    explicit acquire_full_buckets(
        Reader reader, std::shared_ptr<bucket_source<T>> buffer_provider,
        arg::batch_size<std::size_t> batch_size,
        access_tracker<acquire_access> tracker, LiveDownstream live_downstream,
        BatchDownstream batch_downstream)
        : reader(std::move(reader)), bsource(std::move(buffer_provider)),
          bsize(batch_size.value), live_downstream(std::move(live_downstream)),
          batch_downstream(std::move(batch_downstream)),
          trk(std::move(tracker)) {
        if (not bsource)
            throw std::invalid_argument(
                "acquire_full_buckets buffer_provider must not be null");
        if constexpr (not std::is_same_v<LiveDownstream, null_sink>) {
            if (not bsource->supports_shared_views())
                throw std::invalid_argument(
                    "acquire_full_buckets buffer_provider must support shared views");
        }
        if (bsize == 0)
            throw std::invalid_argument(
                "acquire_full_buckets batch size must be positive");
 
        trk.register_access_factory([](auto &tracker) {
            auto *self = LIBTCSPC_OBJECT_FROM_TRACKER(acquire_full_buckets,
                                                      trk, tracker);
            return acquire_access([self] { self->halt(); });
        });
    }
 
    // Custom move because we have a mutex. Move only works when not running.
    ~acquire_full_buckets() = default;
 
    acquire_full_buckets(acquire_full_buckets const &) = delete;
    auto operator=(acquire_full_buckets const &) = delete;
 
    acquire_full_buckets(acquire_full_buckets &&other) noexcept
        : reader(std::move(other.reader)), bsource(std::move(other.bsource)),
          bsize(other.bsize), halted(other.halted),
          live_downstream(std::move(other.live_downstream)),
          batch_downstream(std::move(other.batch_downstream)),
          trk(std::move(other.trk)) {}
 
    auto operator=(acquire_full_buckets &&rhs) noexcept
        -> acquire_full_buckets & {
        reader = std::move(rhs.reader);
        bsource = std::move(rhs.bsource);
        bsize = rhs.bsize;
        halted = rhs.halted;
        live_downstream = std::move(rhs.live_downstream);
        batch_downstream = std::move(rhs.batch_downstream);
        trk = std::move(rhs.trk);
        return *this;
    }
 
    [[nodiscard]] auto introspect_node() const -> processor_info {
        return processor_info(this, "acquire_full_buckets");
    }
 
    [[nodiscard]] auto introspect_graph() const -> processor_graph {
        return merge_processor_graphs(
            live_downstream.introspect_graph().push_entry_point(this),
            batch_downstream.introspect_graph().push_entry_point(this));
    }
 
    void flush() {
        bucket<T> b;
        std::size_t filled = 0;
        {
            auto lock = std::unique_lock(halt_mutex);
            while (not halted) {
                lock.unlock();
                auto const start_time = std::chrono::steady_clock::now();
                if (b.empty()) {
                    b = bsource->bucket_of_size(bsize);
                    filled = 0;
                }
                auto const unfilled = std::span(b).subspan(filled);
                std::optional<std::size_t> const read = reader(unfilled);
                if (not read)
                    return flush_downstreams(std::move(b), filled);
                if constexpr (not std::is_same_v<LiveDownstream, null_sink>)
                    emit_live(b, filled, *read);
                filled += *read;
                if (filled == bsize) {
                    emit_batch(std::move(b));
                    b = {};
                }
                lock.lock();
                if (filled < bsize) {
                    halt_cv.wait_until(lock, start_time + slow_acq_sleep,
                                       [&] { return halted; });
                }
            }
        }
        throw acquisition_halted();
    }
};
 
} // namespace internal

template <typename T, typename Reader, typename Downstream>
auto acquire(Reader reader, std::shared_ptr<bucket_source<T>> buffer_provider,
             arg::batch_size<std::size_t> batch_size,
             access_tracker<acquire_access> tracker, Downstream downstream) {
    return internal::acquire<T, Reader, Downstream>(
        std::move(reader), std::move(buffer_provider), batch_size,
        std::move(tracker), std::move(downstream));
}

template <typename T, typename Reader, typename LiveDownstream,
          typename BatchDownstream>
auto acquire_full_buckets(Reader reader,
                          std::shared_ptr<bucket_source<T>> buffer_provider,
                          arg::batch_size<std::size_t> batch_size,
                          access_tracker<acquire_access> tracker,
                          LiveDownstream live_downstream,
                          BatchDownstream batch_downstream) {
    return internal::acquire_full_buckets<T, Reader, LiveDownstream,
                                          BatchDownstream>(
        std::move(reader), std::move(buffer_provider), batch_size,
        std::move(tracker), std::move(live_downstream),
        std::move(batch_downstream));
}

template <typename T> struct null_reader {
    auto operator()(std::span<T> /* buffer */) -> std::optional<std::size_t> {
        return std::nullopt;
    }
};

template <typename T> struct stuck_reader {
    auto operator()(std::span<T> /* buffer */) -> std::optional<std::size_t> {
        return 0;
    }
};
 
} // namespace tcspc