Jungfraujoch/image_puller/ImagePuller.h

// SPDX-FileCopyrightText: 2025 Filip Leonarski, Paul Scherrer Institute <filip.leonarski@psi.ch>
// SPDX-License-Identifier: GPL-3.0-only

#pragma once

#include <atomic>
#include <chrono>
#include <string>
#include <memory>

#include "../common/ZMQWrappers.h"
#include "../frame_serialize/CBORStream2Deserializer.h"
#include "../common/ThreadSafeFIFO.h"
#include "../common/JfjochTCP.h"
#include "ReceiveRingBuffer.h"

struct PullerAckMessage {
    TCPFrameType ack_for = TCPFrameType::DATA;
    bool ok = true;
    bool fatal = false;
    uint64_t run_number = 0;
    uint32_t socket_number = 0;
    uint64_t image_number = 0;
    uint64_t processed_images = 0;
    TCPAckCode error_code = TCPAckCode::None;
    std::string error_text;
};

struct RawFrame {
    TcpFrameHeader header{};
    bool end = false;
};

struct ImagePullerOutput {
    // The message as it came off the socket, in the puller's ring buffer. It stays there for as long as
    // any copy of this struct is alive - the deserialised image in cbor points into it, and republish
    // sends it as is.
    std::shared_ptr<ReceiveBufferRegion> raw;
    std::shared_ptr<RawFrame> tcp_msg;
    std::shared_ptr<CBORStream2DeserializerOutput> cbor;
    // When the message came off the socket - before it was queued or deserialised, so a consumer that
    // falls behind does not shift it.
    std::chrono::steady_clock::time_point received{};
};

struct ImagePullerStatistics {
    size_t outside_fifo_max_util = 0;
};

class ImagePuller {
protected:
    ThreadSafeFIFO<ImagePullerOutput> outside_fifo; // assuming 1 file is <1000 files, this should be enough to buffer
    // Every received message is held here until consumed; the queues count messages, this bounds bytes.
    // Size 0 means no ring: messages are held on the heap, bounded by the queues alone, as they always
    // were - the broker's own DECTRIS puller runs that way, jfjoch_writer sets a size.
    ReceiveRingBuffer ring;

    // Bytes republish holds right now. Republish keeps the OLDEST not yet sent images, and the ring is
    // reclaimed in arrival order, so an image waiting for a slow republish consumer pins everything
    // received after it - the writer would then run at the republish rate. What republish may hold is
    // therefore capped at a share of the ring; an image over the cap is not republished, exactly as one
    // that finds the republish queue full. Without a ring there is no cap.
    std::atomic<size_t> repub_bytes{0};
    bool RepubCanHold(size_t bytes) const {
        return ring.GetCapacity() == 0 || repub_bytes.load() + bytes <= ring.GetCapacity() / 4;
    }
public:
    constexpr static size_t DefaultQueueSize = 1000;
    constexpr static auto PollTimeout = std::chrono::milliseconds(100);

    explicit ImagePuller(size_t fifo_size = DefaultQueueSize, size_t ring_size = 0);
    virtual ~ImagePuller() = default;

    [[nodiscard]] std::optional<ImagePullerOutput> PollImage(std::chrono::microseconds timeout = PollTimeout);
    [[nodiscard]] size_t GetMaxFifoUtilization() const;
    [[nodiscard]] size_t GetCurrentFifoUtilization() const;
    [[nodiscard]] size_t GetRingBufferCapacity() const;
    [[nodiscard]] size_t GetMaxRingBufferUtilization() const;

    virtual bool SupportsAck() const { return false; }
    virtual bool SendAck(const PullerAckMessage &) { return false; }

    virtual void Disconnect() = 0;
};