[improvement](memory) enable exception for get_free_block to avoid oom

dataroaring · Jun 8, 2024 · 26e2b42 · 26e2b42
1 parent 1a2c4bc
commit 26e2b42
Show file tree

Hide file tree

Showing 3 changed files with 10 additions and 8 deletions.
diff --git a/be/src/vec/exec/scan/scanner_context.cpp b/be/src/vec/exec/scan/scanner_context.cpp
@@ -149,18 +149,19 @@ std::string ScannerContext::parent_name() {
     return _local_state->get_name();
 }
 
-vectorized::BlockUPtr ScannerContext::get_free_block(bool force) {
-    vectorized::BlockUPtr block = nullptr;
+Status ScannerContext::get_free_block(bool force, vectorized::BlockUPtr& block) {
+    block = nullptr;
     if (_free_blocks.try_dequeue(block)) {
         DCHECK(block->mem_reuse());
         _free_blocks_memory_usage -= block->allocated_bytes();
         _free_blocks_memory_usage_mark->set(_free_blocks_memory_usage);
     } else if (_free_blocks_memory_usage < _max_bytes_in_queue || force) {
         _newly_create_free_blocks_num->update(1);
-        block = vectorized::Block::create_unique(_output_tuple_desc->slots(), 0,
-                                                 true /*ignore invalid slots*/);
+        RETURN_IF_CATCH_EXCEPTION(block = vectorized::Block::create_unique(_output_tuple_desc->slots(), 0,
+                                                true /*ignore invalid slots*/));
     }
-    return block;
+
+    return Status::OK();
 }
 
 void ScannerContext::return_free_block(vectorized::BlockUPtr block) {

diff --git a/be/src/vec/exec/scan/scanner_context.h b/be/src/vec/exec/scan/scanner_context.h
@@ -121,7 +121,7 @@ class ScannerContext : public std::enable_shared_from_this<ScannerContext>,
     }
     Status init();
 
-    vectorized::BlockUPtr get_free_block(bool force);
+    Status get_free_block(bool force, vectorized::BlockUPtr& block);
     void return_free_block(vectorized::BlockUPtr block);
     inline void inc_free_block_usage(size_t usage) {
         _free_blocks_memory_usage += usage;

diff --git a/be/src/vec/exec/scan/scanner_scheduler.cpp b/be/src/vec/exec/scan/scanner_scheduler.cpp
@@ -251,8 +251,9 @@ void ScannerScheduler::_scanner_scan(std::shared_ptr<ScannerContext> ctx,
             eos = true;
             break;
         }
-        BlockUPtr free_block = ctx->get_free_block(first_read);
-        if (free_block == nullptr) {
+        BlockUPtr free_block = nullptr;
+        status = ctx->get_free_block(first_read, free_block));
+        if (free_block == nullptr || !status.ok()) {
             break;
         }
         status = scanner->get_block_after_projects(state, free_block.get(), &eos);