iobuf support reserve_aligned

Author: yanglimingcn

src/butil/iobuf.cpp

@@ -1371,6 +1371,41 @@ IOBuf::Area IOBuf::reserve(size_t count) {
     return result;
 }
 
+IOBuf::Area IOBuf::reserve_aligned(size_t n, size_t alignment) {
+    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_block();
+        CHECK(alignment <= b->cap)
+            << "alignment must not greater than block capacity";
+        if (BAIDU_UNLIKELY(!b)) {
+            return INVALID_AREA;
+        }
+        const size_t remainder =
+            reinterpret_cast<uintptr_t>(b->data + b->size) % alignment;
+        const size_t align_offset =
+            (remainder != 0) ? alignment - remainder : 0;
+        if (b->left_space() - align_offset < alignment) {
+            b->size = b->cap;
+            continue;
+        }
+        b->size += align_offset;
+        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,63 @@ 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);
+        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);
+        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);
+        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);
+        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);
+    }
+}
+
 } // namespace