[k2] make init instance memory size adjustable (#1610)

* Fixed diagnostic in cases when memory is not enough
* Made initial instance memory size adjustable
* Made min instance extra memory size adjustable
---------

Signed-off-by: Petr Shumilov <p.shumilov@vkteam.ru>

Author: PetrShumilov

runtime-common/core/allocator/runtime-allocator.h

@@ -13,7 +13,7 @@ struct RuntimeAllocator final : vk::not_copyable {
   static RuntimeAllocator& get() noexcept;
 
   RuntimeAllocator() = default;
-  RuntimeAllocator(size_t script_mem_size, size_t oom_handling_mem_size);
+  RuntimeAllocator(size_t script_mem_size, size_t min_extra_mem_size, size_t oom_handling_mem_size);
 
   void init(void* buffer, size_t script_mem_size, size_t oom_handling_mem_size);
   void free();
@@ -28,5 +28,12 @@ struct RuntimeAllocator final : vk::not_copyable {
   void* realloc_global_memory(void* mem, size_t new_size, size_t old_size) noexcept;
   void free_global_memory(void* mem, size_t size) noexcept;
 
+private:
+  void request_extra_memory(size_t requested_size, size_t min_extra_mem_size) noexcept;
+
+public:
   memory_resource::unsynchronized_pool_resource memory_resource;
+
+private:
+  size_t min_extra_mem_size{0};
 };

runtime-common/core/memory-resource/unsynchronized_pool_resource.cpp

@@ -10,8 +10,6 @@
 
 namespace memory_resource {
 
-constexpr size_t unsynchronized_pool_resource::MAX_CHUNK_BLOCK_SIZE_;
-
 void unsynchronized_pool_resource::init(void* buffer, size_t buffer_size, size_t oom_handling_buffer_size) noexcept {
   monotonic_buffer_resource::init(buffer, buffer_size);
 

runtime-common/core/memory-resource/unsynchronized_pool_resource.h

@@ -30,7 +30,7 @@ class unsynchronized_pool_resource : private monotonic_buffer_resource {
   void* allocate(size_t size) noexcept {
     void* mem = nullptr;
     const auto aligned_size = details::align_for_chunk(size);
-    if (aligned_size < MAX_CHUNK_BLOCK_SIZE_) {
+    if (aligned_size < MAX_CHUNK_BLOCK_SIZE) {
       mem = try_allocate_small_piece(aligned_size);
       if (!mem) {
         mem = allocate_small_piece_from_fallback_resource(aligned_size);
@@ -123,7 +123,7 @@ class unsynchronized_pool_resource : private monotonic_buffer_resource {
   void put_memory_back(void* mem, size_t size) noexcept {
     const bool from_extra_pool = (extra_memory_head_ != &extra_memory_tail_) && is_memory_from_extra_pool(mem, size);
     if (from_extra_pool || !monotonic_buffer_resource::put_memory_back(mem, size)) {
-      if (size < MAX_CHUNK_BLOCK_SIZE_) {
+      if (size < MAX_CHUNK_BLOCK_SIZE) {
         size_t chunk_id = details::get_chunk_id(size);
         free_chunks_[chunk_id].put_mem(mem);
         ++stats_.small_memory_pieces;
@@ -134,15 +134,18 @@ class unsynchronized_pool_resource : private monotonic_buffer_resource {
     }
   }
 
+public:
+  static constexpr size_t MAX_CHUNK_BLOCK_SIZE{static_cast<size_t>(16U * 1024U)};
+
+private:
   details::memory_chunk_tree huge_pieces_;
   monotonic_buffer_resource fallback_resource_;
   size_t oom_handling_memory_size_{0};
 
   extra_memory_pool* extra_memory_head_{nullptr};
   extra_memory_pool extra_memory_tail_{sizeof(extra_memory_pool)};
 
-  static constexpr size_t MAX_CHUNK_BLOCK_SIZE_{static_cast<size_t>(16U * 1024U)};
-  std::array<details::memory_chunk_list, details::get_chunk_id(MAX_CHUNK_BLOCK_SIZE_)> free_chunks_;
+  std::array<details::memory_chunk_list, details::get_chunk_id(MAX_CHUNK_BLOCK_SIZE)> free_chunks_;
 };
 
 } // namespace memory_resource

runtime-light/allocator/allocator-state.h

@@ -11,6 +11,8 @@
 #include "runtime-common/core/allocator/runtime-allocator.h"
 #include "runtime-light/stdlib/diagnostics/logs.h"
 
+inline constexpr auto DEFAULT_MIN_EXTRA_MEMORY_POOL_SIZE{static_cast<size_t>(64 * 1024U * 1024U)};
+
 class AllocatorState final : private vk::not_copyable {
   uint32_t m_libc_alloc_allowed{};
 
@@ -27,8 +29,8 @@ class AllocatorState final : private vk::not_copyable {
     return m_libc_alloc_allowed != 0;
   }
 
-  AllocatorState(size_t script_mem_size, size_t oom_handling_mem_size) noexcept
-      : allocator(script_mem_size, oom_handling_mem_size) {}
+  AllocatorState(size_t script_mem_size, size_t min_extra_mem_size, size_t oom_handling_mem_size) noexcept
+      : allocator(script_mem_size, min_extra_mem_size, oom_handling_mem_size) {}
 
   static const AllocatorState& get() noexcept;
   static AllocatorState& get_mutable() noexcept;

runtime-light/allocator/runtime-light-allocator.cpp

@@ -10,25 +10,12 @@
 #include "runtime-light/k2-platform/k2-api.h"
 #include "runtime-light/stdlib/diagnostics/logs.h"
 
-namespace {
-// TODO: make it depend on max chunk size, e.g. MIN_EXTRA_MEM_SIZE = f(MAX_CHUNK_SIZE);
-constexpr auto MIN_EXTRA_MEM_SIZE = static_cast<size_t>(1024U * 1024U); // extra mem size should be greater than max chunk block size
-constexpr auto EXTRA_MEMORY_MULTIPLIER = 2;
-
-void request_extra_memory(size_t requested_size) noexcept {
-  const size_t extra_mem_size{std::max(MIN_EXTRA_MEM_SIZE, EXTRA_MEMORY_MULTIPLIER * requested_size)};
-  auto& allocator{RuntimeAllocator::get()};
-  auto* extra_mem{allocator.alloc_global_memory(extra_mem_size)};
-  allocator.memory_resource.add_extra_memory(new (extra_mem) memory_resource::extra_memory_pool{extra_mem_size});
-}
-
-} // namespace
-
 RuntimeAllocator& RuntimeAllocator::get() noexcept {
   return AllocatorState::get_mutable().allocator;
 }
 
-RuntimeAllocator::RuntimeAllocator(size_t script_mem_size, size_t oom_handling_mem_size) {
+RuntimeAllocator::RuntimeAllocator(size_t script_mem_size, size_t min_extra_mem_size, size_t oom_handling_mem_size)
+    : min_extra_mem_size(min_extra_mem_size) {
   kphp::log::debug("create runtime allocator -> {:p}: script memory -> {}, oom handling size -> {}", reinterpret_cast<void*>(this), script_mem_size,
                    oom_handling_mem_size);
   void* buffer{alloc_global_memory(script_mem_size)};
@@ -57,7 +44,7 @@ void* RuntimeAllocator::alloc_script_memory(size_t size) noexcept {
   kphp::log::assertion(size != 0);
   void* mem{memory_resource.allocate(size)};
   if (mem == nullptr) [[unlikely]] {
-    request_extra_memory(size);
+    request_extra_memory(size, min_extra_mem_size);
     mem = memory_resource.allocate(size);
     kphp::log::assertion(mem != nullptr);
   }
@@ -68,7 +55,7 @@ void* RuntimeAllocator::alloc0_script_memory(size_t size) noexcept {
   kphp::log::assertion(size != 0);
   void* mem{memory_resource.allocate0(size)};
   if (mem == nullptr) [[unlikely]] {
-    request_extra_memory(size);
+    request_extra_memory(size, min_extra_mem_size);
     mem = memory_resource.allocate0(size);
     kphp::log::assertion(mem != nullptr);
   }
@@ -79,7 +66,7 @@ void* RuntimeAllocator::realloc_script_memory(void* old_mem, size_t new_size, si
   kphp::log::assertion(new_size > old_size);
   void* new_mem{memory_resource.reallocate(old_mem, new_size, old_size)};
   if (new_mem == nullptr) [[unlikely]] {
-    request_extra_memory(new_size * 2);
+    request_extra_memory(new_size * 2, min_extra_mem_size);
     new_mem = memory_resource.reallocate(old_mem, new_size, old_size);
     kphp::log::assertion(new_mem != nullptr);
   }
@@ -113,3 +100,20 @@ void* RuntimeAllocator::realloc_global_memory(void* old_mem, size_t new_size, si
 void RuntimeAllocator::free_global_memory(void* mem, size_t /*unused*/) noexcept {
   k2::free(mem);
 }
+
+void RuntimeAllocator::request_extra_memory(size_t requested_size, size_t min_extra_mem_size) noexcept {
+  // Extra mem size have to be greater than max chunk block
+  min_extra_mem_size = std::max(min_extra_mem_size, memory_resource::unsynchronized_pool_resource::MAX_CHUNK_BLOCK_SIZE);
+
+  // If the requested size is greater than or equal to half of min_extra_mem_size, it’s more efficient to allocate a multiple of the requested size.
+  size_t extra_mem_size{std::max(min_extra_mem_size, 2 * requested_size)};
+
+  // Take into account internal layout of `memory_resource::extra_memory_pool`
+  extra_mem_size += sizeof(memory_resource::extra_memory_pool);
+
+  kphp::log::debug("requested extra memory pool with size {} bytes, will be allocated {} bytes", requested_size, extra_mem_size);
+
+  auto* extra_mem{alloc_global_memory(extra_mem_size)};
+  kphp::log::assertion(extra_mem != nullptr);
+  memory_resource.add_extra_memory(new (extra_mem) memory_resource::extra_memory_pool{extra_mem_size});
+}

runtime-light/runtime-light.cpp

@@ -15,11 +15,9 @@
 #define VISIBILITY_DEFAULT __attribute__((visibility("default")))
 
 VISIBILITY_DEFAULT ImageState* k2_create_image() {
-  kphp::log::debug("start image state creation");
+  kphp::log::debug("start image state creation, requested {} bytes", sizeof(ImageState));
   auto* image_state_ptr{static_cast<ImageState*>(k2::alloc(sizeof(ImageState)))};
-  if (image_state_ptr == nullptr) [[unlikely]] {
-    kphp::log::error("can't allocate enough memory for image state");
-  }
+  kphp::log::assertion(image_state_ptr != nullptr); // Not enough memory for image state
   kphp::log::debug("finish image state creation");
   return image_state_ptr;
 }
@@ -32,11 +30,9 @@ VISIBILITY_DEFAULT void k2_init_image() {
 }
 
 VISIBILITY_DEFAULT ComponentState* k2_create_component() {
-  kphp::log::debug("start component state creation");
+  kphp::log::debug("start component state creation, requested {} bytes", sizeof(ComponentState));
   auto* component_state_ptr{static_cast<ComponentState*>(k2::alloc(sizeof(ComponentState)))};
-  if (component_state_ptr == nullptr) [[unlikely]] {
-    kphp::log::error("can't allocate enough memory for component state");
-  }
+  kphp::log::assertion(component_state_ptr != nullptr); // Not enough memory for component state
   kphp::log::debug("finish component state creation");
   return component_state_ptr;
 }
@@ -48,11 +44,9 @@ VISIBILITY_DEFAULT void k2_init_component() {
 }
 
 VISIBILITY_DEFAULT InstanceState* k2_create_instance() {
-  kphp::log::debug("start instance state creation");
+  kphp::log::debug("start instance state creation, requested {} bytes", sizeof(InstanceState));
   auto* instance_state_ptr{static_cast<InstanceState*>(k2::alloc(sizeof(InstanceState)))};
-  if (instance_state_ptr == nullptr) [[unlikely]] {
-    kphp::log::error("can't allocate enough memory for instance state");
-  }
+  kphp::log::assertion(instance_state_ptr != nullptr); // Not enough memory for instance state
   kphp::log::debug("finish instance state creation");
   return instance_state_ptr;
 }

runtime-light/state/component-state.cpp

@@ -4,14 +4,18 @@
 
 #include "runtime-light/state/component-state.h"
 
+#include <charconv>
 #include <cstddef>
+#include <cstdint>
 #include <cstring>
 #include <iterator>
 #include <memory>
+#include <optional>
 #include <span>
 #include <string_view>
 #include <sys/mman.h>
 #include <sys/stat.h>
+#include <system_error>
 #include <utility>
 
 #include "runtime-common/core/runtime-core.h"
@@ -20,6 +24,24 @@
 #include "runtime-light/stdlib/diagnostics/logs.h"
 #include "runtime-light/stdlib/file/resource.h"
 
+namespace {
+
+std::optional<uint64_t> parse_uint64(std::string_view value_view) noexcept {
+  // Num of symbols in max u64 (18_446_744_073_709_551_615) is 20
+  if (value_view.empty() || value_view.size() > 20) {
+    return std::nullopt;
+  }
+
+  uint64_t res{};
+  if (const auto [ptr, err_code]{std::from_chars(value_view.begin(), value_view.end(), /* out paremeter */ res)};
+      err_code != std::errc{} || ptr != value_view.end()) {
+    return std::nullopt;
+  }
+  return res;
+}
+
+} // namespace
+
 void ComponentState::parse_env() noexcept {
   for (auto i = 0; i < envc; ++i) {
     const auto [env_key, env_value]{k2::env_fetch(i)};
@@ -92,6 +114,24 @@ void ComponentState::parse_exit_after_response_arg(std::string_view value_view)
   }
 }
 
+void ComponentState::parse_initial_instance_memory_size_arg(std::string_view value_view) noexcept {
+  const auto parsed{parse_uint64(value_view)};
+  if (!parsed) {
+    kphp::log::error("couldn't parse initial instance memory size, got {}", value_view);
+  }
+  initial_instance_memory_size = *parsed;
+  kphp::log::info("set initial instance memory size to {} bytes", initial_instance_memory_size);
+}
+
+void ComponentState::parse_min_instance_extra_memory_size_arg(std::string_view value_view) noexcept {
+  const auto parsed{parse_uint64(value_view)};
+  if (!parsed) {
+    kphp::log::error("couldn't parse min instance extra memory size, got {}", value_view);
+  }
+  min_instance_extra_memory_size = *parsed;
+  kphp::log::info("set min instance extra memory size to {} bytes", min_instance_extra_memory_size);
+}
+
 void ComponentState::parse_args() noexcept {
   for (auto i = 0; i < argc; ++i) {
     const auto [arg_key, arg_value]{k2::arg_fetch(i)};
@@ -108,6 +148,10 @@ void ComponentState::parse_args() noexcept {
       parse_cluster_name_arg(value_view);
     } else if (key_view == EXIT_AFTER_RESPONSE_ARG) {
       parse_exit_after_response_arg(value_view);
+    } else if (key_view == INITIAL_INSTANCE_MEMORY_SIZE_ARG) {
+      parse_initial_instance_memory_size_arg(value_view);
+    } else if (key_view == MIN_INSTANCE_EXTRA_MEMORY_SIZE_ARG) {
+      parse_min_instance_extra_memory_size_arg(value_view);
     } else {
       kphp::log::warning("unexpected argument format: {}", key_view);
     }

runtime-light/state/component-state.h

@@ -19,7 +19,7 @@
 #include "runtime-light/stdlib/kml/kml-state.h"
 
 struct ComponentState final : private vk::not_copyable {
-  AllocatorState component_allocator_state{INIT_COMPONENT_ALLOCATOR_SIZE, 0};
+  AllocatorState component_allocator_state{INIT_COMPONENT_ALLOCATOR_SIZE, DEFAULT_MIN_EXTRA_MEMORY_POOL_SIZE, 0};
   KmlComponentState kml_component_state; // This member does not hold any KPHP types, so setting a reference counter is unnecessary.
 
   const uint32_t argc{k2::args_count()};
@@ -29,6 +29,8 @@ struct ComponentState final : private vk::not_copyable {
   mixed runtime_config;
   string cluster_name{DEFAULT_CLUSTER_NAME.data(), DEFAULT_CLUSTER_NAME.size()};
   bool exit_after_response{};
+  uint64_t initial_instance_memory_size{INIT_INSTANCE_ALLOCATOR_SIZE};
+  uint64_t min_instance_extra_memory_size{DEFAULT_MIN_EXTRA_MEMORY_POOL_SIZE};
 
   ComponentState() noexcept {
     parse_env();
@@ -59,7 +61,10 @@ struct ComponentState final : private vk::not_copyable {
   static constexpr std::string_view CLUSTER_NAME_ARG = "cluster-name";
   static constexpr std::string_view DEFAULT_CLUSTER_NAME = "default";
   static constexpr std::string_view EXIT_AFTER_RESPONSE_ARG = "exit-after-response";
-  static constexpr auto INIT_COMPONENT_ALLOCATOR_SIZE = static_cast<size_t>(1024U * 1024U); // 1MB
+  static constexpr std::string_view INITIAL_INSTANCE_MEMORY_SIZE_ARG = "initial-instance-memory-size";
+  static constexpr std::string_view MIN_INSTANCE_EXTRA_MEMORY_SIZE_ARG = "min-instance-extra-memory-size";
+  static constexpr auto INIT_COMPONENT_ALLOCATOR_SIZE = static_cast<size_t>(1024U * 1024U);      // 1MB
+  static constexpr auto INIT_INSTANCE_ALLOCATOR_SIZE = static_cast<size_t>(64U * 1024U * 1024U); // 64MB
 
   void parse_env() noexcept;
 
@@ -74,4 +79,8 @@ struct ComponentState final : private vk::not_copyable {
   void parse_cluster_name_arg(std::string_view) noexcept;
 
   void parse_exit_after_response_arg(std::string_view) noexcept;
+
+  void parse_initial_instance_memory_size_arg(std::string_view) noexcept;
+
+  void parse_min_instance_extra_memory_size_arg(std::string_view) noexcept;
 };

runtime-light/state/image-state.h

@@ -29,7 +29,7 @@
 #include "runtime-light/stdlib/visitors/shape-visitors.h"
 
 struct ImageState final : private vk::not_copyable {
-  AllocatorState image_allocator_state{INIT_IMAGE_ALLOCATOR_SIZE, 0};
+  AllocatorState image_allocator_state{INIT_IMAGE_ALLOCATOR_SIZE, DEFAULT_MIN_EXTRA_MEMORY_POOL_SIZE, 0};
 
   uint32_t pid{k2::getpid()};
   uid_t uid{k2::getuid()};

runtime-light/state/instance-state.h

@@ -88,7 +88,7 @@ struct InstanceState final : vk::not_copyable {
     return instance_kind_;
   }
 
-  AllocatorState instance_allocator_state{INIT_INSTANCE_ALLOCATOR_SIZE, 0};
+  AllocatorState instance_allocator_state{ComponentState::get().initial_instance_memory_size, ComponentState::get().min_instance_extra_memory_size, 0};
 
   kphp::log::contextual_tags instance_tags;
   kphp::coro::io_scheduler io_scheduler;
@@ -135,6 +135,4 @@ struct InstanceState final : vk::not_copyable {
 
   enum image_kind image_kind_ { image_kind::invalid };
   enum instance_kind instance_kind_ { instance_kind::invalid };
-
-  static constexpr auto INIT_INSTANCE_ALLOCATOR_SIZE = static_cast<size_t>(128U * 1024U * 1024U);
 };