terebi/Sources/Core/network_input.hpp

#pragma once

extern "C" {
#include <libavformat/avformat.h>
}
#include <array>
#include <algorithm>
#include <atomic>
#include <chrono>
#include <cmath>
#include <cerrno>
#include <condition_variable>
#include <cstdint>
#include <limits>
#include <mutex>
#include <new>
#include <thread>

namespace playback {
// Only network inputs need a reader: a blocked socket must never own the player worker.
// Format metadata is immutable after start(); the reader exclusively owns demux and seek I/O.
class NetworkInput {
public:
    static constexpr std::size_t packetLimit = 64;
    static constexpr std::size_t byteLimit = 16 * 1024 * 1024;
    AVFormatContext *format = nullptr;

    NetworkInput(const std::atomic<bool> &stopping, const std::atomic<uint64_t> &epoch,
                 uint64_t activeEpoch, std::condition_variable &playerWake)
        : stopping(stopping), epoch(epoch), activeEpoch(activeEpoch), playerWake(playerWake) {}
    ~NetworkInput() {
        stop();
        for (auto &lane : lanes)
            for (auto *packet : lane.packets) av_packet_free(&packet);
        av_packet_free(&scratch);
        avformat_close_input(&format);
    }
    NetworkInput(const NetworkInput &) = delete;
    NetworkInput &operator=(const NetworkInput &) = delete;

