#pragma once

#include <balancer/kernel/io/iobase.h>
#include <balancer/kernel/memory/chunks.h>

#include <util/datetime/base.h>

namespace NSrvKernel {
    class TCountInput : public IIoInput {
    public:
        TCountInput(IIoInput* slave) noexcept
            : Slave_(slave)
        {}

        size_t Readed() const noexcept {
            return Readed_;
        }

    private:
        TError DoRecv(TChunkList& lst, TInstant deadline) noexcept override {
            Y_PROPAGATE_ERROR(Slave_->Recv(lst, deadline));
            Readed_ += lst.size();
            return {};
        }

        void UnRecv(TChunkList&& lst) noexcept override {
            Readed_ -= lst.size();
            IIoInput::UnRecv(std::move(lst));
        }

    private:
        IIoInput* const Slave_ = nullptr;
        size_t Readed_ = 0;
    };
}
