DataBuffer.hxx

Go to the documentation of this file.

 
#ifndef _UTILS_DATABUFFER_HXX_
#define _UTILS_DATABUFFER_HXX_
 
#include "utils/Buffer.hxx"
#include "utils/LinkedObject.hxx"
#include "utils/macros.h"
 
#ifdef GTEST
// #define DEBUG_DATA_BUFFER_FREE
#endif
 
class DataBufferPool;
 
#ifdef DEBUG_DATA_BUFFER_FREE
class DataBuffer;
static void check_db_ownership(DataBuffer *p);
#endif
 
class DataBuffer : public Buffer<uint8_t[]>
{
public:
    void set_size(uint16_t s)
    {
        size_ = s;
    }
 
    void set_next(DataBuffer *n)
    {
        QMember::next = n;
    }
 
    DataBuffer *next()
    {
        return static_cast<DataBuffer *>(QMember::next);
    }
 
    uint8_t *data()
    {
        return static_cast<uint8_t *>(&data_[0]);
    }
 
    DataBuffer *ref()
    {
        return static_cast<DataBuffer *>(Buffer<uint8_t[]>::ref());
    }
 
    DataBuffer *ref_all(unsigned total_size, DataBuffer **tail = nullptr,
        unsigned *tail_size = nullptr)
    {
        DataBuffer *curr = this;
        do
        {
            HASSERT(curr);
            curr->ref();
            if (total_size > curr->size())
            {
                total_size -= curr->size();
            }
            else
            {
                if (tail)
                {
                    *tail = curr;
                }
                if (tail_size)
                {
                    *tail_size = total_size;
                }
                total_size = 0;
            }
            curr = curr->next();
        } while (total_size > 0);
        return this;
    }
 
    void unref_all(unsigned total_size)
    {
        DataBuffer *curr = this;
        while (true)
        {
            HASSERT(curr);
            if (total_size > curr->size())
            {
                DataBuffer *next = curr->next();
                total_size -= curr->size();
                curr->unref();
                curr = next;
            }
            else
            {
                curr->unref();
                break;
            }
        }
    }
 
    DataBuffer *get_read_pointer(
        unsigned skip, uint8_t **ptr, unsigned *available)
    {
        DataBuffer *curr = this;
        while (curr->size() <= skip)
        {
            skip -= curr->size();
            curr = curr->next();
            HASSERT(curr);
        }
        *ptr = curr->data() + skip;
        *available = curr->size() - skip;
        return curr->next();
    }
 
#ifdef DEBUG_DATA_BUFFER_FREE
    void unref()
    {
        if (references() == 1)
        {
            check_db_ownership(this);
        }
        Buffer::unref();
    }
#endif
 
private:
    friend class DataBufferPool;
 
    DataBuffer(DataBufferPool *p)
        : Buffer<uint8_t[]>((Pool *)p)
    {
    }
}; // class DataBuffer
 
using DataBufferPtr = std::unique_ptr<DataBuffer, BufferDelete<uint8_t[]>>;
 
class LinkedDataBufferPtr;
 
class LinkedDataBufferPtr
#ifdef DEBUG_DATA_BUFFER_FREE
    : public LinkedObject<LinkedDataBufferPtr>