    static int interrupt(void *opaque) {
        auto &input = *static_cast<NetworkInput *>(opaque);
        return input.cancelled.load(std::memory_order_relaxed) ||
            input.stopping.load(std::memory_order_relaxed) ||
            input.epoch.load(std::memory_order_relaxed) != input.activeEpoch ||
            input.serial.load(std::memory_order_acquire) != input.operationSerial.load(std::memory_order_relaxed) ||
            now() >= input.deadline.load(std::memory_order_relaxed);
    }
    void armIO() {
        operationSerial.store(serial.load(std::memory_order_acquire), std::memory_order_relaxed);
        deadline.store(now() + 15'000'000'000LL, std::memory_order_relaxed);
    }
    void start(std::array<int, 2> indices, double timelineOrigin) {
        origin = timelineOrigin;
        for (int slot = 0; slot < 2; ++slot) {
            lanes[slot].index = indices[slot];
            if (indices[slot] < 0) continue;
            for (auto &packet : lanes[slot].packets) {
                packet = av_packet_alloc();
                if (!packet) throw std::bad_alloc();
            }
        }
        scratch = av_packet_alloc();
        if (!scratch) throw std::bad_alloc();
        reader = std::thread([this] { run(); });
    }
    void stop() {
        cancelled.store(true, std::memory_order_relaxed);
        changed.notify_all();
        if (reader.joinable()) reader.join();
    }
    void seek(int streamIndex, int64_t timestamp) {
        std::lock_guard lock(mutex);
        for (auto &lane : lanes) clear(lane);
        scratchSlot = -1;
        ended = false;
        failure = 0;
        seekIndex = streamIndex;
        seekTimestamp = timestamp;
        seeking = true;
        serial.fetch_add(1, std::memory_order_release);
        changed.notify_all();
    }
    void discardAudio(bool discard) {
        std::lock_guard lock(mutex);
        skipAudio = discard;
        if (discard) clear(lanes[0]);
        if (discard && scratchSlot == 0) scratchSlot = -1;
        changed.notify_all();
    }
    void finish() {
        std::lock_guard lock(mutex);
        for (auto &lane : lanes) clear(lane);
        scratchSlot = -1;
        ended = true;
        seeking = false;
        failure = 0;
        serial.fetch_add(1, std::memory_order_release);
        changed.notify_all();
    }
    bool seekPending() const {
        std::lock_guard lock(mutex);
        return seeking;
    }
    int error() const {
        std::lock_guard lock(mutex);
        return failure;
    }
    double bufferedEnd(int slot) const {
        std::lock_guard lock(mutex);
        if (failure || seeking) return 0;
        const auto &lane = lanes[slot];
        const double queued = lane.count ?
            lane.ends[(lane.first + lane.count - 1) % packetLimit] : 0;
        return std::max(queued, scratchSlot == slot ? scratchEnd : 0);
    }
    // 1: packet, 0: not yet available, AVERROR_EOF: this input is drained.
    int take(int slot, AVPacket *destination, double &bufferedEnd) {
        std::lock_guard lock(mutex);
        if (failure) return failure;
        bufferedEnd = 0;
        auto &lane = lanes[slot];
        if (lane.count) {
            auto *packet = lane.packets[lane.first];
            bufferedEnd = lane.ends[lane.first];
            lane.bytes -= packetBytes(packet);
            av_packet_move_ref(destination, packet);
            lane.first = (lane.first + 1) % packetLimit;
            --lane.count;
            changed.notify_one();
            return 1;
        }
        return ended ? AVERROR_EOF : 0;
    }

private:
    struct Lane {
        std::array<AVPacket *, packetLimit> packets{};
        std::array<double, packetLimit> ends{};
        int index = -1;
        std::size_t first = 0, count = 0, bytes = 0;
    };
    const std::atomic<bool> &stopping;
    const std::atomic<uint64_t> &epoch;
    const uint64_t activeEpoch;
    std::condition_variable &playerWake;
    mutable std::mutex mutex;
    std::condition_variable changed;
    std::array<Lane, 2> lanes;
    AVPacket *scratch = nullptr;
    std::atomic<bool> cancelled{false};
    std::atomic<uint64_t> serial{0}, operationSerial{0};
    std::atomic<int64_t> deadline{std::numeric_limits<int64_t>::max()};
    bool seeking = false, ended = false, skipAudio = false;
    int seekIndex = -1, failure = 0;
    int64_t seekTimestamp = 0;
    int scratchSlot = -1;
    double scratchEnd = 0, origin = 0;
    std::thread reader;

    static int64_t now() {
        return std::chrono::duration_cast<std::chrono::nanoseconds>(
            std::chrono::steady_clock::now().time_since_epoch()).count();
    }
    static std::size_t packetBytes(const AVPacket *packet) {
        if (packet->size < 0) return byteLimit + 1;
        std::size_t bytes = static_cast<std::size_t>(packet->size);
        if (bytes > byteLimit) return byteLimit + 1;
        for (int i = 0; i < packet->side_data_elems; ++i) {
            if (packet->side_data[i].size > byteLimit - bytes) return byteLimit + 1;
            bytes += packet->side_data[i].size;
        }
        return bytes;
    }
    double packetEnd(int slot, const AVPacket *packet) const {
        if (packet->pts == AV_NOPTS_VALUE || (packet->flags & AV_PKT_FLAG_DISCARD)) return 0;
        int64_t timestamp = packet->pts;
        if (slot == 1) {
            // PTS alone can jump to a future reference frame before its B-frames arrive.
            // Unknown DTS contributes no encoded-video estimate; decoded frames still count.
            if (packet->dts == AV_NOPTS_VALUE) return 0;
            timestamp = std::min(timestamp, packet->dts);
        }
        const double tick = av_q2d(format->streams[lanes[slot].index]->time_base);
        const double end = static_cast<double>(timestamp) * tick - origin +
            static_cast<double>(std::max<int64_t>(packet->duration, 0)) * tick;
        return std::isfinite(end) && end > 0 ? end : 0;
    }
    static void clear(Lane &lane) {
        while (lane.count) {
            av_packet_unref(lane.packets[lane.first]);
            lane.first = (lane.first + 1) % packetLimit;
            --lane.count;
        }
        lane.first = lane.bytes = 0;
    }
    bool stoppingNow() const {
        return cancelled.load(std::memory_order_relaxed) ||
            stopping.load(std::memory_order_relaxed) || epoch.load(std::memory_order_relaxed) != activeEpoch;
    }
    void run() noexcept {
        try {
            while (!stoppingNow()) {
                uint64_t current;
                int index;
                int64_t timestamp;
                bool doSeek;
                {
                    std::unique_lock lock(mutex);
                    changed.wait_for(lock, std::chrono::milliseconds(20), [this] {
                        return stoppingNow() || seeking || !ended;
                    });
                    if (stoppingNow()) return;
                    if (ended && !seeking) continue;
                    current = serial.load(std::memory_order_acquire);
                    doSeek = seeking;
                    index = seekIndex;
                    timestamp = seekTimestamp;
                }
                operationSerial.store(current, std::memory_order_relaxed);
                deadline.store(now() + 15'000'000'000LL, std::memory_order_relaxed);
                if (doSeek) {
                    const int result = index >= 0 ?
                        av_seek_frame(format, index, timestamp, AVSEEK_FLAG_BACKWARD) :
                        avformat_seek_file(format, -1, std::numeric_limits<int64_t>::min(), timestamp, timestamp, 0);
                    if (result >= 0 && format->pb) {
                        // An interrupted read leaves AVIO's error sticky even after a successful format seek.
                        format->pb->error = 0;
                        format->pb->eof_reached = 0;
                    }
                    std::lock_guard lock(mutex);
                    if (serial.load(std::memory_order_acquire) != current) continue;
                    seeking = false;
                    if (result < 0) { failure = result; ended = true; }
                    playerWake.notify_one();
                    if (result < 0) continue;
                }
                const int result = av_read_frame(format, scratch);
                std::unique_lock lock(mutex);
                if (stoppingNow()) return;
                if (serial.load(std::memory_order_acquire) != current) {
                    av_packet_unref(scratch);
                    continue;
                }
                if (result == AVERROR(EAGAIN)) {
                    changed.wait_for(lock, std::chrono::milliseconds(5));
                    continue;
                }
                if (result < 0) {
                    av_packet_unref(scratch);
                    ended = true;
                    if (result != AVERROR_EOF) failure = result;
                    playerWake.notify_one();
                    continue;
                }
                int slot = -1;
                for (int i = 0; i < 2; ++i)
                    if (lanes[i].index >= 0 && scratch->stream_index == lanes[i].index) slot = i;
                if (slot < 0 || (slot == 0 && skipAudio)) { av_packet_unref(scratch); continue; }
                auto &lane = lanes[slot];
                const auto bytes = packetBytes(scratch);
                if (bytes > byteLimit) {
                    failure = AVERROR(ENOBUFS);
                    ended = true;
                    av_packet_unref(scratch);
                    playerWake.notify_one();
                    continue;
                }
                scratchSlot = slot;
                scratchEnd = packetEnd(slot, scratch);
                changed.wait(lock, [&] {
                    return stoppingNow() || serial.load(std::memory_order_acquire) != current ||
                        (slot == 0 && skipAudio) ||
                        (lane.count < packetLimit && lane.bytes + bytes <= byteLimit);
                });
                if (stoppingNow()) return;
                if (serial.load(std::memory_order_acquire) != current || (slot == 0 && skipAudio)) {
                    av_packet_unref(scratch);
                    scratchSlot = -1;
                    continue;
                }
                lane.bytes += bytes;
                lane.ends[(lane.first + lane.count) % packetLimit] = scratchEnd;
                av_packet_move_ref(lane.packets[(lane.first + lane.count) % packetLimit], scratch);
                ++lane.count;
                scratchSlot = -1;
                playerWake.notify_one();
            }
        } catch (...) {
            std::lock_guard lock(mutex);
            failure = AVERROR(ENOMEM);
            ended = true;
            playerWake.notify_one();
        }
    }
};
}