path: root/library/cpp/actors/core/event_pb.h

#pragma once

#include "event.h"
#include "event_load.h"

#include <google/protobuf/io/zero_copy_stream.h>
#include <google/protobuf/arena.h>
#include <library/cpp/actors/protos/actors.pb.h>
#include <util/generic/deque.h>
#include <util/system/context.h>
#include <util/system/filemap.h>
#include <array>

namespace NActors {
 
    class TRopeStream : public NProtoBuf::io::ZeroCopyInputStream { 
        TRope::TConstIterator Iter; 
        const size_t Size; 
 
    public:
        TRopeStream(TRope::TConstIterator iter, size_t size) 
            : Iter(iter) 
            , Size(size) 
        {} 

        bool Next(const void** data, int* size) override; 
        void BackUp(int count) override; 
        bool Skip(int count) override; 
        int64_t ByteCount() const override {
            return TotalByteCount;
        }

    private: 
        int64_t TotalByteCount = 0;
    };

    class TChunkSerializer : public NProtoBuf::io::ZeroCopyOutputStream { 
    public:
        TChunkSerializer() = default; 
        virtual ~TChunkSerializer() = default; 
 
        virtual bool WriteRope(const TRope *rope) = 0; 
        virtual bool WriteString(const TString *s) = 0; 
    };

    class TAllocChunkSerializer final : public TChunkSerializer { 
    public:
        bool Next(void** data, int* size) override; 
        void BackUp(int count) override; 
        int64_t ByteCount() const override {
            return Buffers->GetSize();
        }
        bool WriteAliasedRaw(const void* data, int size) override; 

        // WARNING: these methods require owner to retain ownership and immutability of passed objects 
        bool WriteRope(const TRope *rope) override; 
        bool WriteString(const TString *s) override; 
 
        inline TIntrusivePtr<TEventSerializedData> Release(bool extendedFormat) { 
            if (extendedFormat) { 
                Buffers->SetExtendedFormat(); 
            } 
            return std::move(Buffers); 
        }

    protected:
        TIntrusivePtr<TEventSerializedData> Buffers = new TEventSerializedData; 
        TRope Backup; 
    };

    class TCoroutineChunkSerializer final : public TChunkSerializer, protected ITrampoLine { 
    public:
        using TChunk = std::pair<const char*, size_t>; 
 
        TCoroutineChunkSerializer();
        ~TCoroutineChunkSerializer();

        void SetSerializingEvent(const IEventBase *event); 
        void Abort(); 
        std::pair<TChunk*, TChunk*> FeedBuf(void* data, size_t size); 
        bool IsComplete() const {
            return !Event; 
        }
        bool IsSuccessfull() const {
            return SerializationSuccess;
        }
        const IEventBase *GetCurrentEvent() const { 
            return Event; 
        } 

        bool Next(void** data, int* size) override;
        void BackUp(int count) override;
        int64_t ByteCount() const override {
            return TotalSerializedDataSize;
        }
        bool WriteAliasedRaw(const void* data, int size) override;
        bool AllowsAliasing() const override;

        bool WriteRope(const TRope *rope) override; 
        bool WriteString(const TString *s) override; 
 
    protected:
        void DoRun() override;
        void Resume();
        bool Produce(const void *data, size_t size); 

        i64 TotalSerializedDataSize;
        TMappedAllocation Stack;
        TContClosure SelfClosure;
        TContMachineContext InnerContext;
        TContMachineContext *BufFeedContext = nullptr; 
        char *BufferPtr; 
        size_t SizeRemain; 
        static constexpr size_t MaxChunks = 16; 
        TChunk Chunks[MaxChunks]; 
        size_t NumChunks = 0; 
        const IEventBase *Event = nullptr; 
        bool CancelFlag = false; 
        bool AbortFlag; 
        bool SerializationSuccess; 
        bool Finished = false; 
    };

#ifdef ACTORLIB_HUGE_PB_SIZE 
    static const size_t EventMaxByteSize = 140 << 20; // (140MB) 
#else 
    static const size_t EventMaxByteSize = 67108000; 
#endif 
 
