iobuf support reserve_aligned

Author: yanglimingcn

src/butil/iobuf.cpp

@@ -1371,6 +1371,50 @@ 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_block();
+        if (BAIDU_UNLIKELY(!b)) {
+            return INVALID_AREA;
+        }
+        const size_t remainder =
+            reinterpret_cast<uintptr_t>(b->data + b->size) % alignment;
+        const size_t align_add = (remainder != 0) ? alignment - remainder : 0;
+        const size_t align_left =
+            (b->left_space() > align_add) ? b->left_space() - align_add : 0;
+        if (b->size == 0 && align_left < alignment) {
+            LOG(ERROR) << "iobuf block size " << b->cap
+                       << " is not suitable for alignment " << alignment;
+            return INVALID_AREA;
+        }
+        if (align_left < alignment) {
+            b->size = b->cap;
+            continue;
+        }
+        b->size += align_add;
+        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,79 @@ 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);
+    }
+    {
+        LOG(INFO) << "0000";
+        butil::IOBuf buf;
+        auto area = buf.reserve_aligned(4096, 4095);
+        ASSERT_EQ(area, butil::IOBuf::INVALID_AREA);
+    }
+    {
+        LOG(INFO) << "1111";
+        butil::IOBuf buf;
+        auto area = buf.reserve_aligned(4096 * 10, 8192);
+        ASSERT_EQ(area, butil::IOBuf::INVALID_AREA);
+    }
+}
+
 } // namespace