iobuf support reserve_aligned

Author: yanglimingcn

src/butil/iobuf.cpp

@@ -438,6 +438,51 @@ IOBuf::Block* share_tls_block() {
     return new_block;
 }
 
+// Get a memory aligned block from TLS. if TLS is not found, create one.
+IOBuf::Block* share_tls_aligned_block(size_t alignment) {
+    // find one block if align satisfied, then aligned the block offset
+    auto find_and_align_block = [](IOBuf::Block* const b, size_t alignment) {
+        const size_t remainder =
+            reinterpret_cast<uintptr_t>(b->data + b->size) % alignment;
+        size_t adder = 0;
+        if (remainder != 0) {
+            adder = alignment - remainder;
+        }
+        size_t left_space = 0;
+        if (b->left_space() > adder) {
+            left_space = b->left_space() - adder;
+        }
+        if (left_space < alignment) {
+            return false;
+        }
+        b->size += adder;  // align block offset
+        return true;
+    };
+    TLSData& tls_data = g_tls_data;
+    IOBuf::Block* b = tls_data.block_head;
+    while (b && b->full()) {
+        IOBuf::Block* const saved_next = b->u.portal_next;
+        b->dec_ref();
+        tls_data.block_head = saved_next;
+        --tls_data.num_blocks;
+        b = saved_next;
+    }
+    while (b) {
+        IOBuf::Block* const saved_next = b->u.portal_next;
+        if (find_and_align_block(b, alignment)) {
+            return b;
+        }
+        b = saved_next;
+    }
+    if (!b) {
+        b = create_block(alignment * 2 + sizeof(IOBuf::Block));  // may be NULL
+        if (b) {
+            find_and_align_block(b, alignment);  // always return true
+        }
+    }
+    return b;
+}
+
 // Return one block to TLS.
 inline void release_tls_block(IOBuf::Block* b) {
     if (!b) {
@@ -1371,6 +1416,35 @@ IOBuf::Area IOBuf::reserve(size_t count) {
     return result;
 }
 
+IOBuf::Area IOBuf::reserve_aligned(size_t n, size_t alignment) {
+    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;
+    }
+    size_t count = (n + alignment - 1) & ~(alignment - 1);
+    IOBuf::Area result = INVALID_AREA;
+    size_t total_nc = 0;
+    while (total_nc < count) {  // excluded count == 0
+        IOBuf::Block* b = iobuf::share_tls_aligned_block(alignment);
+        if (BAIDU_UNLIKELY(!b)) {
+            return INVALID_AREA;
+        }
+        const size_t nc =
+            std::min(count - total_nc, b->left_space()) & ~(alignment - 1);
+        const IOBuf::BlockRef r = {(uint32_t)b->size, (uint32_t)nc, b};
+        _push_back_ref(r);
+        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);
+        }
+        total_nc += nc;
+        b->size += nc;
+    }
+    return result;
+}
+
 int IOBuf::unsafe_assign(Area area, const void* data) {
     if (area == INVALID_AREA || data == NULL) {
         LOG(ERROR) << "Invalid parameters";

src/butil/iobuf.h

@@ -273,6 +273,11 @@ friend class IOBufCutter;
     // NOTE: reserve(0) returns INVALID_AREA.
     Area reserve(size_t n);
 
+    // Reserve `n' aligned uninitialized bytes at back-side.
+    // Returns an object representing the reserved area, INVALID_AREA on failure.
+    // NOTE: reserve(0) returns INVALID_AREA, alignment must be a power of 2.
+    Area reserve_aligned(size_t n, size_t alignment);
+
     // [EXTREMELY UNSAFE]
     // Copy `data' to the reserved `area'. `data' must be as long as the
     // reserved size.

test/iobuf_unittest.cpp

@@ -1785,4 +1785,87 @@ TEST_F(IOBufTest, acquire_tls_block) {
     ASSERT_NE(butil::iobuf::block_cap(b), butil::iobuf::block_size(b));
 }
 
+TEST_F(IOBufTest, reserve_aligned) {
+    {
+        butil::IOBuf buf;
+        auto area = buf.reserve_aligned(1024, 16);
+        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::IOBuf buf;
+        auto area = buf.reserve_aligned(1024, 4096);
+        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::IOBuf buf;
+        auto area = buf.reserve_aligned(8191, 4096);
+        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::IOBuf buf;
+        auto area = buf.reserve_aligned(4096 * 10 - 1, 4096);
+        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::IOBuf buf;
+        auto area = buf.reserve_aligned(4096, 4095);
+        ASSERT_EQ(area, butil::IOBuf::INVALID_AREA);
+    }
+    {
+        butil::IOBuf buf;
+        auto area = buf.reserve_aligned(4096 * 10 + 1, 8192);
+        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);
+    }
+}
+
 } // namespace