diff --git a/src/butil/iobuf.cpp b/src/butil/iobuf.cpp index 382bcfb8ff..56dd6b235b 100644 --- a/src/butil/iobuf.cpp +++ b/src/butil/iobuf.cpp @@ -358,7 +358,7 @@ void remove_tls_block_chain() { // Get a (non-full) block from TLS. // Notice that the block is not removed from TLS. -inline IOBuf::Block* share_tls_block() { +IOBuf::Block* share_tls_block() { TLSData& tls_data = g_tls_data; IOBuf::Block* const b = tls_data.block_head; if (b != NULL && !b->full()) { @@ -366,9 +366,13 @@ inline IOBuf::Block* share_tls_block() { } IOBuf::Block* new_block = NULL; if (b) { - new_block = b->portal_next; - b->dec_ref(); - --tls_data.num_blocks; + new_block = b; + while (new_block && new_block->full()) { + IOBuf::Block* const saved_next = new_block->portal_next; + new_block->dec_ref(); + --tls_data.num_blocks; + new_block = saved_next; + } } else if (!tls_data.registered) { tls_data.registered = true; // Only register atexit at the first time diff --git a/test/iobuf_unittest.cpp b/test/iobuf_unittest.cpp index 9322c54699..cec4019fa4 100644 --- a/test/iobuf_unittest.cpp +++ b/test/iobuf_unittest.cpp @@ -47,6 +47,7 @@ extern IOBuf::Block* get_tls_block_head(); extern int get_tls_block_count(); extern void remove_tls_block_chain(); extern IOBuf::Block* acquire_tls_block(); +extern IOBuf::Block* share_tls_block(); extern void release_tls_block_chain(IOBuf::Block* b); extern uint32_t block_cap(IOBuf::Block const* b); extern uint32_t block_size(IOBuf::Block const* b); @@ -1656,6 +1657,31 @@ TEST_F(IOBufTest, append_user_data_and_share) { ASSERT_EQ(data, my_free_params); } +TEST_F(IOBufTest, share_tls_block) { + butil::iobuf::remove_tls_block_chain(); + butil::IOBuf::Block* b = butil::iobuf::acquire_tls_block(); + ASSERT_EQ(0, butil::iobuf::block_size(b)); + + butil::IOBuf::Block* b2 = butil::iobuf::share_tls_block(); + butil::IOBuf buf; + for (size_t i = 0; i < butil::iobuf::block_cap(b2); i++) { + buf.push_back('x'); + } + // after pushing to b2, b2 is full but it is still head of tls block. + ASSERT_NE(b, b2); + butil::iobuf::release_tls_block_chain(b); + ASSERT_EQ(b, butil::iobuf::share_tls_block()); + // After releasing b, now tls block is b(not full) -> b2(full) -> NULL + for (size_t i = 0; i < butil::iobuf::block_cap(b); i++) { + buf.push_back('x'); + } + // now tls block is b(full) -> b2(full) -> NULL + butil::IOBuf::Block* head_block = butil::iobuf::share_tls_block(); + ASSERT_EQ(0, butil::iobuf::block_size(head_block)); + ASSERT_NE(b, head_block); + ASSERT_NE(b2, head_block); +} + TEST_F(IOBufTest, acquire_tls_block) { butil::iobuf::remove_tls_block_chain(); butil::IOBuf::Block* b = butil::iobuf::acquire_tls_block();