iobuf support reserve_aligned

src/butil/iobuf.cpp

@@ -355,6 +355,23 @@ inline IOBuf::Block* create_block() {
     return create_block(IOBuf::DEFAULT_BLOCK_SIZE);
 }
 
+inline IOBuf::Block* create_block_aligned(size_t block_size, size_t alignment) {
+    if (block_size > 0xFFFFFFFFULL) {
+        LOG(FATAL) << "block_size=" << block_size << " is too large";
+        return NULL;
+    }
+    char* mem = (char*)iobuf::blockmem_allocate(block_size);
+    if (mem == NULL) {
+        return NULL;
+    }
+    char* data = mem + sizeof(IOBuf::Block);
+    // change data pointer & data size make align satisfied
+    size_t adder = (-reinterpret_cast<uintptr_t>(data)) & (alignment - 1);
+    size_t size =
+        (block_size - sizeof(IOBuf::Block) - adder) & ~(alignment - 1);
+    return new (mem) IOBuf::Block(data + adder, size);
+}
+
 // === Share TLS blocks between appending operations ===
 // Max number of blocks in each TLS. This is a soft limit namely
 // release_tls_block_chain() may exceed this limit sometimes.
@@ -1785,6 +1802,45 @@ void IOPortal::return_cached_blocks_impl(Block* b) {
     iobuf::release_tls_block_chain(b);
 }
 
+IOBuf::Area IOReserveAlignedBuf::reserve(size_t count) {
+    IOBuf::Area result = INVALID_AREA;
+    if (_reserved == true) {
+        LOG(ERROR) << "Already call reserved";
+        return result;
+    }
+    _reserved = true;
+    bool is_power_two = _alignment > 0 && (_alignment & (_alignment - 1));
+    if (is_power_two != 0) {
+        LOG(ERROR) << "Invalid alignment, must power of two";
+        return INVALID_AREA;
+    }
+    count = (count + _alignment - 1) & ~(_alignment - 1);
+    size_t total_nc = 0;
+    while (total_nc < count) {
+        const auto block_size =
+            std::max(_alignment, 4096UL) * 2 + sizeof(IOBuf::Block);
+        auto b = iobuf::create_block_aligned(block_size, _alignment);
+        if (BAIDU_UNLIKELY(!b)) {
+            LOG(ERROR) << "Create block failed";
+            return result;
+        }
+        const size_t nc = std::min(count - total_nc, b->left_space());
+        const IOBuf::BlockRef r = {(uint32_t)b->size, (uint32_t)nc, b};
+        _push_back_ref(r);
+        // aligned block is not from tls, release block ref
+        b->dec_ref();
+        if (total_nc == 0) {
+            // Encode the area at first time. Notice that the pushed ref may
+            // be merged with existing ones.
+            result = make_area(_ref_num() - 1, _back_ref().length - nc, count);
+        }
+        // add total nc
+        total_nc += nc;
+        b->size += nc;
+    };
+    return result;
+}
+
 //////////////// IOBufCutter ////////////////
 
 IOBufCutter::IOBufCutter(butil::IOBuf* buf)

src/butil/iobuf.h

@@ -489,6 +489,17 @@ class IOPortal : public IOBuf {
     Block* _block;
 };
 
+class IOReserveAlignedBuf : public IOBuf {
+public:
+    IOReserveAlignedBuf(size_t alignment)
+        : _alignment(alignment), _reserved(false) {}
+    Area reserve(size_t count);
+
+private:
+    size_t _alignment;
+    bool _reserved;
+};
+
 // Specialized utility to cut from IOBuf faster than using corresponding
 // methods in IOBuf.
 // Designed for efficiently parsing data from IOBuf.

test/iobuf_unittest.cpp

@@ -22,6 +22,7 @@
 #include <fcntl.h>                     // O_RDONLY
 #include <stdlib.h>
 #include <memory>
+#include <cstring>
 #include <butil/files/temp_file.h>      // TempFile
 #include <butil/containers/flat_map.h>
 #include <butil/macros.h>