    template <typename TEv, typename TRecord /*protobuf record*/, ui32 TEventType, typename TRecHolder>
    class TEventPBBase: public TEventBase<TEv, TEventType> , public TRecHolder {
        // a vector of data buffers referenced by record; if filled, then extended serialization mechanism applies 
        TVector<TRope> Payload; 
 
    public:
        using TRecHolder::Record;

    public:
        using ProtoRecordType = TRecord;

        TEventPBBase() = default;

        explicit TEventPBBase(const TRecord& rec)
        {
            Record = rec;
        }

        explicit TEventPBBase(TRecord&& rec)
        {
            Record = std::move(rec);
        }

        TString ToStringHeader() const override {
            return Record.GetTypeName();
        }

        TString ToString() const override {
            return Record.ShortDebugString(); 
        }

        bool IsSerializable() const override { 
            return true; 
        } 
 
        bool IsExtendedFormat() const override { 
            return static_cast<bool>(Payload); 
        } 
 
        bool SerializeToArcadiaStream(TChunkSerializer* chunker) const override {
            // serialize payload first 
            if (Payload) { 
                void *data; 
                int size = 0; 
                auto append = [&](const char *p, size_t len) { 
                    while (len) { 
                        if (size) { 
                            const size_t numBytesToCopy = std::min<size_t>(size, len); 
                            memcpy(data, p, numBytesToCopy); 
                            data = static_cast<char*>(data) + numBytesToCopy; 
                            size -= numBytesToCopy; 
                            p += numBytesToCopy; 
                            len -= numBytesToCopy; 
                        } else if (!chunker->Next(&data, &size)) { 
                            return false; 
                        } 
                    } 
                    return true; 
                }; 
                auto appendNumber = [&](size_t number) { 
                    char buf[MaxNumberBytes]; 
                    return append(buf, SerializeNumber(number, buf)); 
                }; 
                char marker = PayloadMarker; 
                append(&marker, 1); 
                if (!appendNumber(Payload.size())) { 
                    return false; 
                } 
                for (const TRope& rope : Payload) { 
                    if (!appendNumber(rope.GetSize())) { 
                        return false; 
                    } 
                    if (rope) { 
                        if (size) { 
                            chunker->BackUp(std::exchange(size, 0)); 
                        } 
                        if (!chunker->WriteRope(&rope)) { 
                            return false; 
                        } 
                    } 
                } 
                if (size) { 
                    chunker->BackUp(size); 
                } 
            } 
 
            return Record.SerializeToZeroCopyStream(chunker);
        }

        ui32 CalculateSerializedSize() const override {
            ssize_t result = Record.ByteSize(); 
            if (result >= 0 && Payload) { 
                ++result; // marker 
                char buf[MaxNumberBytes]; 
                result += SerializeNumber(Payload.size(), buf); 
                for (const TRope& rope : Payload) { 
                    result += SerializeNumber(rope.GetSize(), buf); 
                    result += rope.GetSize(); 
                } 
            } 
            return result;
        }

        static IEventBase* Load(TIntrusivePtr<TEventSerializedData> input) {
            THolder<TEventPBBase> ev(new TEv());
            if (!input->GetSize()) { 
                Y_PROTOBUF_SUPPRESS_NODISCARD ev->Record.ParseFromString(TString());
            } else { 
                TRope::TConstIterator iter = input->GetBeginIter(); 
                ui64 size = input->GetSize(); 
 
                if (input->IsExtendedFormat()) { 
                    // check marker 
                    if (!iter.Valid() || *iter.ContiguousData() != PayloadMarker) { 
                        Y_FAIL("invalid event"); 
                    } 
                    // skip marker 
                    iter += 1; 
                    --size; 
                    // parse number of payload ropes 
                    size_t numRopes = DeserializeNumber(iter, size); 
                    if (numRopes == Max<size_t>()) { 
                        Y_FAIL("invalid event"); 
                    } 
                    while (numRopes--) { 
                        // parse length of the rope 
                        const size_t len = DeserializeNumber(iter, size); 
                        if (len == Max<size_t>() || size < len) { 
                            Y_FAIL("invalid event len# %zu size# %" PRIu64, len, size); 
                        } 
                        // extract the rope 
                        TRope::TConstIterator begin = iter; 
                        iter += len; 
                        size -= len; 
                        ev->Payload.emplace_back(begin, iter); 
                    } 
                } 
 
                // parse the protobuf 
                TRopeStream stream(iter, size); 
                if (!ev->Record.ParseFromZeroCopyStream(&stream)) { 
                    Y_FAIL("Failed to parse protobuf event type %" PRIu32 " class %s", TEventType, TypeName(ev->Record).data());
                } 
            }
            ev->CachedByteSize = input->GetSize(); 
            return ev.Release(); 
        }

