add size in bthread_keytable_pool_t

src/bthread/key.cpp

@@ -275,9 +275,9 @@ class BAIDU_CACHELINE_ALIGNMENT KeyTableList {
         return temp;
     }
 
-    void move_first_n_to_target(KeyTable** target, uint32_t size) {
+    int move_first_n_to_target(KeyTable** target, uint32_t size) {
         if (size > _length || _head == NULL) {
-            return;
+            return 0;
         }
 
         KeyTable* current = _head;
@@ -301,6 +301,7 @@ class BAIDU_CACHELINE_ALIGNMENT KeyTableList {
                 _tail = NULL;
             }
         }
+        return count;
     }
 
     inline uint32_t get_length() {
@@ -343,8 +344,10 @@ KeyTable* borrow_keytable(bthread_keytable_pool_t* pool) {
             if (list) {
                 for (uint32_t i = 0; i < FLAGS_borrow_from_globle_size; ++i) {
                     if (p) {
-                        list->get()->append(p);
                         pool->free_keytables = p->next;
+                        list->get()->append(p);
+                        p = (KeyTable*)pool->free_keytables;
+                        --pool->size;
                     } else {
                         break;
                     }
@@ -387,8 +390,10 @@ void return_keytable(bthread_keytable_pool_t* pool, KeyTable* kt) {
         pthread_rwlock_unlock(&pool->rwlock);
         pthread_rwlock_wrlock(&pool->rwlock);
         if (!pool->destroyed) {
-            list->get()->move_first_n_to_target((KeyTable**)(&pool->free_keytables),
-                                                FLAGS_key_table_list_size / 2);
+            int out = list->get()->move_first_n_to_target(
+                (KeyTable**)(&pool->free_keytables),
+                FLAGS_key_table_list_size / 2);
+            pool->size += out;
         }
     }
     pthread_rwlock_unlock(&pool->rwlock);
@@ -440,6 +445,7 @@ int bthread_keytable_pool_init(bthread_keytable_pool_t* pool) {
     pthread_rwlock_init(&pool->rwlock, NULL);
     pool->list = new butil::ThreadLocal<bthread::KeyTableList>();
     pool->free_keytables = NULL;
+    pool->size = 0;
     pool->destroyed = 0;
     return 0;
 }
@@ -452,6 +458,7 @@ int bthread_keytable_pool_destroy(bthread_keytable_pool_t* pool) {
     bthread::KeyTable* saved_free_keytables = NULL;
     pthread_rwlock_wrlock(&pool->rwlock);
     pool->destroyed = 1;
+    pool->size = 0;
     delete (butil::ThreadLocal<bthread::KeyTableList>*)pool->list;
     saved_free_keytables = (bthread::KeyTable*)pool->free_keytables;
     pool->list = NULL;
@@ -489,10 +496,7 @@ int bthread_keytable_pool_getstat(bthread_keytable_pool_t* pool,
         return EINVAL;
     }
     pthread_rwlock_wrlock(&pool->rwlock);
-    size_t count = 0;
-    bthread::KeyTable* p = (bthread::KeyTable*)pool->free_keytables;
-    for (; p; p = p->next, ++count) {}
-    stat->nfree = count;
+    stat->nfree = pool->size;
     pthread_rwlock_unlock(&pool->rwlock);
     return 0;
 }
@@ -554,6 +558,7 @@ void bthread_keytable_pool_reserve(bthread_keytable_pool_t* pool,
         }
         kt->next = (bthread::KeyTable*)pool->free_keytables;
         pool->free_keytables = kt;
+        ++pool->size;
         pthread_rwlock_unlock(&pool->rwlock);
         if (data == NULL) {
             break;

src/bthread/types.h

@@ -88,6 +88,7 @@ typedef struct {
     pthread_rwlock_t rwlock;
     void* list;
     void* free_keytables;
+    size_t size;
     int destroyed;
 } bthread_keytable_pool_t;
 

test/bthread_key_unittest.cpp

@@ -441,7 +441,7 @@ static void usleep_thread_impl(PoolData2* data) {
         PoolData2* data_new = new PoolData2();
         ASSERT_EQ(0, bthread_setspecific(data->key, data_new));
     }
-    bthread_usleep(1000L * 1000L);
+    bthread_usleep(1000L);
     int length = get_thread_local_keytable_list_length(&test_pool);
     ASSERT_LE((size_t)length, bthread::FLAGS_key_table_list_size);
 }
@@ -487,7 +487,9 @@ TEST(KeyTest, frequently_borrow_keytable_when_using_pool) {
     bthread_t bth;
     ASSERT_EQ(0, bthread_start_urgent(&bth, &data.attr, run_launch_many_bthreads, &data));
     ASSERT_EQ(0, bthread_join(bth, NULL));
-
+    std::cout << "Free keytable size is "
+              << bthread_keytable_pool_size(&test_pool)
+              << " use keytable size is 25000" << std::endl;
     ASSERT_EQ(0, bthread_keytable_pool_destroy(&test_pool));
     ASSERT_EQ(0, bthread_key_delete(data.key));
 }
@@ -539,10 +541,12 @@ TEST(KeyTest, borrow_and_return_keytable_when_using_pool) {
     bthread_t borrow_bth[8];
     bthread_t return_bth[8];
     for (size_t i = 0; i < arraysize(borrow_bth); ++i) {
-        ASSERT_EQ(0, bthread_start_background(&borrow_bth[i], &attr, borrow_thread, NULL));
+        ASSERT_EQ(0, bthread_start_background(&borrow_bth[i], &attr,
+                                              borrow_thread, NULL));
     }
     for (size_t i = 0; i < arraysize(return_bth); ++i) {
-        ASSERT_EQ(0, bthread_start_background(&return_bth[i], &attr, return_thread, NULL));
+        ASSERT_EQ(0, bthread_start_background(&return_bth[i], &attr,
+                                              return_thread, NULL));
     }
 
     for (size_t i = 0; i < arraysize(borrow_bth); ++i) {
@@ -610,6 +614,7 @@ TEST(KeyTest, use_bthread_mutex_in_dtor) {
               *(lid_seqs.end() - 1) - *lid_seqs.begin());
 
     ASSERT_EQ(0, bthread_key_delete(key));
+    ASSERT_EQ(0, bthread_mutex_destroy(&mu));
 }
 
 }  // namespace