writer.h

#ifndef PROIO_WRITER_H
#define PROIO_WRITER_H

#include <condition_variable>
#include <cstring>
#include <mutex>
#include <string>
#include <thread>

#include "google/protobuf/descriptor.h"
#include "google/protobuf/descriptor.pb.h"
#include "google/protobuf/io/zero_copy_stream_impl.h"

#include "event.h"
#include "proio/proto/proio.pb.h"

namespace proio {
typedef proto::BucketHeader_CompType Compression;
const Compression LZ4 = proto::BucketHeader::LZ4;
const Compression GZIP = proto::BucketHeader::GZIP;
const Compression UNCOMPRESSED = proto::BucketHeader::NONE;

const std::size_t minBucketWriteWindow = 0x100000;

class BucketOutputStream : public google::protobuf::io::ZeroCopyOutputStream {
   public:
    BucketOutputStream();
    virtual ~BucketOutputStream();

    bool Next(void **data, int *size) override;
    void BackUp(int count) override;
    google::protobuf::int64 ByteCount() const override;
    bool AllowsAliasing() const override;

    uint8_t *Bytes();
    void Reset();
    void Reset(uint64_t size);
    void WriteTo(google::protobuf::io::ZeroCopyOutputStream *stream);
    void SetOffset(uint64_t offset);

   private:
    std::vector<uint8_t> bytes;
    uint64_t offset;
};

class Writer : public std::mutex {
   public:
    Writer(int fd);
    Writer(std::string filename);
    virtual ~Writer();

    void Flush();
    void Push(Event *event);
    void PushMetadata(std::string name, const std::string &data);
    void PushMetadata(std::string name, const char *data);
    void SetCompression(Compression alg = GZIP, int level = -1) {
        compression = alg;
        complevel = level;
    }
    void SetBucketDumpThreshold(uint64_t thres = 0x1000000) { bucketDumpThres = thres; }

   private:
    BucketOutputStream *bucket;
    google::protobuf::io::FileOutputStream *fileStream;
    uint64_t bucketEvents;
    Compression compression;
    int complevel;
    BucketOutputStream *compBucket;
    uint64_t bucketDumpThres;
    proto::BucketHeader *header;
    std::map<std::string, std::shared_ptr<const std::string>> metadata;
    std::set<const google::protobuf::FileDescriptor *> writtenFDs;

    void initBucket();

    std::thread streamWriteThread;
    typedef struct {
        bool isValid;

        BucketOutputStream *compBucket;
        proto::BucketHeader *header;
        google::protobuf::io::FileOutputStream *fileStream;

        std::mutex doJobMutex;
        std::condition_variable doJobCond;
        std::mutex workerReadyMutex;
        std::condition_variable workerReadyCond;
    } WriteJob;
    WriteJob streamWriteJob;
    std::unique_lock<std::mutex> workerReadyLock;

    static void streamWrite(WriteJob *job);
};

const uint8_t magicBytes[] = {0xe1, 0xc1, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                              0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00};

const class SerializationError : public std::exception {
    virtual const char *what() const throw() { return "Failed to serialize message"; }
} serializationError;

const class FileCreationError : public std::exception {
    virtual const char *what() const throw() { return "Failed to creating file for writing"; }
} fileCreationError;

const class LZ4FrameCreationError : public std::exception {
    virtual const char *what() const throw() { return "Failed to create LZ4 frame"; }
} lz4FrameCreationError;
}  // namespace proio

#endif  // PROIO_WRITER_H