[k2] add AllocatorState and libc_alloc_guard (#1316)

apolyakov · web-flow · commit 46d481be7245 · 2025-05-16T18:45:04.000+03:00
diff --git a/runtime-light/allocator-wrapper/libc-alloc-wrapper.cpp b/runtime-light/allocator-wrapper/libc-alloc-wrapper.cpp
@@ -4,20 +4,34 @@
 
 #include <cstddef>
 
+#include "runtime-common/core/allocator/script-malloc-interface.h"
+#include "runtime-light/allocator/allocator-state.h"
 #include "runtime-light/utils/logs.h"
 
-extern "C" void* __wrap_malloc([[maybe_unused]] size_t size) noexcept {
-  kphp::log::error("unexpected use of malloc");
+extern "C" void* __wrap_malloc(size_t size) noexcept {
+  if (!AllocatorState::get().libc_alloc_allowed()) [[unlikely]] {
+    kphp::log::error("unexpected use of malloc");
+  }
+  return kphp::memory::script::alloc(size);
 }
 
-extern "C" void __wrap_free([[maybe_unused]] void* ptr) noexcept {
-  kphp::log::error("unexpected use of free");
+extern "C" void __wrap_free(void* ptr) noexcept {
+  if (!AllocatorState::get().libc_alloc_allowed()) [[unlikely]] {
+    kphp::log::error("unexpected use of free");
+  }
+  kphp::memory::script::free(ptr);
 }
 
-extern "C" void* __wrap_calloc([[maybe_unused]] size_t nmemb, [[maybe_unused]] size_t size) noexcept {
-  kphp::log::error("unexpected use of calloc");
+extern "C" void* __wrap_calloc(size_t nmemb, size_t size) noexcept {
+  if (!AllocatorState::get().libc_alloc_allowed()) [[unlikely]] {
+    kphp::log::error("unexpected use of calloc");
+  }
+  return kphp::memory::script::calloc(nmemb, size);
 }
 
 extern "C" void* __wrap_realloc([[maybe_unused]] void* ptr, [[maybe_unused]] size_t size) noexcept {
-  kphp::log::error("unexpected use of realloc");
+  if (!AllocatorState::get().libc_alloc_allowed()) [[unlikely]] {
+    kphp::log::error("unexpected use of realloc");
+  }
+  return kphp::memory::script::realloc(ptr, size);
 }
diff --git a/runtime-light/allocator/allocator-state.cpp b/runtime-light/allocator/allocator-state.cpp
@@ -0,0 +1,26 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#include "runtime-light/allocator/allocator-state.h"
+
+#include "runtime-light/k2-platform/k2-api.h"
+#include "runtime-light/state/component-state.h"
+#include "runtime-light/state/image-state.h"
+#include "runtime-light/state/instance-state.h"
+#include "runtime-light/utils/logs.h"
+
+const AllocatorState& AllocatorState::get() noexcept {
+  if (const auto* instance_state_ptr{k2::instance_state()}; instance_state_ptr != nullptr) [[likely]] {
+    return instance_state_ptr->instance_allocator_state;
+  } else if (const auto* component_state_ptr{k2::component_state()}; component_state_ptr != nullptr) {
+    return component_state_ptr->component_allocator_state;
+  } else if (const auto* image_state_ptr{k2::image_state()}; image_state_ptr != nullptr) {
+    return image_state_ptr->image_allocator_state;
+  }
+  kphp::log::error("can't find allocator state");
+}
+
+AllocatorState& AllocatorState::get_mutable() noexcept {
+  return const_cast<AllocatorState&>(AllocatorState::get());
+}
diff --git a/runtime-light/allocator/allocator-state.h b/runtime-light/allocator/allocator-state.h
@@ -0,0 +1,35 @@
+// Compiler for PHP (aka KPHP)
+// Copyright (c) 2025 LLC «V Kontakte»
+// Distributed under the GPL v3 License, see LICENSE.notice.txt
+
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+
+#include "common/mixin/not_copyable.h"
+#include "runtime-common/core/allocator/runtime-allocator.h"
+#include "runtime-light/utils/logs.h"
+
+class AllocatorState final : private vk::not_copyable {
+  uint32_t m_libc_alloc_allowed{};
+
+public:
+  RuntimeAllocator allocator;
+
+  void enable_libc_alloc() noexcept {
+    ++m_libc_alloc_allowed;
+  }
+  void disable_libc_alloc() noexcept {
+    kphp::log::assertion(m_libc_alloc_allowed-- != 0);
+  }
+  bool libc_alloc_allowed() const noexcept {
+    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) {}
+
+  static const AllocatorState& get() noexcept;
+  static AllocatorState& get_mutable() noexcept;
+};
diff --git a/runtime-light/allocator/allocator.cmake b/runtime-light/allocator/allocator.cmake
@@ -1 +1,2 @@
-set(RUNTIME_LIGHT_ALLOCATOR_SRC allocator/runtime-light-allocator.cpp)
+set(RUNTIME_LIGHT_ALLOCATOR_SRC allocator/allocator-state.cpp
+                                allocator/runtime-light-allocator.cpp)
diff --git a/runtime-light/allocator/allocator.h b/runtime-light/allocator/allocator.h
@@ -8,9 +8,29 @@
 #include <cstddef>
 
 #include "runtime-common/core/allocator/script-allocator-managed.h"
+#include "runtime-light/allocator/allocator-state.h"
 
 template<std::derived_from<ScriptAllocatorManaged> T, typename... Args>
 requires std::constructible_from<T, Args...>
 auto make_unique_on_script_memory(Args&&... args) noexcept {
   return std::make_unique<T>(std::forward<Args>(args)...);
 }
+
+namespace kphp::memory {
+
+struct libc_alloc_guard final {
+  libc_alloc_guard() noexcept {
+    AllocatorState::get_mutable().enable_libc_alloc();
+  }
+
+  ~libc_alloc_guard() {
+    AllocatorState::get_mutable().disable_libc_alloc();
+  }
+
+  libc_alloc_guard(const libc_alloc_guard&) = delete;
+  libc_alloc_guard(libc_alloc_guard&&) = delete;
+  libc_alloc_guard& operator=(const libc_alloc_guard&) = delete;
+  libc_alloc_guard& operator=(libc_alloc_guard&&) = delete;
+};
+
+} // namespace kphp::memory
diff --git a/runtime-light/allocator/runtime-light-allocator.cpp b/runtime-light/allocator/runtime-light-allocator.cpp
@@ -6,10 +6,8 @@
 #include <cstddef>
 #include <cstring>
 
+#include "runtime-light/allocator/allocator-state.h"
 #include "runtime-light/k2-platform/k2-api.h"
-#include "runtime-light/state/component-state.h"
-#include "runtime-light/state/image-state.h"
-#include "runtime-light/state/instance-state.h"
 #include "runtime-light/utils/logs.h"
 
 namespace {
@@ -27,14 +25,7 @@ void request_extra_memory(size_t requested_size) noexcept {
 } // namespace
 
 RuntimeAllocator& RuntimeAllocator::get() noexcept {
-  if (k2::instance_state() != nullptr) [[likely]] {
-    return InstanceState::get().allocator;
-  } else if (k2::component_state() != nullptr) {
-    return ComponentState::get_mutable().allocator;
-  } else if (k2::image_state() != nullptr) {
-    return ImageState::get_mutable().allocator;
-  }
-  kphp::log::error("no available allocators");
+  return AllocatorState::get_mutable().allocator;
 }
 
 RuntimeAllocator::RuntimeAllocator(size_t script_mem_size, size_t oom_handling_mem_size) {
diff --git a/runtime-light/state/component-state.h b/runtime-light/state/component-state.h
@@ -9,26 +9,20 @@
 #include <string_view>
 
 #include "common/mixin/not_copyable.h"
-#include "runtime-common/core/allocator/runtime-allocator.h"
 #include "runtime-common/core/runtime-core.h"
+#include "runtime-light/allocator/allocator-state.h"
 #include "runtime-light/k2-platform/k2-api.h"
 
 struct ComponentState final : private vk::not_copyable {
-  RuntimeAllocator allocator;
+  AllocatorState component_allocator_state{INIT_COMPONENT_ALLOCATOR_SIZE, 0};
 
-  const uint32_t argc;
-  const uint32_t envc;
+  const uint32_t argc{k2::args_count()};
+  const uint32_t envc{k2::env_count()};
+  array<string> ini_opts{array_size{argc, false}};
+  array<mixed> env{array_size{envc, false}};
   mixed runtime_config;
-  array<string> ini_opts;
-  array<mixed> env;
-
-  ComponentState() noexcept
-      : allocator(INIT_COMPONENT_ALLOCATOR_SIZE, 0),
-        argc(k2::args_count()),
-        envc(k2::env_count()),
-        ini_opts(array_size{argc, false}) /* overapproximation */
-        ,
-        env(array_size{envc, false}) {
+
+  ComponentState() noexcept {
     parse_env();
     parse_args();
   }
diff --git a/runtime-light/state/image-state.h b/runtime-light/state/image-state.h
@@ -11,8 +11,8 @@
 
 #include "common/mixin/not_copyable.h"
 #include "common/php-functions.h"
-#include "runtime-common/core/allocator/runtime-allocator.h"
 #include "runtime-common/core/runtime-core.h"
+#include "runtime-light/allocator/allocator-state.h"
 #include "runtime-light/k2-platform/k2-api.h"
 #include "runtime-light/stdlib/math/math-state.h"
 #include "runtime-light/stdlib/rpc/rpc-client-state.h"
@@ -21,10 +21,10 @@
 #include "runtime-light/utils/logs.h"
 
 struct ImageState final : private vk::not_copyable {
-  RuntimeAllocator allocator;
+  AllocatorState image_allocator_state{INIT_IMAGE_ALLOCATOR_SIZE, 0};
 
   char* c_linear_mem{nullptr};
-  uint32_t pid{};
+  uint32_t pid{k2::getpid()};
   string uname_info_s;
   string uname_info_n;
   string uname_info_r;
@@ -38,9 +38,7 @@ struct ImageState final : private vk::not_copyable {
   StringImageState string_image_state;
   MathImageState math_image_state;
 
-  ImageState() noexcept
-      : allocator(INIT_IMAGE_ALLOCATOR_SIZE, 0),
-        pid(k2::getpid()) {
+  ImageState() noexcept {
     utsname uname_info{};
     if (const auto errc{k2::uname(std::addressof(uname_info))}; errc != k2::errno_ok) [[unlikely]] {
       kphp::log::error("can't get uname, error code: {}", errc);
diff --git a/runtime-light/state/instance-state.h b/runtime-light/state/instance-state.h
@@ -11,9 +11,9 @@
 #include <utility>
 
 #include "common/mixin/not_copyable.h"
-#include "runtime-common/core/allocator/runtime-allocator.h"
 #include "runtime-common/core/runtime-core.h"
 #include "runtime-common/core/std/containers.h"
+#include "runtime-light/allocator/allocator-state.h"
 #include "runtime-light/core/globals/php-script-globals.h"
 #include "runtime-light/coroutine/task.h"
 #include "runtime-light/k2-platform/k2-api.h"
@@ -62,7 +62,7 @@ struct InstanceState final : vk::not_copyable {
   using list = kphp::stl::list<T, kphp::memory::script_allocator>;
 
   InstanceState() noexcept
-      : allocator(INIT_INSTANCE_ALLOCATOR_SIZE, 0) {}
+      : instance_allocator_state(INIT_INSTANCE_ALLOCATOR_SIZE, 0) {}
 
   ~InstanceState() = default;
 
@@ -106,7 +106,7 @@ struct InstanceState final : vk::not_copyable {
   void release_stream(uint64_t) noexcept;
   void release_all_streams() noexcept;
 
-  RuntimeAllocator allocator;
+  AllocatorState instance_allocator_state;
 
   CoroutineScheduler scheduler;
   ForkInstanceState fork_instance_state;

Original file line number	Diff line number	Diff line change
`@@ -1 +1,2 @@`
`1`		`-set(RUNTIME_LIGHT_ALLOCATOR_SRC allocator/runtime-light-allocator.cpp)`
	`1`	`+set(RUNTIME_LIGHT_ALLOCATOR_SRC allocator/allocator-state.cpp`
	`2`	`+ allocator/runtime-light-allocator.cpp)`