SyncStream.hxx

Go to the documentation of this file.

 
#ifndef _UTILS_SYNCSTREAM_HXX_
#define _UTILS_SYNCSTREAM_HXX_
 
#include <memory>
#include <stdint.h>
#include <string.h>
#include <unistd.h>
 
#include "utils/macros.h"
#include "utils/logging.h"
 
class SyncStream
{
public:
    virtual ~SyncStream()
    {
    }
 
    virtual ssize_t write(const void *data, size_t len) = 0;
 
    virtual int finalize(int status)
    {
        return 0;
    }
 
    ssize_t write_all(const void *data, size_t len)
    {
        auto *d = to_8(data);
        size_t written = 0;
        while (len > 0)
        {
            auto ret = write(d, len);
            if (ret < 0)
            {
                return ret;
            }
            if (ret == 0)
            {
                return written;
            }
            written += ret;
            d += ret;
            len -= ret;
        }
        return written;
    }
 
protected:
    static const uint8_t *to_8(const void *d)
    {
        return static_cast<const uint8_t *>(d);
    }
 
    static uint8_t *to_8(void *d)
    {
        return static_cast<uint8_t *>(d);
    }
};
 
class HeaderStream : public SyncStream
{
public:
    HeaderStream(void *header, size_t header_size)
        : data_(to_8(header))
        , remaining_(header_size)
    {
    }
 
    ssize_t write(const void *data, size_t len) override
    {
        if (remaining_ == 0)
        {
            return 0;
        }
        if (len > remaining_)
        {
            len = remaining_;
        }
        memcpy(data_, data, len);
        remaining_ -= len;
        data_ += len;
        return len;
    }
 
private:
    uint8_t *data_;
    size_t remaining_;
};
 
class WrappedStream : public SyncStream
{
public:
    WrappedStream(SyncStream *delegate)
        : delegate_(delegate)
    {
    }
 
    void set_delegate(SyncStream *delegate)
    {
        if (delegate_)
        {
            // TODO: discards error value.
            delegate_->finalize(-1);
        }
        delegate_.reset(delegate);
    }
 
    int finalize(int status) override
    {
        int ret = 0;
        if (delegate_)
        {
            ret = delegate_->finalize(status);
            delegate_.reset();
        }
        return ret;
    }
 
protected:
    std::unique_ptr<SyncStream> delegate_;
};
 
class MaxLengthStream : public WrappedStream
{
public:
    MaxLengthStream(size_t length, SyncStream *delegate)
        : WrappedStream(delegate)
        , remaining_(length)
    {
    }
 
    ~MaxLengthStream() {
        LOG(INFO, "deleting maxlengthstream remaining=%d", (int) remaining_);
    }
    
    ssize_t write(const void *data, size_t len) override
    {
        if (remaining_ == 0)
        {
            return 0;
        }
        if (len > remaining_)
        {
            len = remaining_;
        }
        ssize_t ret = delegate_->write(data, len);
        if (ret <= 0)
        {
            return ret;
        }
        remaining_ -= ret;
        return ret;
    }
 
private:
    size_t remaining_;
};
 
class MinWriteStream : public WrappedStream
{
public:
    MinWriteStream(
        unsigned min_write_length, uint8_t fill_byte, SyncStream *delegate)
        : WrappedStream(delegate)
        , buffer_(nullptr)
        , bufLength_(0)
        , minWriteLength_(min_write_length)
        , fillByte_(fill_byte)
        , needsFill_(1)
    {
    }
 
    MinWriteStream(
        unsigned min_write_length, SyncStream *delegate)
        : WrappedStream(delegate)
        , buffer_(nullptr)
        , bufLength_(0)
        , minWriteLength_(min_write_length)
        , needsFill_(0)
    {
    }
    
    ~MinWriteStream()
    {
        LOG(INFO, "deleting minwritestream l=%d", (int) minWriteLength_);
        delete[] buffer_;
    }
 
    ssize_t write(const void *data, size_t len) override
    {
        if (len == 0)
            return 0; // not sure what to do here
        if (bufLength_)
        {
            // There is some data in the buffer. Try to complete and flush it.
            HASSERT(buffer_);
            size_t cp = len;
            if (cp + bufLength_ > minWriteLength_)
            {
                cp = minWriteLength_ - bufLength_;
            }
            memcpy(buffer_ + bufLength_, data, cp);
            bufLength_ += cp;
            if (bufLength_ >= minWriteLength_)
            {
                // Flush
                auto ret = delegate_->write_all(buffer_, bufLength_);
                if (ret <= 0)
                    return ret;
                HASSERT(ret == (int)bufLength_);
                bufLength_ = 0;
            }
            return cp;
        }
        if (len < minWriteLength_)
        {
            // Too small write. Must copy stuff to the buffer.
            if (!buffer_)
            {
                buffer_ = (uint8_t *)new uint32_t[minWriteLength_ / 4];
            }
            memcpy(buffer_, data, len);
            bufLength_ = len;
            return len;
        }
        auto ret = delegate_->write(data, len);
        return ret;
    }
 
    int finalize(int status) override
    {
        if (bufLength_)
        {
            memset(
                buffer_ + bufLength_, fillByte_, minWriteLength_ - bufLength_);
            auto ret = delegate_->write(buffer_, minWriteLength_);
            if (ret < 0)
                return ret;
            if (ret == 0)
                return -1; // dropped data.
        }
        return delegate_->finalize(status);
    }
 
private:
    uint8_t *buffer_;
    unsigned bufLength_;
    unsigned minWriteLength_;
    uint8_t fillByte_;
    uint8_t needsFill_;
};
 
class StringAppendStream : public SyncStream {
public:
    StringAppendStream(std::string* output)
        : output_(output) {}
 
    void set_output(std::string* output) {
        output_ = output;
    }
    
    ssize_t write(const void *data, size_t len) override {
        output_->append((const char*)data, len);
        return len;
    }
 
private:
    std::string* output_;
};
 
class DelegateStream : public SyncStream
{
public:
    ssize_t write(const void *data, size_t len) override
    {
        if (len == 0)
            return 0;
        while (true)
        {
            if (!delegate_)
            {
                return 0;
            }
            ssize_t ret = delegate_->write(data, len);
            if (ret != 0)
            {
                // Both positive (data actually written) as well as negative
                // (error) will be returned directly.
                return ret;
            }
            on_eof();
        }
    }
 
protected:
    virtual void on_eof()
    {
        delegate_->finalize(0);
        delegate_.reset();
    }
 
    std::unique_ptr<SyncStream> delegate_;
};
 
#endif // _UTILS_SYNCSTREAM_HXX_