#endif
{
public:
    LinkedDataBufferPtr()
    {
    }
 
    ~LinkedDataBufferPtr()
    {
        reset();
    }
 
    LinkedDataBufferPtr(LinkedDataBufferPtr &&o)
        : head_(o.head_)
        , tail_(o.tail_)
        , size_(o.size_)
        , skip_(o.skip_)
        , free_(o.free_)
    {
        o.clear();
    }
 
    void operator=(LinkedDataBufferPtr &&o)
    {
        reset();
        head_ = o.head_;
        tail_ = o.tail_;
        size_ = o.size_;
        skip_ = o.skip_;
        free_ = o.free_;
        o.clear();
    }
 
    LinkedDataBufferPtr(const LinkedDataBufferPtr &) = delete;
    void operator=(const LinkedDataBufferPtr &) = delete;
 
    void reset(const LinkedDataBufferPtr &o, ssize_t size = -1)
    {
        reset();
        if (size < 0)
        {
            size = o.size_;
        }
        if ((size_t)size > o.size_)
        {
            size = o.size_;
        }
        if (!size)
        {
            // Nothing to copy, this will be an empty buffer.
            return;
        }
        skip_ = o.skip_;
        size_ = size;
        // Takes references, keeping the tail and tail size.
        unsigned tail_size = 0;
        head_ = o.head_->ref_all(o.skip_ + size, &tail_, &tail_size);
        HASSERT(tail_size > 0);
        free_ = -tail_size;
    }
 
    void reset(DataBuffer *buf)
    {
        reset();
        head_ = tail_ = buf;
        skip_ = 0;
        free_ = buf->size();
        size_ = 0;
        buf->set_size(0);
    }
 
    void reset(DataBuffer *buf, unsigned skip, unsigned size)
    {
        reset();
        head_ = buf;
        skip_ = skip;
        size_ = size;
        free_ = -int(skip + size);
        tail_ = buf;
    }
 
    void append_empty_buffer(DataBuffer *buf)
    {
        if (!head_)
        {
            reset(buf);
            return;
        }
        HASSERT(free_ >= 0); // appendable
        HASSERT(tail_);
        // Note: if free_ was > 0, there were some unused bytes in the tail
        // buffer. However, as part of the append operation, we lose these
        // bytes as capacity. The new free part will be only in the newly
        // appended tail_ buffer. This is because free_ can never span more
        // than one buffer.
        free_ = buf->size();
        buf->set_size(0);
        HASSERT(!tail_->next());
        tail_->set_next(buf);
        tail_ = buf;
    }
 
    void reset()
    {
        if (head_)
        {
            auto *h = head_;
            size_t len = size_ + skip_;
            clear();
            h->unref_all(len);
            return;
        }
        else
        {
            clear();
        }
    }
 
    uint8_t *data_write_pointer()
    {
        if (!tail_)
        {
            return nullptr;
        }
        return tail_->data() + tail_->size();
    }
 
    void data_write_advance(size_t len)
    {
        HASSERT(free_ >= 0 && ((int)len <= free_));
        free_ -= len;
        tail_->set_size(tail_->size() + len);
        size_ += len;
    }
 
    const uint8_t *data_read_pointer(size_t *len)
    {
        if (!head_ || !size_)
        {
            *len = 0;
            return nullptr;
        }
        unsigned avail = 0;
        uint8_t *p;
        head_->get_read_pointer(skip_, &p, &avail);
        if (avail > size_)
        {
            avail = size_;
        }
        *len = avail;
        return p;
    }
 
    void data_read_advance(size_t len)
    {
        HASSERT(len <= size());
        skip_ += len;
        size_ -= len;
        while (head_ && skip_ >= head_->size())
        {
            if (head_ == tail_)
            {
                if (free() > 0)
                {
                    // We can still write into this buffer, do not unref it.
                    break;
                }
                else
                {
                    // We're ending up with an empty linkedbuffer.
                    auto *b = head_;
                    clear();
                    b->unref();
                    return;
                }
            }
            skip_ -= head_->size();
            auto *b = head_;
            auto *next_head = head_->next();
            head_ = next_head;
            if (!head_)
            {
                tail_ = nullptr;
            }
            b->unref();
        }
    }
 
    DataBuffer *head() const
    {
        return head_;
    }
 
    DataBuffer *tail() const
    {
        return tail_;
    }
    
    unsigned skip() const
    {
        return skip_;
    }
 
    unsigned size() const
    {
        return size_;
    }
 
    size_t free() const
    {
        if (free_ < 0)
        {
            return 0;
        }
        return free_;
    }
 
    LinkedDataBufferPtr transfer_head(size_t len)
    {
        LinkedDataBufferPtr ret;
        ret.head_ = head_;
        ret.tail_ = tail_;
        ret.skip_ = skip_;
        ret.size_ = len;
 
        HASSERT(tail_); // always true when we have a buffer
        HASSERT(len <= size_);
 
        size_t bytes_left = size_ - len;
 
        // tail_->size() is the previously used bytes in the tail buffer. The
        // number of bytes not transferred shall fit into this. There must be
        // however at least one byte in the tail buffer that *was* transferred.
        HASSERT(bytes_left < tail_->size());
 
        size_t bytes_transferred_from_tail_buffer = tail_->size() - bytes_left;
 
        // Since the tail is now in both the transferred chain as well as in
        // the current chain, it needs an extra ref. We keep that ref.
        head_ = tail_->ref();
        size_ = bytes_left;
        skip_ = bytes_transferred_from_tail_buffer;
        HASSERT(skip_ > 0);
        // Saves the end offset of the tail buffer in ret.
        ret.free_ = -bytes_transferred_from_tail_buffer;
        // this->free_ remains as is.
        return ret;
    }
 
    void append_to(std::string *recvd) const
    {
        DataBuffer *head = head_;
        unsigned skip = skip_;
        size_t len = size_;
        recvd->reserve(recvd->size() + len);
        while (len > 0)
        {
            uint8_t *ptr;
            unsigned available;
            head = head->get_read_pointer(skip, &ptr, &available);
            if (available > len)
            {
                available = len;
            }
            recvd->append((char *)ptr, available);
            len -= available;
            skip = 0;
        }
    }
 
    bool try_append_from(const LinkedDataBufferPtr &o, bool add_link = false)
    {
        if (!o.size())
        {
            return true; // zero bytes, nothing to do.
        }
        if (!size_)
        {
            // We are empty, so anything can be appended.
            reset(o);
            return true;
        }
        if (free_ >= 0)
        {
            // writeable buffer, cannot append.
            return false;
        }
        HASSERT(o.head());
        if (o.head() != tail_) // Buffer does not start in the same chain where
                               // we end.
        {
            HASSERT(tail_); // else we went into the !size_ branch above
 
            // Checks if the end of the tail buffer is already reached. This
            // means that we don't depend on the value of free_ anymore for
            // correctness. We also check that o starts at the beginning of the
            // head buffer.
            if (tail_->size() == (size_t)-free_ && o.skip() == 0)
            {
                if (tail_->next() == o.head())
                {
                    // link already exists
                }
                else if (add_link && tail_->next() == nullptr)
                {
                    tail_->set_next(o.head());
                }
                else
                {
                    return false;
                }
            }
            else
            {
                return false;
            }
        }
        else if (-free_ != (int)o.skip())
        {
            // Not back-to-back.
            return false;
        }
        // Now we're good, so take over the extra buffers.
        // Acquire extra references
        o.head_->ref_all(o.skip() + o.size());
        if (tail_ == o.head())
        {
            HASSERT(o.head_->references() > 1);
            // Release duplicate reference between the two chains.
            o.head_->unref();
        }
        tail_ = o.tail_;
        if (o.free_ < 0)
        {
            free_ = o.free_;
        }
        else
        {
            free_ = -tail_->size();
        }
        size_ += o.size_;
        return true;
    }
 
private:
    void clear()
    {
        head_ = tail_ = nullptr;
        skip_ = free_ = size_ = 0;
    }
 
    DataBuffer *head_ {nullptr};
    DataBuffer *tail_ {nullptr};
    size_t size_ {0};
    uint16_t skip_ {0};
    int16_t free_ {0};
};
 