        size_t GetCachedByteSize() const { 
            if (CachedByteSize == 0) { 
                CachedByteSize = CalculateSerializedSize(); 
            }
            return CachedByteSize;
        }

        ui32 CalculateSerializedSizeCached() const override { 
            return GetCachedByteSize(); 
        } 
 
        void InvalidateCachedByteSize() {
            CachedByteSize = 0;
        }

    public: 
        void ReservePayload(size_t size) {
            Payload.reserve(size);
        }

        ui32 AddPayload(TRope&& rope) { 
            const ui32 id = Payload.size(); 
            Payload.push_back(std::move(rope)); 
            InvalidateCachedByteSize(); 
            return id; 
        } 
 
        const TRope& GetPayload(ui32 id) const { 
            Y_VERIFY(id < Payload.size()); 
            return Payload[id]; 
        } 
 
        ui32 GetPayloadCount() const { 
            return Payload.size(); 
        } 
 
        void StripPayload() { 
            Payload.clear(); 
        } 
 
    protected: 
        mutable size_t CachedByteSize = 0;
 
        static constexpr char PayloadMarker = 0x07; 
        static constexpr size_t MaxNumberBytes = (sizeof(size_t) * CHAR_BIT + 6) / 7; 
 
        static size_t SerializeNumber(size_t num, char *buffer) { 
            char *begin = buffer; 
            do { 
                *buffer++ = (num & 0x7F) | (num >= 128 ? 0x80 : 0x00); 
                num >>= 7; 
            } while (num); 
            return buffer - begin; 
        } 
 
        static size_t DeserializeNumber(const char **ptr, const char *end) { 
            const char *p = *ptr; 
            size_t res = 0; 
            size_t offset = 0; 
            for (;;) { 
                if (p == end) { 
                    return Max<size_t>(); 
                } 
                const char byte = *p++; 
                res |= (static_cast<size_t>(byte) & 0x7F) << offset; 
                offset += 7; 
                if (!(byte & 0x80)) { 
                    break; 
                } 
            } 
            *ptr = p; 
            return res; 
        } 
 
        static size_t DeserializeNumber(TRope::TConstIterator& iter, ui64& size) { 
            size_t res = 0; 
            size_t offset = 0; 
            for (;;) { 
                if (!iter.Valid()) { 
                    return Max<size_t>(); 
                } 
                const char byte = *iter.ContiguousData(); 
                iter += 1; 
                --size; 
                res |= (static_cast<size_t>(byte) & 0x7F) << offset; 
                offset += 7; 
                if (!(byte & 0x80)) { 
                    break; 
                } 
            } 
            return res; 
        } 
    };

    // Protobuf record not using arena
    template <typename TRecord>
    struct TRecordHolder {
        TRecord Record;
    };

    // Protobuf arena and a record allocated on it
    template <typename TRecord, size_t InitialBlockSize, size_t MaxBlockSize>
    struct TArenaRecordHolder {
        google::protobuf::Arena PbArena;
        TRecord& Record;

        static const google::protobuf::ArenaOptions GetArenaOptions() {
            google::protobuf::ArenaOptions opts;
            opts.initial_block_size = InitialBlockSize;
            opts.max_block_size = MaxBlockSize;
            return opts;
        }

