rdma rumtime extend block restriction

Author: yanglimingcn

src/brpc/rdma/block_pool.cpp

@@ -36,7 +36,7 @@ DEFINE_int32(rdma_memory_pool_initial_size_mb, 1024,
              "Initial size of memory pool for RDMA (MB)");
 DEFINE_int32(rdma_memory_pool_increase_size_mb, 1024,
              "Increased size of memory pool for RDMA (MB)");
-DEFINE_int32(rdma_memory_pool_max_regions, 1, "Max number of regions");
+DEFINE_int32(rdma_memory_pool_max_regions, 3, "Max number of regions");
 DEFINE_int32(rdma_memory_pool_buckets, 4, "Number of buckets to reduce race");
 DEFINE_int32(rdma_memory_pool_tls_cache_num, 128, "Number of cached block in tls");
 DEFINE_bool(rdma_memory_pool_user_specified_memory, false,
@@ -96,6 +96,7 @@ struct GlobalInfo {
     std::vector<IdleNode*> idle_list[BLOCK_SIZE_COUNT];
     std::vector<butil::Mutex*> lock[BLOCK_SIZE_COUNT];
     std::vector<size_t> idle_size[BLOCK_SIZE_COUNT];
+    int region_num[BLOCK_SIZE_COUNT];
     butil::Mutex extend_lock;
 };
 static GlobalInfo* g_info = NULL;
@@ -132,6 +133,7 @@ static void* ExtendBlockPoolImpl(void* region_base, size_t region_size,
                                  int block_type) {
     if (g_region_num == FLAGS_rdma_memory_pool_max_regions) {
         LOG(INFO) << "Memory pool reaches max regions";
+        free(region_base);
         errno = ENOMEM;
         return NULL;
     }
@@ -167,6 +169,7 @@ static void* ExtendBlockPoolImpl(void* region_base, size_t region_size,
         g_info->idle_list[block_type][i] = node[i];
         g_info->idle_size[block_type][i] += node[i]->len;
     }
+    g_info->region_num[block_type]++;
 
     return region_base;
 }
@@ -200,6 +203,16 @@ static void* ExtendBlockPool(size_t region_size, int block_type) {
     return ExtendBlockPoolImpl(region_base, region_size, block_type);
 }
 
+static bool CanExendBlockRuntime(int block_type) {
+    auto rc = FLAGS_rdma_memory_pool_buckets == 1 &&
+              g_info->region_num[block_type] == 1;
+    if (!rc) {
+        LOG_EVERY_SECOND(ERROR)
+            << "Runtime extend memory only support one bucket per block type";
+    }
+    return rc;
+}
+
 void* ExtendBlockPoolByUser(void* region_base, size_t region_size,
                             int block_type) {
     if (FLAGS_rdma_memory_pool_user_specified_memory == false) {
@@ -215,9 +228,7 @@ void* ExtendBlockPoolByUser(void* region_base, size_t region_size,
     BAIDU_SCOPED_LOCK(*g_info->lock[block_type][index]);
     BAIDU_SCOPED_LOCK(g_info->extend_lock);
 
-    if (g_region_num > 1 && FLAGS_rdma_memory_pool_buckets > 1) {
-        LOG_EVERY_SECOND(ERROR)
-            << "Runtime extend memory only support single bucket";
+    if (CanExendBlockRuntime(block_type) == false) {
         return NULL;
     }
     region_size =
@@ -300,6 +311,7 @@ bool InitBlockPool(RegisterCallback cb) {
         if (g_info->idle_size[i].size() != g_buckets) {
             return false;
         }
+        g_info->region_num[i] = 0;
         for (size_t j = 0; j < g_buckets; ++j) {
             g_info->lock[i][j] = new (std::nothrow) butil::Mutex;
             if (!g_info->lock[i][j]) {
@@ -349,6 +361,9 @@ static void* AllocBlockFrom(int block_type) {
         BAIDU_SCOPED_LOCK(g_info->extend_lock);
         node = g_info->idle_list[block_type][index];
         if (!node) {
+            if (CanExendBlockRuntime(block_type) == false) {
+                return NULL;
+            }
             // There is no block left, extend a new region
             if (!ExtendBlockPool(FLAGS_rdma_memory_pool_increase_size_mb,
                                  block_type)) {