Fix acquire tls block

src/butil/iobuf.cpp

@@ -282,10 +282,14 @@ IOBuf::Block* get_portal_next(IOBuf::Block const* b) {
     return b->portal_next;
 }
 
-uint32_t block_cap(IOBuf::Block const *b) {
+uint32_t block_cap(IOBuf::Block const* b) {
     return b->cap;
 }
 
+uint32_t block_size(IOBuf::Block const* b) {
+    return b->size;
+}
+
 inline IOBuf::Block* create_block(const size_t block_size) {
     if (block_size > 0xFFFFFFFFULL) {
         LOG(FATAL) << "block_size=" << block_size << " is too large";
@@ -417,7 +421,7 @@ inline void release_tls_block(IOBuf::Block *b) {
 
 // Return chained blocks to TLS.
 // NOTE: b MUST be non-NULL and all blocks linked SHOULD not be full.
-inline void release_tls_block_chain(IOBuf::Block* b) {
+void release_tls_block_chain(IOBuf::Block* b) {
     TLSData& tls_data = g_tls_data;
     size_t n = 0;
     if (tls_data.num_blocks >= MAX_BLOCKS_PER_THREAD) {
@@ -451,13 +455,13 @@ inline void release_tls_block_chain(IOBuf::Block* b) {
 }
 
 // Get and remove one (non-full) block from TLS. If TLS is empty, create one.
-inline IOBuf::Block* acquire_tls_block() {
+IOBuf::Block* acquire_tls_block() {
     TLSData& tls_data = g_tls_data;
     IOBuf::Block* b = tls_data.block_head;
     if (!b) {
         return create_block();
     }
-    if (b->full()) {
+    while (b->full()) {
         IOBuf::Block* const saved_next = b->portal_next;
         b->dec_ref();
         tls_data.block_head = saved_next;

test/iobuf_unittest.cpp

@@ -33,7 +33,11 @@ extern uint32_t block_cap(butil::IOBuf::Block const* b);
 extern IOBuf::Block* get_tls_block_head();
 extern int get_tls_block_count();
 extern void remove_tls_block_chain();
-IOBuf::Block* get_portal_next(IOBuf::Block const* b);
+extern IOBuf::Block* acquire_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);
+extern IOBuf::Block* get_portal_next(IOBuf::Block const* b);
 }
 }
 
@@ -1639,4 +1643,28 @@ TEST_F(IOBufTest, append_user_data_and_share) {
     ASSERT_EQ(data, my_free_params);
 }
 
+TEST_F(IOBufTest, acquire_tls_block) {
+    butil::iobuf::remove_tls_block_chain();
+    butil::IOBuf::Block* b = butil::iobuf::acquire_tls_block();
+    const size_t block_cap = butil::iobuf::block_cap(b);
+    butil::IOBuf buf;
+    for (size_t i = 0; i < block_cap; i++) {
+        buf.append("x");
+    }
+    ASSERT_EQ(1, butil::iobuf::get_tls_block_count());
+    butil::IOBuf::Block* head = butil::iobuf::get_tls_block_head();
+    ASSERT_EQ(butil::iobuf::block_cap(head), butil::iobuf::block_size(head));
+    butil::iobuf::release_tls_block_chain(b);
+    ASSERT_EQ(2, butil::iobuf::get_tls_block_count());
+    for (size_t i = 0; i < block_cap; i++) {
+        buf.append("x");
+    }
+    ASSERT_EQ(2, butil::iobuf::get_tls_block_count());
+    head = butil::iobuf::get_tls_block_head();
+    ASSERT_EQ(butil::iobuf::block_cap(head), butil::iobuf::block_size(head));
+    b = butil::iobuf::acquire_tls_block();
+    ASSERT_EQ(0, butil::iobuf::get_tls_block_count());
+    ASSERT_NE(butil::iobuf::block_cap(b), butil::iobuf::block_size(b));
+}
+
 } // namespace