        TArenaRecordHolder()
            : PbArena(GetArenaOptions())
            , Record(*google::protobuf::Arena::CreateMessage<TRecord>(&PbArena))
        {}
    };

    template <typename TEv, typename TRecord, ui32 TEventType>
    class TEventPB : public TEventPBBase<TEv, TRecord, TEventType, TRecordHolder<TRecord> > {
        typedef TEventPBBase<TEv, TRecord, TEventType, TRecordHolder<TRecord> > TPbBase;
        // NOTE: No extra fields allowed: TEventPB must be a "template typedef"
    public:
        using TPbBase::TPbBase;
    };

    template <typename TEv, typename TRecord, ui32 TEventType, size_t InitialBlockSize = 512, size_t MaxBlockSize = 16*1024>
    using TEventPBWithArena = TEventPBBase<TEv, TRecord, TEventType, TArenaRecordHolder<TRecord, InitialBlockSize, MaxBlockSize> >;

    template <typename TEv, typename TRecord, ui32 TEventType>
    class TEventShortDebugPB: public TEventPB<TEv, TRecord, TEventType> {
    public:
        using TBase = TEventPB<TEv, TRecord, TEventType>;
        TEventShortDebugPB() = default;
        explicit TEventShortDebugPB(const TRecord& rec)
            : TBase(rec)
        {
        }
        explicit TEventShortDebugPB(TRecord&& rec)
            : TBase(std::move(rec))
        {
        }
        TString ToString() const override {
            return TypeName<TEv>() + " { " + TBase::Record.ShortDebugString() + " }";
        }
    };

    template <typename TEv, typename TRecord, ui32 TEventType>
    class TEventPreSerializedPB: public TEventPB<TEv, TRecord, TEventType> {
    protected:
        using TBase = TEventPB<TEv, TRecord, TEventType>;
        using TSelf = TEventPreSerializedPB<TEv, TRecord, TEventType>;
        using TBase::Record;

    public:
        TString PreSerializedData; // already serialized PB data (using message::SerializeToString)

        TEventPreSerializedPB() = default;

        explicit TEventPreSerializedPB(const TRecord& rec)
            : TBase(rec)
        {
        }

        explicit TEventPreSerializedPB(TRecord&& rec)
            : TBase(std::move(rec))
        {
        }

        // when remote event received locally this method will merge preserialized data
        const TRecord& GetRecord() {
            TRecord& base(TBase::Record);
            if (!PreSerializedData.empty()) {
                TRecord copy;
                Y_PROTOBUF_SUPPRESS_NODISCARD copy.ParseFromString(PreSerializedData);
                copy.MergeFrom(base);
                base.Swap(&copy);
                PreSerializedData.clear();
            }
            return TBase::Record;
        }

        const TRecord& GetRecord() const {
            return const_cast<TSelf*>(this)->GetRecord();
        }

        TRecord* MutableRecord() {
            GetRecord(); // Make sure PreSerializedData is parsed
            return &(TBase::Record);
        }

        TString ToString() const override {
            return GetRecord().ShortDebugString(); 
        }

        bool SerializeToArcadiaStream(TChunkSerializer* chunker) const override {
            return chunker->WriteString(&PreSerializedData) && TBase::SerializeToArcadiaStream(chunker);
        }

        ui32 CalculateSerializedSize() const override {
            return PreSerializedData.size() + TBase::CalculateSerializedSize();
        }

        size_t GetCachedByteSize() const {
            return PreSerializedData.size() + TBase::GetCachedByteSize();
        }
 
        ui32 CalculateSerializedSizeCached() const override { 
            return GetCachedByteSize(); 
        } 
    };

    inline TActorId ActorIdFromProto(const NActorsProto::TActorId& actorId) {
        return TActorId(actorId.GetRawX1(), actorId.GetRawX2());
    }

    inline void ActorIdToProto(const TActorId& src, NActorsProto::TActorId* dest) {
        Y_VERIFY_DEBUG(dest);
        dest->SetRawX1(src.RawX1());
        dest->SetRawX2(src.RawX2());
    }
}