@@ -1785,4 +1786,108 @@ TEST_F(IOBufTest, acquire_tls_block) {
     ASSERT_NE(butil::iobuf::block_cap(b), butil::iobuf::block_size(b));
 }
 
+TEST_F(IOBufTest, reserve_aligned) {
+    {
+        butil::IOReserveAlignedBuf buf(16);
+        auto area = buf.reserve(1024);
+        ASSERT_NE(area, butil::IOBuf::INVALID_AREA);
+        butil::IOBufAsZeroCopyInputStream wrapper(buf);
+        const void* data;
+        int size;
+        int total_size = 0;
+        while (wrapper.Next(&data, &size)) {
+            ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 16, 0);
+            ASSERT_EQ(size % 16, 0);
+            total_size += size;
+        }
+        ASSERT_EQ(total_size, 1024);
+    }
+    {
+        butil::IOReserveAlignedBuf buf(4096);
+        auto area = buf.reserve(1024);
+        ASSERT_NE(area, butil::IOBuf::INVALID_AREA);
+        butil::IOBufAsZeroCopyInputStream wrapper(buf);
+        const void* data;
+        int size;
+        int total_size = 0;
+        while (wrapper.Next(&data, &size)) {
+            ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
+            ASSERT_EQ(size % 4096, 0);
+            total_size += size;
+        }
+        ASSERT_EQ(total_size, 4096);
+    }
+    {
+        butil::IOReserveAlignedBuf buf(4096);
+        auto area = buf.reserve(8191);
+        ASSERT_NE(area, butil::IOBuf::INVALID_AREA);
+        butil::IOBufAsZeroCopyInputStream wrapper(buf);
+        const void* data;
+        int size;
+        int total_size = 0;
+        while (wrapper.Next(&data, &size)) {
+            ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
+            ASSERT_EQ(size % 4096, 0);
+            total_size += size;
+        }
+        ASSERT_EQ(total_size, 8192);
+    }
+    {
+        butil::IOReserveAlignedBuf buf(4096);
+        auto area = buf.reserve(4096 * 10 - 1);
+        ASSERT_NE(area, butil::IOBuf::INVALID_AREA);
+        butil::IOBufAsZeroCopyInputStream wrapper(buf);
+        const void* data;
+        int size;
+        int total_size = 0;
+        while (wrapper.Next(&data, &size)) {
+            ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
+            ASSERT_EQ(size % 4096, 0);
+            total_size += size;
+        }
+        ASSERT_EQ(total_size, 4096 * 10);
+    }
+    {
+        butil::IOReserveAlignedBuf buf(4095);
+        auto area = buf.reserve(4096);
+        ASSERT_EQ(area, butil::IOBuf::INVALID_AREA);
+    }
+    {
+        butil::IOReserveAlignedBuf buf(8192);
+        auto area = buf.reserve(4096 * 10 + 1);
+        ASSERT_NE(area, butil::IOBuf::INVALID_AREA);
+        butil::IOBufAsZeroCopyInputStream wrapper(buf);
+        const void* data;
+        int size;
+        int total_size = 0;
+        while (wrapper.Next(&data, &size)) {
+            ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
+            ASSERT_EQ(size % 4096, 0);
+            total_size += size;
+        }
+        ASSERT_EQ(total_size, 4096 * 10 + 8192);
+    }
+    {
+        butil::IOReserveAlignedBuf buf(4096);
+        auto area = buf.reserve(1024 * 1024 * 3);
+        ASSERT_NE(area, butil::IOBuf::INVALID_AREA);
+        butil::IOBufAsZeroCopyInputStream wrapper(buf);
+        const void* data;
+        int size;
+        int count = 0;
+        int total_size = 0;
+        std::stringstream ss;
+        while (wrapper.Next(&data, &size)) {
+            ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
+            ASSERT_EQ(size % 4096, 0);
+            std::string str(size, 'A' + count++);
+            ss << str;
+            std::memcpy(const_cast<void*>(data), str.data(), str.size());
+            total_size += size;
+        }
+        ASSERT_EQ(total_size, 3145728);
+        ASSERT_EQ(ss.str(), buf.to_string());
+    }
+}
+
 } // namespace