#ifdef DEBUG_DATA_BUFFER_FREE
void check_db_ownership(DataBuffer *b)
{
    AtomicHolder h(LinkedDataBufferPtr::head_mu());
    for (LinkedDataBufferPtr *l = LinkedDataBufferPtr::link_head(); l;
         l = l->link_next())
    {
        ssize_t total = l->skip() + l->size();
        for (DataBuffer *curr = l->head(); total > 0;)
        {
            HASSERT(curr != b);
            total -= curr->size();
            curr = curr->next();
        }
    }
}
#endif
 
class DataBufferPool : public Pool
{
public:
    DataBufferPool(unsigned payload_size)
        : payloadSize_(payload_size)
    {
        HASSERT(payload_size <= 65535u - sizeof(BufferBase));
    }
 
#ifdef GTEST
    uint16_t *payload_size_override()
    {
        return &payloadSize_;
    }
#endif
 
    size_t free_items() override
    {
        return base_pool()->free_items(alloc_size());
    }
 
    size_t free_items(size_t size) override
    {
        return base_pool()->free_items(size);
    }
 
    void alloc(DataBuffer **result)
    {
#ifdef DEBUG_BUFFER_MEMORY
        g_current_alloc = &&alloc;
    alloc:
#endif
        *result = static_cast<DataBuffer *>(
            base_pool()->alloc_untyped(alloc_size(), nullptr));
        if (*result)
        {
            new (*result) DataBuffer(this);
            (*result)->size_ = payload_size();
        }
    }
 
    void alloc(DataBufferPtr *result)
    {
        DataBuffer *b;
        alloc(&b);
        result->reset(b);
    }
 
private:
    BufferBase *alloc_untyped(size_t size, Executable *flow) override
    {
        DIE("DataBufferPool does not support this type of allocation.");
    }
 
    void free(BufferBase *item) override
    {
        // Restores the correct size for assigning it to the right freelist
        // bucket.
        item->size_ = alloc_size();
        // Clears the next pointer as we are not using these for queues.
        item->next = nullptr;
        base_pool()->free(item);
    }
 
    Pool *base_pool()
    {
        return mainBufferPool;
    }
 
    uint16_t alloc_size()
    {
        return sizeof(BufferBase) + payloadSize_;
    }
 
    uint16_t payload_size()
    {
        return payloadSize_;
    }
 
    uint16_t payloadSize_;
};
 
#endif // _UTILS_DATABUFFER_HXX_