Merge pull request #66526 from alexkats/gwp

Don't start GWP allocations until init is finished

programs/keeper/Keeper.cpp

@@ -52,6 +52,10 @@
 #    include <Server/CertificateReloader.h>
 #endif
 
+#if USE_GWP_ASAN
+#    include <Common/GWPAsan.h>
+#endif
+
 #include <Server/ProtocolServerAdapter.h>
 #include <Server/KeeperTCPHandlerFactory.h>
 
@@ -639,6 +643,10 @@ try
         tryLogCurrentException(log, "Disabling cgroup memory observer because of an error during initialization");
     }
 
+#if USE_GWP_ASAN
+    GWPAsan::initFinished();
+#endif
+
 
     LOG_INFO(log, "Ready for connections.");
 

programs/server/Server.cpp

@@ -2213,6 +2213,7 @@ try
         CannotAllocateThreadFaultInjector::setFaultProbability(server_settings.cannot_allocate_thread_fault_injection_probability);
 
 #if USE_GWP_ASAN
+        GWPAsan::initFinished();
         GWPAsan::setForceSampleProbability(server_settings.gwp_asan_force_sample_probability);
 #endif
 

src/Client/ClientBase.cpp

@@ -81,6 +81,10 @@
 #include <Common/config_version.h>
 #include "config.h"
 
+#if USE_GWP_ASAN
+#    include <Common/GWPAsan.h>
+#endif
+
 
 namespace fs = std::filesystem;
 using namespace std::literals;
@@ -3264,6 +3268,11 @@ void ClientBase::init(int argc, char ** argv)
     fatal_log = createLogger("ClientBase", fatal_channel_ptr.get(), Poco::Message::PRIO_FATAL);
     signal_listener = std::make_unique<SignalListener>(nullptr, fatal_log);
     signal_listener_thread.start(*signal_listener);
+
+#if USE_GWP_ASAN
+    GWPAsan::initFinished();
+#endif
+
 }
 
 }

src/Common/Allocator.cpp

@@ -68,7 +68,7 @@ void * allocNoTrack(size_t size, size_t alignment)
 {
     void * buf;
 #if USE_GWP_ASAN
-    if (unlikely(GWPAsan::GuardedAlloc.shouldSample()))
+    if (unlikely(GWPAsan::shouldSample()))
     {
         if (void * ptr = GWPAsan::GuardedAlloc.allocate(size, alignment))
         {
@@ -185,7 +185,7 @@ void * Allocator<clear_memory_, populate>::realloc(void * buf, size_t old_size,
     }
 
 #if USE_GWP_ASAN
-    if (unlikely(GWPAsan::GuardedAlloc.shouldSample()))
+    if (unlikely(GWPAsan::shouldSample()))
     {
         auto trace_alloc = CurrentMemoryTracker::alloc(new_size);
         if (void * ptr = GWPAsan::GuardedAlloc.allocate(new_size, alignment))

src/Common/GWPAsan.cpp

@@ -217,6 +217,13 @@ void printReport([[maybe_unused]] uintptr_t fault_address)
         reinterpret_cast<void **>(trace.data()), 0, trace_length, [&](const auto line) { LOG_FATAL(logger, fmt::runtime(line)); });
 }
 
+std::atomic<bool> init_finished = false;
+
+void initFinished()
+{
+    init_finished.store(true, std::memory_order_relaxed);
+}
+
 std::atomic<double> force_sample_probability = 0.0;
 
 void setForceSampleProbability(double value)

src/Common/GWPAsan.h

@@ -19,12 +19,30 @@ bool isGWPAsanError(uintptr_t fault_address);
 
 void printReport(uintptr_t fault_address);
 
+extern std::atomic<bool> init_finished;
+
+void initFinished();
+
 extern std::atomic<double> force_sample_probability;
 
 void setForceSampleProbability(double value);
 
+/**
+ * We'd like to postpone sampling allocations under the startup is finished. There are mainly
+ * two reasons for that:
+ *
+ * - To avoid complex issues with initialization order
+ * - Don't waste MaxSimultaneousAllocations on global objects as it's not useful
+*/
+inline bool shouldSample()
+{
+    return init_finished.load(std::memory_order_relaxed) && GuardedAlloc.shouldSample();
+}
+
 inline bool shouldForceSample()
 {
+    if (!init_finished.load(std::memory_order_relaxed))
+        return false;
     std::bernoulli_distribution dist(force_sample_probability.load(std::memory_order_relaxed));
     return dist(thread_local_rng);
 }

src/Common/memory.h

@@ -37,7 +37,7 @@ requires DB::OptionalArgument<TAlign...>
 inline ALWAYS_INLINE void * newImpl(std::size_t size, TAlign... align)
 {
 #if USE_GWP_ASAN
-    if (unlikely(GWPAsan::GuardedAlloc.shouldSample()))
+    if (unlikely(GWPAsan::shouldSample()))
     {
         if constexpr (sizeof...(TAlign) == 1)
         {
@@ -83,7 +83,7 @@ inline ALWAYS_INLINE void * newImpl(std::size_t size, TAlign... align)
 inline ALWAYS_INLINE void * newNoExcept(std::size_t size) noexcept
 {
 #if USE_GWP_ASAN
-    if (unlikely(GWPAsan::GuardedAlloc.shouldSample()))
+    if (unlikely(GWPAsan::shouldSample()))
     {
         if (void * ptr = GWPAsan::GuardedAlloc.allocate(size))
         {
@@ -102,7 +102,7 @@ inline ALWAYS_INLINE void * newNoExcept(std::size_t size) noexcept
 inline ALWAYS_INLINE void * newNoExcept(std::size_t size, std::align_val_t align) noexcept
 {
 #if USE_GWP_ASAN
-    if (unlikely(GWPAsan::GuardedAlloc.shouldSample()))
+    if (unlikely(GWPAsan::shouldSample()))
     {
         if (void * ptr = GWPAsan::GuardedAlloc.allocate(size, alignToSizeT(align)))
         {

tests/queries/0_stateless/02440_mutations_finalization.sql

@@ -1,3 +1,4 @@
+-- Tags: no-parallel
 
 create table mut (n int) engine=ReplicatedMergeTree('/test/02440/{database}/mut', '1') order by tuple();
 set insert_keeper_fault_injection_probability=0;