Merge pull request #25198 from vespa-engine/geirst/more-robust-initialization-of-proton-components

Make initialization of proton components more robust regarding resour…

7 files changed, 191 insertions, 10 deletions

diff --git a/searchcore/CMakeLists.txt b/searchcore/CMakeLists.txt
index 0f7d90b1491..d6c353e46c9 100644
--- a/searchcore/CMakeLists.txt
+++ b/searchcore/CMakeLists.txt
@@ -151,6 +151,7 @@ vespa_define_module(
     src/tests/proton/server/disk_mem_usage_metrics
     src/tests/proton/server/disk_mem_usage_sampler
     src/tests/proton/server/health_adapter
+    src/tests/proton/server/initialize_threads_calculator
     src/tests/proton/server/memory_flush_config_updater
     src/tests/proton/server/memoryflush
     src/tests/proton/server/shared_threading_service
diff --git a/searchcore/src/tests/proton/server/initialize_threads_calculator/CMakeLists.txt b/searchcore/src/tests/proton/server/initialize_threads_calculator/CMakeLists.txt
new file mode 100644
index 00000000000..e3ce152384c
--- /dev/null
+++ b/searchcore/src/tests/proton/server/initialize_threads_calculator/CMakeLists.txt
@@ -0,0 +1,9 @@
+# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+vespa_add_executable(searchcore_initialize_threads_calculator_test_app TEST
+    SOURCES
+    initialize_threads_calculator_test.cpp
+    DEPENDS
+    searchcore_server
+    GTest::GTest
+)
+vespa_add_test(NAME searchcore_initialize_threads_calculator_test_app COMMAND searchcore_initialize_threads_calculator_test_app)
diff --git a/searchcore/src/tests/proton/server/initialize_threads_calculator/initialize_threads_calculator_test.cpp b/searchcore/src/tests/proton/server/initialize_threads_calculator/initialize_threads_calculator_test.cpp
new file mode 100644
index 00000000000..7aad69076ca
--- /dev/null
+++ b/searchcore/src/tests/proton/server/initialize_threads_calculator/initialize_threads_calculator_test.cpp
@@ -0,0 +1,64 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include <vespa/searchcore/proton/server/initialize_threads_calculator.h>
+#include <vespa/searchlib/test/directory_handler.h>
+#include <vespa/vespalib/gtest/gtest.h>
+#include <vespa/vespalib/util/threadstackexecutor.h>
+
+using namespace proton;
+using vespalib::ThreadStackExecutor;
+
+class InitializeThreadsCalculatorTest : public search::test::DirectoryHandler, public testing::Test {
+public:
+    InitializeThreadsCalculatorTest() : DirectoryHandler("tmp") {}
+};
+
+void
+expect_successful_init(uint32_t exp_threads)
+{
+    InitializeThreadsCalculator i("tmp", 9);
+    EXPECT_EQ(exp_threads, i.num_threads());
+    EXPECT_TRUE(i.threads().get() != nullptr);
+    EXPECT_EQ(exp_threads, dynamic_cast<const ThreadStackExecutor&>(*i.threads()).getNumThreads());
+    i.init_done();
+    EXPECT_TRUE(i.threads().get() == nullptr);
+}
+
+void
+expect_aborted_init(uint32_t exp_threads, uint32_t cfg_threads = 9)
+{
+    InitializeThreadsCalculator i("tmp", cfg_threads);
+    EXPECT_EQ(exp_threads, i.num_threads());
+    EXPECT_TRUE(i.threads().get() != nullptr);
+    EXPECT_EQ(exp_threads, dynamic_cast<const ThreadStackExecutor&>(*i.threads()).getNumThreads());
+}
+
+TEST_F(InitializeThreadsCalculatorTest, initialize_threads_unchanged_when_init_is_successful)
+{
+    expect_successful_init(9);
+    // The previous init was successful,
+    // so we still use the configured number of initialize threads.
+    expect_successful_init(9);
+}
+
+TEST_F(InitializeThreadsCalculatorTest, initialize_threads_cut_in_half_when_init_is_aborted)
+{
+    expect_aborted_init(9);
+    expect_aborted_init(4);
+    expect_aborted_init(2);
+    expect_aborted_init(1);
+    expect_aborted_init(1);
+}
+
+TEST_F(InitializeThreadsCalculatorTest, zero_initialize_threads_is_special)
+{
+    {
+        InitializeThreadsCalculator i("tmp", 0);
+        EXPECT_EQ(0, i.num_threads());
+        EXPECT_TRUE(i.threads().get() == nullptr);
+    }
+    expect_aborted_init(1, 0);
+    expect_aborted_init(1, 0);
+}
+
+GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchcore/src/vespa/searchcore/proton/server/CMakeLists.txt b/searchcore/src/vespa/searchcore/proton/server/CMakeLists.txt
index 544e693b093..c1f7b24b05c 100644
--- a/searchcore/src/vespa/searchcore/proton/server/CMakeLists.txt
+++ b/searchcore/src/vespa/searchcore/proton/server/CMakeLists.txt
@@ -57,6 +57,7 @@ vespa_add_library(searchcore_server STATIC
     hw_info_explorer.cpp
     idocumentdbowner.cpp
     ifeedview.cpp
+    initialize_threads_calculator.cpp
     ireplayconfig.cpp
     job_tracked_maintenance_job.cpp
     lid_space_compaction_handler.cpp
diff --git a/searchcore/src/vespa/searchcore/proton/server/initialize_threads_calculator.cpp b/searchcore/src/vespa/searchcore/proton/server/initialize_threads_calculator.cpp
new file mode 100644
index 00000000000..727e4bb1e58
--- /dev/null
+++ b/searchcore/src/vespa/searchcore/proton/server/initialize_threads_calculator.cpp
@@ -0,0 +1,70 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#include "initialize_threads_calculator.h"
+#include <vespa/vespalib/util/cpu_usage.h>
+#include <vespa/vespalib/util/size_literals.h>
+#include <vespa/vespalib/util/threadstackexecutor.h>
+#include <fstream>
+
+using vespalib::CpuUsage;
+using CpuCategory = vespalib::CpuUsage::Category;
+
+namespace {
+
+void
+write(const vespalib::string& path, uint32_t num_threads)
+{
+    std::ofstream file;
+    file.open(path);
+    file << num_threads;
+    file.close();
+}
+
+uint32_t
+read(const vespalib::string& path)
+{
+    std::ifstream file;
+    file.open(path);
+    uint32_t result;
+    file >> result;
+    file.close();
+    return result;
+}
+
+VESPA_THREAD_STACK_TAG(proton_initialize_executor)
+
+const vespalib::string file_name = "initialize-threads.txt";
+
+}
+
+namespace proton {
+
+InitializeThreadsCalculator::InitializeThreadsCalculator(const vespalib::string& base_dir,
+                                                         uint32_t configured_num_threads)
+    : _path(base_dir + "/" + file_name),
+      _num_threads(configured_num_threads),
+      _threads()
+{
+    if (std::filesystem::exists(_path)) {
+        _num_threads = read(_path.c_str());
+        _num_threads = std::max(1u, (_num_threads / 2));
+        std::filesystem::remove(_path);
+    }
+    write(_path.c_str(), _num_threads);
+    if (_num_threads > 0) {
+        _threads = std::make_shared<vespalib::ThreadStackExecutor>(_num_threads, 128_Ki,
+                                                                   CpuUsage::wrap(proton_initialize_executor, CpuCategory::SETUP));
+    }
+}
+
+InitializeThreadsCalculator::~InitializeThreadsCalculator() = default;
+
+void
+InitializeThreadsCalculator::init_done()
+{
+    std::filesystem::remove(_path);
+    _threads = InitializeThreads();
+}
+
+}
+
diff --git a/searchcore/src/vespa/searchcore/proton/server/initialize_threads_calculator.h b/searchcore/src/vespa/searchcore/proton/server/initialize_threads_calculator.h
new file mode 100644
index 00000000000..509db24ebf3
--- /dev/null
+++ b/searchcore/src/vespa/searchcore/proton/server/initialize_threads_calculator.h
@@ -0,0 +1,36 @@
+// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/stllike/string.h>
+#include <vespa/vespalib/util/threadexecutor.h>
+#include <filesystem>
+#include <memory>
+
+namespace proton {
+
+/**
+ * Class that is used to calculate the number of threads to use
+ * during the initialization of proton components.
+ *
+ * The number of threads is cut in half each time the initialization of proton components is aborted,
+ * e.g. due to running out of memory.
+ * This adjustment should ensure that we eventually are able to initialize and start proton.
+ */
+ class InitializeThreadsCalculator {
+ private:
+     using InitializeThreads = std::shared_ptr<vespalib::ThreadExecutor>;
+     std::filesystem::path _path;
+     uint32_t _num_threads;
+     InitializeThreads _threads;
+
+ public:
+     InitializeThreadsCalculator(const vespalib::string& base_dir,
+                                 uint32_t configured_num_threads);
+     ~InitializeThreadsCalculator();
+     uint32_t num_threads() const { return _num_threads; }
+     InitializeThreads threads() const { return _threads; }
+     void init_done();
+ };
+
+}
diff --git a/searchcore/src/vespa/searchcore/proton/server/proton.cpp b/searchcore/src/vespa/searchcore/proton/server/proton.cpp
index 5c82b6e74ae..636d898ed5d 100644
--- a/searchcore/src/vespa/searchcore/proton/server/proton.cpp
+++ b/searchcore/src/vespa/searchcore/proton/server/proton.cpp
@@ -3,9 +3,11 @@
 #include "proton.h"
 #include "disk_mem_usage_sampler.h"
 #include "document_db_explorer.h"
+#include "documentdbconfig.h"
 #include "fileconfigmanager.h"
 #include "flushhandlerproxy.h"
 #include "hw_info_explorer.h"
+#include "initialize_threads_calculator.h"
 #include "memoryflush.h"
 #include "persistencehandlerproxy.h"
 #include "prepare_restart_handler.h"
@@ -15,7 +17,6 @@
 #include "resource_usage_explorer.h"
 #include "searchhandlerproxy.h"
 #include "simpleflush.h"
-#include "documentdbconfig.h"
 
 #include <vespa/document/base/exceptions.h>
 #include <vespa/document/datatype/documenttype.h>
@@ -147,7 +148,6 @@ struct MetricsUpdateHook : metrics::UpdateHook
 
 const vespalib::string CUSTOM_COMPONENT_API_PATH = "/state/v1/custom/component";
 
-VESPA_THREAD_STACK_TAG(proton_initialize_executor)
 VESPA_THREAD_STACK_TAG(proton_close_executor)
 
 void ensureWritableDir(const vespalib::string &dirName) {
@@ -351,14 +351,12 @@ Proton::init(const BootstrapConfig::SP & configSnapshot)
 
     vespalib::string fileConfigId;
     _compile_cache_executor_binding = vespalib::eval::CompileCache::bind(_shared_service->shared_raw());
-    InitializeThreads initializeThreads;
-    if (protonConfig.initialize.threads > 0) {
-        initializeThreads = std::make_shared<vespalib::ThreadStackExecutor>(protonConfig.initialize.threads, 128_Ki,
-                                                                            CpuUsage::wrap(proton_initialize_executor, CpuCategory::SETUP));
-        _initDocumentDbsInSequence = (protonConfig.initialize.threads == 1);
-    }
-    _protonConfigurer.applyInitialConfig(initializeThreads);
-    initializeThreads.reset();
+
+    InitializeThreadsCalculator calc(protonConfig.basedir, protonConfig.initialize.threads);
+    LOG(info, "Start initializing components: threads=%u, configured=%u",
+        calc.num_threads(), protonConfig.initialize.threads);
+    _initDocumentDbsInSequence = (calc.num_threads() == 1);
+    _protonConfigurer.applyInitialConfig(calc.threads());
 
     _prepareRestartHandler = std::make_unique<PrepareRestartHandler>(*_flushEngine);
     RPCHooks::Params rpcParams(*this, protonConfig.rpcport, _configUri, protonConfig.slobrokconfigid,
@@ -367,6 +365,8 @@ Proton::init(const BootstrapConfig::SP & configSnapshot)
     _metricsEngine->addExternalMetrics(_rpcHooks->proto_rpc_adapter_metrics());
 
     waitForInitDone();
+    LOG(info, "Done initializing components");
+    calc.init_done();
 
     _metricsEngine->start(_configUri);
     _stateServer = std::make_unique<vespalib::StateServer>(protonConfig.httpport, _healthAdapter,