Skip to content

Commit ea60451

Browse files
committed
iobuf support reserve_aligned
1 parent 0625945 commit ea60451

3 files changed

Lines changed: 99 additions & 0 deletions

File tree

src/butil/iobuf.cpp

Lines changed: 35 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1371,6 +1371,41 @@ IOBuf::Area IOBuf::reserve(size_t count) {
13711371
return result;
13721372
}
13731373

1374+
IOBuf::Area IOBuf::reserve_aligned(size_t n, size_t alignment) {
1375+
size_t count = (n + alignment - 1) & ~(alignment - 1);
1376+
IOBuf::Area result = INVALID_AREA;
1377+
size_t total_nc = 0;
1378+
while (total_nc < count) { // excluded count == 0
1379+
IOBuf::Block* b = iobuf::share_tls_block();
1380+
CHECK(alignment <= b->cap)
1381+
<< "alignment must not greater than block capacity";
1382+
if (BAIDU_UNLIKELY(!b)) {
1383+
return INVALID_AREA;
1384+
}
1385+
const size_t remainder =
1386+
reinterpret_cast<uintptr_t>(b->data + b->size) % alignment;
1387+
const size_t align_offset =
1388+
(remainder != 0) ? alignment - remainder : 0;
1389+
if (b->left_space() - align_offset < alignment) {
1390+
b->size = b->cap;
1391+
continue;
1392+
}
1393+
b->size += align_offset;
1394+
const size_t nc =
1395+
std::min(count - total_nc, b->left_space()) & ~(alignment - 1);
1396+
const IOBuf::BlockRef r = {(uint32_t)b->size, (uint32_t)nc, b};
1397+
_push_back_ref(r);
1398+
if (total_nc == 0) {
1399+
// Encode the area at first time. Notice that the pushed ref may
1400+
// be merged with existing ones.
1401+
result = make_area(_ref_num() - 1, _back_ref().length - nc, count);
1402+
}
1403+
total_nc += nc;
1404+
b->size += nc;
1405+
}
1406+
return result;
1407+
}
1408+
13741409
int IOBuf::unsafe_assign(Area area, const void* data) {
13751410
if (area == INVALID_AREA || data == NULL) {
13761411
LOG(ERROR) << "Invalid parameters";

src/butil/iobuf.h

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -273,6 +273,11 @@ friend class IOBufCutter;
273273
// NOTE: reserve(0) returns INVALID_AREA.
274274
Area reserve(size_t n);
275275

276+
// Reserve `n' aligned uninitialized bytes at back-side.
277+
// Returns an object representing the reserved area, INVALID_AREA on failure.
278+
// NOTE: reserve(0) returns INVALID_AREA, alignment must be a power of 2.
279+
Area reserve_aligned(size_t n, size_t alignment);
280+
276281
// [EXTREMELY UNSAFE]
277282
// Copy `data' to the reserved `area'. `data' must be as long as the
278283
// reserved size.

test/iobuf_unittest.cpp

Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1785,4 +1785,63 @@ TEST_F(IOBufTest, acquire_tls_block) {
17851785
ASSERT_NE(butil::iobuf::block_cap(b), butil::iobuf::block_size(b));
17861786
}
17871787

1788+
TEST_F(IOBufTest, reserve_aligned) {
1789+
{
1790+
butil::IOBuf buf;
1791+
auto area = buf.reserve_aligned(1024, 16);
1792+
butil::IOBufAsZeroCopyInputStream wrapper(buf);
1793+
const void* data;
1794+
int size;
1795+
int total_size = 0;
1796+
while (wrapper.Next(&data, &size)) {
1797+
ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 16, 0);
1798+
ASSERT_EQ(size % 16, 0);
1799+
total_size += size;
1800+
}
1801+
ASSERT_EQ(total_size, 1024);
1802+
}
1803+
{
1804+
butil::IOBuf buf;
1805+
auto area = buf.reserve_aligned(1024, 4096);
1806+
butil::IOBufAsZeroCopyInputStream wrapper(buf);
1807+
const void* data;
1808+
int size;
1809+
int total_size = 0;
1810+
while (wrapper.Next(&data, &size)) {
1811+
ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
1812+
ASSERT_EQ(size % 4096, 0);
1813+
total_size += size;
1814+
}
1815+
ASSERT_EQ(total_size, 4096);
1816+
}
1817+
{
1818+
butil::IOBuf buf;
1819+
auto area = buf.reserve_aligned(8191, 4096);
1820+
butil::IOBufAsZeroCopyInputStream wrapper(buf);
1821+
const void* data;
1822+
int size;
1823+
int total_size = 0;
1824+
while (wrapper.Next(&data, &size)) {
1825+
ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
1826+
ASSERT_EQ(size % 4096, 0);
1827+
total_size += size;
1828+
}
1829+
ASSERT_EQ(total_size, 8192);
1830+
}
1831+
{
1832+
butil::IOBuf buf;
1833+
auto area = buf.reserve_aligned(4096 * 10 - 1, 4096);
1834+
butil::IOBufAsZeroCopyInputStream wrapper(buf);
1835+
const void* data;
1836+
int size;
1837+
int total_size = 0;
1838+
while (wrapper.Next(&data, &size)) {
1839+
ASSERT_EQ(reinterpret_cast<uintptr_t>(data) % 4096, 0);
1840+
ASSERT_EQ(size % 4096, 0);
1841+
total_size += size;
1842+
}
1843+
ASSERT_EQ(total_size, 4096 * 10);
1844+
}
1845+
}
1846+
17881847
} // namespace

0 commit comments

Comments
 (0)