From 404d7d23cb5deaade0964409f18521649be25dfb Mon Sep 17 00:00:00 2001 From: Gabor Gyimesi Date: Wed, 21 Oct 2026 14:59:15 +0200 Subject: [PATCH 1/6] MINIFICPP-2933 Stop MiNiFi C++ if C2 connection is lost for some time --- C2.md | 10 ++ .../containers/container_linux.py | 13 ++ .../containers/container_protocol.py | 3 + .../containers/container_windows.py | 13 ++ .../src/minifi_behave/steps/checking_steps.py | 7 + bin/minifi.service.in | 1 + conf/minifi.properties.in | 2 + .../tests/features/minifi_c2_server.feature | 20 +++ .../tests/features/steps/steps.py | 25 +++- libminifi/include/FlowController.h | 4 +- libminifi/include/RequestShutdown.h | 32 +++++ libminifi/include/c2/C2Agent.h | 8 +- libminifi/include/c2/C2ConnectionWatchdog.h | 46 +++++++ libminifi/include/c2/C2ExitCodes.h | 25 ++++ libminifi/src/Configuration.cpp | 2 + libminifi/src/FlowController.cpp | 4 +- libminifi/src/c2/C2Agent.cpp | 45 ++++++- libminifi/src/c2/C2ConnectionWatchdog.cpp | 52 ++++++++ .../integration/C2ConnectionTimeoutTest.cpp | 93 +++++++++++++ .../libtest/integration/IntegrationBase.cpp | 10 +- .../libtest/integration/IntegrationBase.h | 6 + .../C2AgentWatchdogConfigurationTests.cpp | 123 ++++++++++++++++++ .../test/unit/C2ConnectionWatchdogTests.cpp | 103 +++++++++++++++ libminifi/test/unit/C2ShutdownRecorder.h | 53 ++++++++ .../minifi-cpp/properties/Configuration.h | 1 + minifi_main/MiNiFiMain.cpp | 16 ++- minifi_main/MiNiFiWindowsService.cpp | 17 ++- 27 files changed, 718 insertions(+), 16 deletions(-) create mode 100644 libminifi/include/RequestShutdown.h create mode 100644 libminifi/include/c2/C2ConnectionWatchdog.h create mode 100644 libminifi/include/c2/C2ExitCodes.h create mode 100644 libminifi/src/c2/C2ConnectionWatchdog.cpp create mode 100644 libminifi/test/integration/C2ConnectionTimeoutTest.cpp create mode 100644 libminifi/test/unit/C2AgentWatchdogConfigurationTests.cpp create mode 100644 libminifi/test/unit/C2ConnectionWatchdogTests.cpp create mode 100644 libminifi/test/unit/C2ShutdownRecorder.h diff --git a/C2.md b/C2.md index f934cd0e77..b338a52271 100644 --- a/C2.md +++ b/C2.md @@ -28,6 +28,7 @@ options defined are located in minifi.properties. - [Base Options](#base-options) - [Flow Id and URL](#flow-id-and-url) - [Agent Identifier Fallback](#agent-identifier-fallback) + - [C2 Connection Timeout](#c2-connection-timeout) - [Metrics](#metrics) - [UpdatePolicies](#updatepolicies) - [Triggers](#triggers) @@ -136,6 +137,15 @@ It is possible to set a persistent fallback agent id. This is needed so that the # in minifi.properties nifi.c2.agent.identifier.fallback=my_fallback_id + +#### C2 Connection Timeout + +There is an option to shut down the agent if the connection to the C2 server exceeds a specified timeout. This is controlled by the `nifi.c2.connection.timeout` property, which defines the maximum time allowed without a successful heartbeat response. If the connection is not established (or is lost) within this period, the agent will trigger a shutdown with exit code 90. If the agent is shut down with exit 90 neither the systemd service nor the Windows service will restart it. The feature is off by default, and can be enabled by setting a non-zero value. If the timeout is set to a value less than or equal to the heartbeat period, it will be automatically adjusted to twice the heartbeat period. + + # in minifi.properties + # timeout for C2 connection, setting to 0 disables the timeout + nifi.c2.connection.timeout=30 min + ### Metrics Command and Control metrics can be used to send metrics through the heartbeat or via the DESCRIBE diff --git a/behave_framework/src/minifi_behave/containers/container_linux.py b/behave_framework/src/minifi_behave/containers/container_linux.py index 034f53c69b..9df8567f68 100644 --- a/behave_framework/src/minifi_behave/containers/container_linux.py +++ b/behave_framework/src/minifi_behave/containers/container_linux.py @@ -330,6 +330,19 @@ def exited(self) -> bool: except Exception: return False + @property + def exit_code(self) -> int: + if not self.container: + return -1 + try: + self.container.reload() + return self.container.attrs["State"].get("ExitCode", -1) + except docker.errors.NotFound: + self.container = None + return -1 + except Exception: + return -1 + def get_number_of_files(self, directory_path: str) -> int: if not self.container: logger.warning("Container not running") diff --git a/behave_framework/src/minifi_behave/containers/container_protocol.py b/behave_framework/src/minifi_behave/containers/container_protocol.py index cc167c33b3..12e832c3b1 100644 --- a/behave_framework/src/minifi_behave/containers/container_protocol.py +++ b/behave_framework/src/minifi_behave/containers/container_protocol.py @@ -44,6 +44,9 @@ def get_logs(self) -> str: ... @property def exited(self) -> bool: ... + @property + def exit_code(self) -> int: ... + def get_number_of_files(self, directory_path: str) -> int: ... def verify_file_contents(self, directory_path: str, expected_contents: list[str]) -> bool: ... diff --git a/behave_framework/src/minifi_behave/containers/container_windows.py b/behave_framework/src/minifi_behave/containers/container_windows.py index 5117562b70..2e5f3d68c1 100644 --- a/behave_framework/src/minifi_behave/containers/container_windows.py +++ b/behave_framework/src/minifi_behave/containers/container_windows.py @@ -318,6 +318,19 @@ def exited(self) -> bool: except Exception: return False + @property + def exit_code(self) -> int: + if not self.container: + return -1 + try: + self.container.reload() + return self.container.attrs["State"].get("ExitCode", -1) + except docker.errors.NotFound: + self.container = None + return -1 + except Exception: + return -1 + def get_number_of_files(self, directory_path: str) -> int: if not self.container: return -1 diff --git a/behave_framework/src/minifi_behave/steps/checking_steps.py b/behave_framework/src/minifi_behave/steps/checking_steps.py index 83ff073b7d..a67f64c4d3 100644 --- a/behave_framework/src/minifi_behave/steps/checking_steps.py +++ b/behave_framework/src/minifi_behave/steps/checking_steps.py @@ -131,6 +131,13 @@ def verify_minifi_logs_do_not_contain_warnings(context: MinifiTestContext): ) +@then("the logs of the exited Minifi container contain the following message: '{message}'") +@then('the logs of the exited Minifi container contain the following message: "{message}"') +def verify_exited_minifi_container_logs_contain_message(context: MinifiTestContext, message: str): + assert context.containers[DEFAULT_MINIFI_CONTAINER_NAME].exited or log_due_to_failure(context) + assert message in context.containers[DEFAULT_MINIFI_CONTAINER_NAME].get_logs() or log_due_to_failure(context) + + @then("the logs of the '{container}' container contain the following message: '{message}' in less than {duration}") @then('the logs of the "{container}" container contain the following message: "{message}" in less than {duration}') def verify_container_logs_contain_message(context: MinifiTestContext, container: str, message: str, duration: str): diff --git a/bin/minifi.service.in b/bin/minifi.service.in index 2c7f84068f..0b8932cccf 100644 --- a/bin/minifi.service.in +++ b/bin/minifi.service.in @@ -12,6 +12,7 @@ RestartSec=3 KillSignal=SIGTERM TimeoutStopSec=20 RestartForceExitStatus=3 +RestartPreventExitStatus=90 [Install] WantedBy=multi-user.target diff --git a/conf/minifi.properties.in b/conf/minifi.properties.in index 722fc1ca71..8b55d19dbc 100644 --- a/conf/minifi.properties.in +++ b/conf/minifi.properties.in @@ -121,6 +121,8 @@ nifi.c2.full.heartbeat=false # nifi.c2.root.class.definitions.metrics.metrics.processorMetrics.name=ProcessorMetric # nifi.c2.root.class.definitions.metrics.metrics.processorMetrics.classes=GetFileMetrics +# nifi.c2.connection.timeout=30 min + ## Specify multiple heartbeat handlers. ## HeartbeatLogger logs the heartbeats on TRACE for debugging. #nifi.c2.agent.heartbeat.reporter.classes=HeartbeatLogger diff --git a/extensions/standard-processors/tests/features/minifi_c2_server.feature b/extensions/standard-processors/tests/features/minifi_c2_server.feature index c35ddeed95..9daba8f7df 100644 --- a/extensions/standard-processors/tests/features/minifi_c2_server.feature +++ b/extensions/standard-processors/tests/features/minifi_c2_server.feature @@ -55,3 +55,23 @@ Feature: MiNiFi can communicate with Apache NiFi MiNiFi C2 server Then the MiNiFi C2 server logs contain the following message: "acknowledged with a state of FULLY_APPLIED(DONE)" in less than 30 seconds And a single file with the content "test" is placed in the "/tmp/output" directory in less than 10 seconds And the Minifi logs do not contain the following message: "Failed to parse json response: The document is empty. at 0" after 0 seconds + + Scenario: MiNiFi shuts down when the connection to the MiNiFi C2 server is lost + Given a GenerateFlowFile processor + And C2 is enabled in MiNiFi + And MiNiFi configuration "nifi.c2.connection.timeout" is set to "10 sec" + And a MiNiFi C2 server is started + When the MiNiFi instance starts up + Then the MiNiFi C2 server logs contain the following message: "acknowledged with a state of FULLY_APPLIED(DONE)" in less than 30 seconds + And the Minifi logs do not contain the following message: "C2 connection timeout exceeded, shutting down..." after 15 seconds + When the MiNiFi C2 server is stopped + Then MiNiFi exits with exit code 90 in less than 60 seconds + And the logs of the exited Minifi container contain the following message: "C2 connection timeout exceeded, shutting down..." + + Scenario: MiNiFi shuts down when the connection to the MiNiFi C2 server is not available + Given a GenerateFlowFile processor + And C2 is enabled in MiNiFi + And MiNiFi configuration "nifi.c2.connection.timeout" is set to "10 sec" + When the MiNiFi instance starts up + Then MiNiFi exits with exit code 90 in less than 60 seconds + And the logs of the exited Minifi container contain the following message: "C2 connection timeout exceeded, shutting down..." diff --git a/extensions/standard-processors/tests/features/steps/steps.py b/extensions/standard-processors/tests/features/steps/steps.py index b4469fbe9a..df4b7940b5 100644 --- a/extensions/standard-processors/tests/features/steps/steps.py +++ b/extensions/standard-processors/tests/features/steps/steps.py @@ -16,12 +16,15 @@ # import humanfriendly -from behave import given, step, then +from behave import given, step, then, when from containers.diag_slave_container import DiagSlave from containers.syslog_container import SyslogContainer from containers.tcp_client_container import TcpClientContainer from minifi_behave.containers.minifi_protocol import set_up_ssl_properties -from minifi_behave.core.helpers import wait_for_condition +from minifi_behave.core.helpers import ( + log_due_to_failure, + wait_for_condition, +) from minifi_behave.core.minifi_test_context import MinifiTestContext from minifi_behave.steps import ( checking_steps, # noqa: F401 @@ -126,6 +129,24 @@ def start_minifi_c2_server(context: MinifiTestContext): assert context.containers["minifi-c2-server"].deploy(context) +@when("the MiNiFi C2 server is stopped") +def stop_minifi_c2_server(context: MinifiTestContext): + context.containers["minifi-c2-server"].stop() + + +@then("MiNiFi exits with exit code {exit_code:d} in less than {duration}") +def verify_minifi_exit_code(context: MinifiTestContext, exit_code: int, duration: str): + duration_seconds = humanfriendly.parse_timespan(duration) + minifi_container = context.get_or_create_default_minifi_container() + assert wait_for_condition( + condition=lambda: minifi_container.exited, + timeout_seconds=duration_seconds, + bail_condition=lambda: False, + context=context, + ) + assert minifi_container.exit_code == exit_code or log_due_to_failure(context) + + @then('the MiNiFi C2 server logs contain the following message: "{log_message}" in less than {duration}') def verify_c2_server_logs_contain_message(context: MinifiTestContext, log_message: str, duration: str): duration_seconds = humanfriendly.parse_timespan(duration) diff --git a/libminifi/include/FlowController.h b/libminifi/include/FlowController.h index 00ff3f756d..5177e29579 100644 --- a/libminifi/include/FlowController.h +++ b/libminifi/include/FlowController.h @@ -46,6 +46,7 @@ #include "minifi-cpp/core/state/nodes/MetricsBase.h" #include "core/state/UpdateController.h" #include "c2/C2Agent.h" +#include "RequestShutdown.h" #include "CronDrivenSchedulingAgent.h" #include "EventDrivenSchedulingAgent.h" #include "properties/Configure.h" @@ -69,7 +70,8 @@ class FlowController : public core::controller::ForwardingControllerServiceProvi FlowController(std::shared_ptr provenance_repo, std::shared_ptr flow_file_repo, std::shared_ptr configure, std::shared_ptr flow_configuration, std::shared_ptr content_repo, std::unique_ptr metrics_publisher_store = nullptr, - std::shared_ptr filesystem = std::make_shared(), std::function request_restart = []{}, + std::shared_ptr filesystem = std::make_shared(), + RequestShutdownCallback request_shutdown = [](ShutdownType, std::optional){}, utils::file::AssetManager* asset_manager = {}, core::BulletinStore* bulletin_store = {}); ~FlowController() override; diff --git a/libminifi/include/RequestShutdown.h b/libminifi/include/RequestShutdown.h new file mode 100644 index 0000000000..7451fb6a0a --- /dev/null +++ b/libminifi/include/RequestShutdown.h @@ -0,0 +1,32 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include + +namespace org::apache::nifi::minifi { + +enum class ShutdownType { + RESTART, + SHUTDOWN +}; + +using RequestShutdownCallback = std::function exit_code)>; + +} // namespace org::apache::nifi::minifi diff --git a/libminifi/include/c2/C2Agent.h b/libminifi/include/c2/C2Agent.h index 691d4be577..27c73a6458 100644 --- a/libminifi/include/c2/C2Agent.h +++ b/libminifi/include/c2/C2Agent.h @@ -43,6 +43,8 @@ #include "utils/file/FileSystem.h" #include "C2Utils.h" #include "utils/file/AssetManager.h" +#include "C2ConnectionWatchdog.h" +#include "RequestShutdown.h" namespace org::apache::nifi::minifi::c2 { @@ -62,7 +64,7 @@ class C2Agent : public state::UpdateController { C2Agent(std::shared_ptr configuration, std::weak_ptr node_reporter, std::shared_ptr filesystem, - std::function request_restart, + minifi::RequestShutdownCallback request_shutdown, utils::file::AssetManager* asset_manager); void initialize(core::controller::ControllerServiceProvider *controller, state::Pausable *pause_handler, state::StateMonitor* update_sink); @@ -235,12 +237,14 @@ class C2Agent : public state::UpdateController { const uint64_t C2RESPONSE_POLL_MS = 100; std::atomic restart_needed_ = false; - std::function request_restart_; + minifi::RequestShutdownCallback request_shutdown_; // time point the last time we performed a heartbeat. std::chrono::steady_clock::time_point last_run_; utils::file::AssetManager* asset_manager_; + std::optional c2_connection_timeout_; + std::optional connection_watchdog_; }; } // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/include/c2/C2ConnectionWatchdog.h b/libminifi/include/c2/C2ConnectionWatchdog.h new file mode 100644 index 0000000000..9c51d96ecc --- /dev/null +++ b/libminifi/include/c2/C2ConnectionWatchdog.h @@ -0,0 +1,46 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include + +#include "RequestShutdown.h" +#include "core/logging/LoggerFactory.h" +#include "minifi-cpp/core/logging/Logger.h" + +namespace org::apache::nifi::minifi::c2 { + +class C2ConnectionWatchdog { + public: + C2ConnectionWatchdog(std::chrono::milliseconds connection_timeout, minifi::RequestShutdownCallback request_shutdown); + + void heartBeatHappened() { + last_heartbeat_ = std::chrono::steady_clock::now(); + } + + private: + std::chrono::milliseconds connection_timeout_; + std::atomic last_heartbeat_ = std::chrono::steady_clock::now(); + minifi::RequestShutdownCallback request_shutdown_; + std::shared_ptr logger_ = core::logging::LoggerFactory::getLogger(); + std::jthread watchdog_thread_; +}; + +} // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/include/c2/C2ExitCodes.h b/libminifi/include/c2/C2ExitCodes.h new file mode 100644 index 0000000000..0a0c709176 --- /dev/null +++ b/libminifi/include/c2/C2ExitCodes.h @@ -0,0 +1,25 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include + +namespace org::apache::nifi::minifi::c2 { + +constexpr int32_t C2_CONNECTION_TIMEOUT_EXIT_CODE = 90; + +} // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/src/Configuration.cpp b/libminifi/src/Configuration.cpp index 73c9519c2e..f236ff814f 100644 --- a/libminifi/src/Configuration.cpp +++ b/libminifi/src/Configuration.cpp @@ -97,6 +97,8 @@ const std::unordered_map provenance_repo, std::shared_ptr flow_file_repo, std::shared_ptr configure, std::shared_ptr flow_configuration, std::shared_ptr content_repo, std::unique_ptr metrics_publisher_store, - std::shared_ptr filesystem, std::function request_restart, + std::shared_ptr filesystem, RequestShutdownCallback request_shutdown, utils::file::AssetManager* asset_manager, core::BulletinStore* bulletin_store) : core::controller::ForwardingControllerServiceProvider(core::className()), running_(false), @@ -82,7 +82,7 @@ FlowController::FlowController(std::shared_ptr provenance_repo if (auto publisher = metrics_publisher_store_->getMetricsPublisher(c2::C2_METRICS_PUBLISHER).lock()) { c2_metrics_publisher = std::dynamic_pointer_cast(publisher); } - c2_agent_ = std::make_unique(configuration_, c2_metrics_publisher, std::move(filesystem), std::move(request_restart), asset_manager); + c2_agent_ = std::make_unique(configuration_, c2_metrics_publisher, std::move(filesystem), std::move(request_shutdown), asset_manager); } if (c2::isControllerSocketEnabled(configuration_)) { diff --git a/libminifi/src/c2/C2Agent.cpp b/libminifi/src/c2/C2Agent.cpp index 67ec8b72a6..59e14a54f4 100644 --- a/libminifi/src/c2/C2Agent.cpp +++ b/libminifi/src/c2/C2Agent.cpp @@ -52,10 +52,29 @@ using namespace std::literals::chrono_literals; namespace org::apache::nifi::minifi::c2 { +namespace { +std::optional parseTimeoutPeriod(const std::shared_ptr& configuration, const std::string& config_param, + const std::shared_ptr& logger) { + return configuration->get(config_param) | utils::andThen([&](const auto& timeout_str) -> std::optional { + auto duration = parsing::parseDuration(timeout_str) | utils::toOptional(); + if (!duration) { + logger->log_error("Failed to parse C2 configuration parameter '{}': '{}'", config_param, timeout_str); + } + if (duration && *duration <= 0ms) { + if (*duration < 0ms) { + logger->log_warn("Negative value for C2 configuration parameter '{}': '{}', ignoring it", config_param, timeout_str); + } + return std::nullopt; + } + return duration; + }); +} +} // namespace + C2Agent::C2Agent(std::shared_ptr configuration, std::weak_ptr node_reporter, std::shared_ptr filesystem, - std::function request_restart, + minifi::RequestShutdownCallback request_shutdown, utils::file::AssetManager* asset_manager) : heart_beat_period_(3s), max_c2_responses(5), @@ -63,7 +82,7 @@ C2Agent::C2Agent(std::shared_ptr configuration, node_reporter_(std::move(node_reporter)), filesystem_(std::move(filesystem)), thread_pool_(2, "C2 threadpool"), - request_restart_(std::move(request_restart)), + request_shutdown_(std::move(request_shutdown)), last_run_(std::chrono::steady_clock::now()), asset_manager_(asset_manager) { if (!configuration_->getAgentClass()) { @@ -109,6 +128,9 @@ void C2Agent::start() { if (controller_running_) { return; } + if (c2_connection_timeout_) { + connection_watchdog_.emplace(*c2_connection_timeout_, request_shutdown_); + } task_ids_.clear(); for (const auto& function : functions_) { utils::Identifier uuid = utils::IdGenerator::getIdGenerator()->generate(); @@ -130,6 +152,7 @@ void C2Agent::stop() { thread_pool_.stopTasks(id.to_string()); } thread_pool_.shutdown(); + connection_watchdog_.reset(); logger_->log_info("C2 agent stopped"); } @@ -179,6 +202,17 @@ void C2Agent::configure(const std::shared_ptr &configure, bool reconf } logger_->log_debug("Using {} as the heartbeat period", heart_beat_period_); + c2_connection_timeout_ = parseTimeoutPeriod(configuration_, std::string{Configuration::nifi_c2_connection_timeout}, logger_); + if (c2_connection_timeout_ && *c2_connection_timeout_ <= heart_beat_period_) { + logger_->log_warn( + "C2 connection timeout is less than the heartbeat period: c2 connection timeout={}, heartbeat period={}. Adjusting it to be twice the " + "heartbeat period", + *c2_connection_timeout_, + heart_beat_period_); + c2_connection_timeout_ = heart_beat_period_ * 2; + logger_->log_debug("Using {} as the C2 connection timeout", *c2_connection_timeout_); + } + std::string heartbeat_reporters; if (configure->get(Configuration::nifi_c2_agent_heartbeat_reporter_classes, "c2.agent.heartbeat.reporter.classes", heartbeat_reporters)) { std::vector reporters = utils::string::splitAndTrim(heartbeat_reporters, ","); @@ -234,6 +268,10 @@ void C2Agent::performHeartBeat() { } C2Payload response = protocol_->consumePayload(payload); + if (connection_watchdog_ && response.getStatus().getState() != state::UpdateState::READ_ERROR) { + connection_watchdog_->heartBeatHappened(); + } + enqueue_c2_server_response(std::move(response)); std::lock_guard lock(heartbeat_mutex); @@ -944,7 +982,8 @@ utils::TaskRescheduleInfo C2Agent::produce() { if (restart_needed_ && requests.empty()) { configuration_->commitChanges(); - request_restart_(); + connection_watchdog_.reset(); + request_shutdown_(minifi::ShutdownType::RESTART, std::nullopt); return utils::TaskRescheduleInfo::Done(); } diff --git a/libminifi/src/c2/C2ConnectionWatchdog.cpp b/libminifi/src/c2/C2ConnectionWatchdog.cpp new file mode 100644 index 0000000000..f82acf97be --- /dev/null +++ b/libminifi/src/c2/C2ConnectionWatchdog.cpp @@ -0,0 +1,52 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include "c2/C2ConnectionWatchdog.h" + +#include +#include +#include + +#include "c2/C2ExitCodes.h" + +using namespace std::literals::chrono_literals; + +namespace org::apache::nifi::minifi::c2 { + +C2ConnectionWatchdog::C2ConnectionWatchdog(std::chrono::milliseconds connection_timeout, minifi::RequestShutdownCallback request_shutdown) + : connection_timeout_(connection_timeout), + request_shutdown_(std::move(request_shutdown)) { + if (connection_timeout_ <= 0ms) { + throw std::invalid_argument("C2 connection timeout must be greater than 0"); + } + + watchdog_thread_ = std::jthread([this](const std::stop_token& st) { + std::condition_variable_any cv; + std::mutex sleep_mutex; + std::unique_lock sleep_lock(sleep_mutex); + while (!st.stop_requested()) { + const auto deadline = last_heartbeat_.load() + connection_timeout_; + if (std::chrono::steady_clock::now() >= deadline) { + logger_->log_error("C2 connection timeout exceeded, shutting down..."); + request_shutdown_(minifi::ShutdownType::SHUTDOWN, C2_CONNECTION_TIMEOUT_EXIT_CODE); + break; + } + cv.wait_until(sleep_lock, st, deadline, [] { return false; }); + } + }); +} + +} // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/test/integration/C2ConnectionTimeoutTest.cpp b/libminifi/test/integration/C2ConnectionTimeoutTest.cpp new file mode 100644 index 0000000000..2ad6447dc9 --- /dev/null +++ b/libminifi/test/integration/C2ConnectionTimeoutTest.cpp @@ -0,0 +1,93 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "c2/C2Agent.h" +#include "c2/C2ExitCodes.h" +#include "integration/HTTPIntegrationBase.h" +#include "integration/ServerAwareHandler.h" +#include "properties/Configuration.h" +#include "unit/Catch.h" +#include "unit/TestBase.h" +#include "unit/TestUtils.h" + +using namespace std::literals::chrono_literals; + +namespace org::apache::nifi::minifi::test { + +class FailingHeartbeatHandler : public ServerAwareHandler { + public: + bool handlePost(CivetServer* /*server*/, struct mg_connection* conn) override { + ++heartbeat_count_; + if (fail_) { + mg_printf(conn, "HTTP/1.1 500 Internal Server Error\r\nContent-Length: 0\r\nConnection: close\r\n\r\n"); + } else { + static const std::string response = R"({"operation": "heartbeat", "requested_operations": []})"; + mg_printf(conn, "HTTP/1.1 200 OK\r\nContent-Type: application/json\r\nContent-Length: %lu\r\nConnection: close\r\n\r\n", response.length()); + mg_printf(conn, "%s", response.c_str()); + } + return true; + } + + std::atomic fail_{false}; + std::atomic heartbeat_count_{0}; +}; + +class VerifyC2ConnectionTimeout : public VerifyC2Base { + public: + VerifyC2ConnectionTimeout(const std::filesystem::path& test_file_path, FailingHeartbeatHandler& handler) + : VerifyC2Base(test_file_path), + handler_(handler) { + } + + void testSetup() override { + LogTestController::getInstance().setDebug(); + VerifyC2Base::testSetup(); + } + + void configureC2() override { + VerifyC2Base::configureC2(); + configuration->set(minifi::Configuration::nifi_c2_agent_heartbeat_period, "100 ms"); + configuration->set(minifi::Configuration::nifi_c2_connection_timeout, "1 sec"); + } + + void runAssertions() override { + // successful heartbeats keep the agent alive well past the connection timeout + std::this_thread::sleep_for(3s); + REQUIRE(handler_.heartbeat_count_ > 0); + REQUIRE(shutdown_requested_count_ == 0); + + handler_.fail_ = true; + REQUIRE(utils::verifyEventHappenedInPollTime(5s, [&] { return shutdown_requested_count_ > 0; })); + CHECK(shutdown_requested_count_ == 1); + CHECK(last_shutdown_exit_code_ == minifi::c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); + CHECK(utils::verifyLogLinePresenceInPollTime(1s, "C2 connection timeout exceeded, shutting down...")); + } + + private: + FailingHeartbeatHandler& handler_; +}; + +TEST_CASE("C2 agent requests shutdown when the C2 server stops responding successfully", "[c2test]") { + FailingHeartbeatHandler handler; + VerifyC2ConnectionTimeout harness(std::filesystem::path(TEST_RESOURCES) / "TestEmpty.yml", handler); + harness.setUrl("http://localhost:0/heartbeat", &handler); + harness.run(); +} + +} // namespace org::apache::nifi::minifi::test diff --git a/libminifi/test/libtest/integration/IntegrationBase.cpp b/libminifi/test/libtest/integration/IntegrationBase.cpp index 8a767546a6..7d8d399873 100644 --- a/libminifi/test/libtest/integration/IntegrationBase.cpp +++ b/libminifi/test/libtest/integration/IntegrationBase.cpp @@ -129,7 +129,13 @@ void IntegrationBase::run() { std::shared_ptr pg(flow_config->getRoot()); queryRootProcessGroup(pg); - const auto request_restart = [&, this] { + const auto request_shutdown = [&, this] (ShutdownType shutdown_type, std::optional exit_code) { + if (shutdown_type == ShutdownType::SHUTDOWN) { + last_shutdown_exit_code_ = exit_code.value_or(0); + ++shutdown_requested_count_; + running = false; + return; + } ++restart_requested_count_; running = true; }; @@ -138,7 +144,7 @@ void IntegrationBase::run() { asset_manager_ = std::make_unique(*configuration); auto metrics_publisher_store = std::make_unique(configuration, repo_metric_sources, flow_config, asset_manager_.get(), bulletin_store_.get()); flowController_ = std::make_unique(test_repo, test_flow_repo, configuration, - std::move(flow_config), content_repo, std::move(metrics_publisher_store), filesystem, request_restart, asset_manager_.get()); + std::move(flow_config), content_repo, std::move(metrics_publisher_store), filesystem, request_shutdown, asset_manager_.get()); flowController_->load(); updateProperties(*flowController_); flowController_->start(); diff --git a/libminifi/test/libtest/integration/IntegrationBase.h b/libminifi/test/libtest/integration/IntegrationBase.h index 9f406d57ad..16c5286e9e 100644 --- a/libminifi/test/libtest/integration/IntegrationBase.h +++ b/libminifi/test/libtest/integration/IntegrationBase.h @@ -59,6 +59,8 @@ class IntegrationBase { key_dir{std::move(other.key_dir)}, state_dir{std::move(other.state_dir)}, restart_requested_count_{other.restart_requested_count_.load()}, + shutdown_requested_count_{other.shutdown_requested_count_.load()}, + last_shutdown_exit_code_{other.last_shutdown_exit_code_.load()}, home_path_{std::move(other.home_path_)}, flow_config_path_{std::move(other.flow_config_path_)} {} @@ -74,6 +76,8 @@ class IntegrationBase { key_dir = std::move(other.key_dir); state_dir = std::move(other.state_dir); restart_requested_count_ = other.restart_requested_count_.load(); + shutdown_requested_count_ = other.shutdown_requested_count_.load(); + last_shutdown_exit_code_ = other.last_shutdown_exit_code_.load(); home_path_ = std::move(other.home_path_); flow_config_path_ = std::move(other.flow_config_path_); return *this; @@ -141,6 +145,8 @@ class IntegrationBase { std::filesystem::path key_dir; std::filesystem::path state_dir; std::atomic restart_requested_count_{0}; + std::atomic shutdown_requested_count_{0}; + std::atomic last_shutdown_exit_code_{0}; std::optional home_path_; FlowConfigPath flow_config_path_; }; diff --git a/libminifi/test/unit/C2AgentWatchdogConfigurationTests.cpp b/libminifi/test/unit/C2AgentWatchdogConfigurationTests.cpp new file mode 100644 index 0000000000..36af760b2e --- /dev/null +++ b/libminifi/test/unit/C2AgentWatchdogConfigurationTests.cpp @@ -0,0 +1,123 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include +#include +#include + +#include "C2ShutdownRecorder.h" +#include "c2/C2Agent.h" +#include "c2/C2ExitCodes.h" +#include "properties/Configure.h" +#include "unit/Catch.h" +#include "unit/TestBase.h" +#include "unit/TestUtils.h" +#include "utils/file/FileSystem.h" + +using namespace std::literals::chrono_literals; + +namespace org::apache::nifi::minifi::test { + +class C2AgentWatchdogConfigurationTestFixture { + public: + C2AgentWatchdogConfigurationTestFixture() { + LogTestController::getInstance().reset(); + LogTestController::getInstance().setTrace(); + configuration_->set(Configuration::nifi_c2_agent_heartbeat_period, "100 ms"); + } + + C2AgentWatchdogConfigurationTestFixture(C2AgentWatchdogConfigurationTestFixture&&) = delete; + C2AgentWatchdogConfigurationTestFixture(const C2AgentWatchdogConfigurationTestFixture&) = delete; + C2AgentWatchdogConfigurationTestFixture& operator=(C2AgentWatchdogConfigurationTestFixture&&) = delete; + C2AgentWatchdogConfigurationTestFixture& operator=(const C2AgentWatchdogConfigurationTestFixture&) = delete; + + ~C2AgentWatchdogConfigurationTestFixture() { + if (agent_) { + agent_->stop(); + } + LogTestController::getInstance().reset(); + } + + void startAgent() { + agent_ = std::make_unique(configuration_, + std::weak_ptr{}, + std::make_shared(), + recorder_.callback(), + nullptr); + agent_->initialize(nullptr, nullptr, nullptr); + agent_->start(); + } + + protected: + std::shared_ptr configuration_ = std::make_shared(); + C2ShutdownRecorder recorder_; + std::unique_ptr agent_; +}; + +TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "No shutdown is requested when the C2 connection timeout is not set", "[c2agentwatchdog]") { + startAgent(); + REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 1s)); + std::this_thread::sleep_for(1s); + REQUIRE(recorder_.count() == 0); +} + +TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "C2 connection timeout of 0 disables the watchdog", "[c2agentwatchdog]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "0 sec"); + startAgent(); + REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 1s)); + std::this_thread::sleep_for(1s); + REQUIRE(recorder_.count() == 0); + REQUIRE_FALSE(LogTestController::getInstance().contains("Failed to parse C2 configuration parameter", 0ms)); +} + +TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "Invalid C2 connection timeout disables the watchdog", "[c2agentwatchdog]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "abc"); + startAgent(); + REQUIRE(LogTestController::getInstance().contains("Failed to parse C2 configuration parameter 'nifi.c2.connection.timeout': 'abc'", 1s)); + std::this_thread::sleep_for(1s); + REQUIRE(recorder_.count() == 0); +} + +TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "C2 connection timeout shorter than the heartbeat period is increased", + "[c2agentwatchdog]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "50 ms"); + startAgent(); + REQUIRE(LogTestController::getInstance().contains("C2 connection timeout is less than the heartbeat period", 1s)); + REQUIRE(LogTestController::getInstance().contains("Using 200ms as the C2 connection timeout", 1s)); +} + +TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "Failed heartbeats do not prevent the C2 connection timeout", "[c2agentwatchdog]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "300 ms"); + startAgent(); + REQUIRE(utils::verifyEventHappenedInPollTime(3s, [&] { return recorder_.count() > 0; }, 10ms)); + REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 0ms)); + REQUIRE(recorder_.count() == 1); + const auto [shutdown_type, exit_code] = recorder_.firstCall(); + CHECK(shutdown_type == ShutdownType::SHUTDOWN); + REQUIRE(exit_code.has_value()); + CHECK(*exit_code == c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); +} + +TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "Stopping the C2 agent stops the watchdog", "[c2agentwatchdog]") { + configuration_->set(Configuration::nifi_c2_connection_timeout, "300 ms"); + startAgent(); + agent_->stop(); + std::this_thread::sleep_for(600ms); + REQUIRE(recorder_.count() == 0); +} + +} // namespace org::apache::nifi::minifi::test diff --git a/libminifi/test/unit/C2ConnectionWatchdogTests.cpp b/libminifi/test/unit/C2ConnectionWatchdogTests.cpp new file mode 100644 index 0000000000..f1b300a874 --- /dev/null +++ b/libminifi/test/unit/C2ConnectionWatchdogTests.cpp @@ -0,0 +1,103 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#include +#include + +#include "C2ShutdownRecorder.h" +#include "c2/C2ConnectionWatchdog.h" +#include "c2/C2ExitCodes.h" +#include "unit/Catch.h" +#include "unit/TestBase.h" +#include "unit/TestUtils.h" + +using namespace std::literals::chrono_literals; + +namespace org::apache::nifi::minifi::test { + +namespace { +void verifyTimeoutShutdownRequested(C2ShutdownRecorder& recorder) { + REQUIRE(utils::verifyEventHappenedInPollTime(2s, [&] { return recorder.count() > 0; }, 10ms)); + REQUIRE(recorder.count() == 1); + const auto [shutdown_type, exit_code] = recorder.firstCall(); + CHECK(shutdown_type == ShutdownType::SHUTDOWN); + REQUIRE(exit_code.has_value()); + CHECK(*exit_code == c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); +} +} // namespace + +TEST_CASE("Test invalid values", "[c2watchdog]") { + LogTestController::getInstance().reset(); + SECTION("Zero timeout value") { + REQUIRE_THROWS_WITH(c2::C2ConnectionWatchdog(0ms, [](ShutdownType, std::optional) {}), + "C2 connection timeout must be greater than 0"); + } + + SECTION("Negative timeout value") { + REQUIRE_THROWS_WITH(c2::C2ConnectionWatchdog(-1ms, [](ShutdownType, std::optional) {}), + "C2 connection timeout must be greater than 0"); + } +} + +TEST_CASE("Test shutdown when no C2 connection is present at all", "[c2watchdog]") { + LogTestController::getInstance().reset(); + C2ShutdownRecorder recorder; + c2::C2ConnectionWatchdog watchdog(400ms, recorder.callback()); + verifyTimeoutShutdownRequested(recorder); + REQUIRE(LogTestController::getInstance().contains("C2 connection timeout exceeded, shutting down...", 1s)); +} + +TEST_CASE("Test shutdown when C2 connection is lost", "[c2watchdog]") { + LogTestController::getInstance().reset(); + C2ShutdownRecorder recorder; + c2::C2ConnectionWatchdog watchdog(400ms, recorder.callback()); + for (int i = 0; i < 10; ++i) { + watchdog.heartBeatHappened(); + REQUIRE(recorder.count() == 0); + std::this_thread::sleep_for(50ms); + } + verifyTimeoutShutdownRequested(recorder); + REQUIRE(LogTestController::getInstance().contains("C2 connection timeout exceeded, shutting down...", 1s)); +} + +TEST_CASE("Shutdown is only requested once", "[c2watchdog]") { + LogTestController::getInstance().reset(); + C2ShutdownRecorder recorder; + c2::C2ConnectionWatchdog watchdog(100ms, recorder.callback()); + verifyTimeoutShutdownRequested(recorder); + std::this_thread::sleep_for(300ms); + REQUIRE(recorder.count() == 1); +} + +TEST_CASE("Watchdog is destroyed promptly even with a long timeout", "[c2watchdog]") { + LogTestController::getInstance().reset(); + C2ShutdownRecorder recorder; + auto watchdog = std::make_unique(1h, recorder.callback()); + const auto destruction_start = std::chrono::steady_clock::now(); + watchdog.reset(); + REQUIRE(std::chrono::steady_clock::now() - destruction_start < 1s); + REQUIRE(recorder.count() == 0); +} + +TEST_CASE("No shutdown is requested after the watchdog is destroyed", "[c2watchdog]") { + LogTestController::getInstance().reset(); + C2ShutdownRecorder recorder; + { c2::C2ConnectionWatchdog watchdog(200ms, recorder.callback()); } + std::this_thread::sleep_for(500ms); + REQUIRE(recorder.count() == 0); +} + +} // namespace org::apache::nifi::minifi::test diff --git a/libminifi/test/unit/C2ShutdownRecorder.h b/libminifi/test/unit/C2ShutdownRecorder.h new file mode 100644 index 0000000000..da50c878df --- /dev/null +++ b/libminifi/test/unit/C2ShutdownRecorder.h @@ -0,0 +1,53 @@ +/** + * Licensed to the Apache Software Foundation (ASF) under one or more + * contributor license agreements. See the NOTICE file distributed with + * this work for additional information regarding copyright ownership. + * The ASF licenses this file to You under the Apache License, Version 2.0 + * (the "License"); you may not use this file except in compliance with + * the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include +#include +#include +#include +#include + +#include "RequestShutdown.h" + +namespace org::apache::nifi::minifi::test { + +class C2ShutdownRecorder { + public: + RequestShutdownCallback callback() { + return [this](ShutdownType shutdown_type, std::optional exit_code) { + std::lock_guard lock(mutex_); + calls_.emplace_back(shutdown_type, exit_code); + }; + } + + size_t count() { + std::lock_guard lock(mutex_); + return calls_.size(); + } + + std::pair> firstCall() { + std::lock_guard lock(mutex_); + return calls_.at(0); + } + + private: + std::mutex mutex_; + std::vector>> calls_; +}; + +} // namespace org::apache::nifi::minifi::test diff --git a/minifi-api/include/minifi-cpp/properties/Configuration.h b/minifi-api/include/minifi-cpp/properties/Configuration.h index 1ec60afa51..9272c2152a 100644 --- a/minifi-api/include/minifi-cpp/properties/Configuration.h +++ b/minifi-api/include/minifi-cpp/properties/Configuration.h @@ -127,6 +127,7 @@ class Configuration : public virtual Properties { static constexpr const char *nifi_c2_rest_request_encoding = "nifi.c2.rest.request.encoding"; static constexpr const char *nifi_c2_flow_info_processor_bulletin_limit = "nifi.c2.flow.info.processor.bulletin.limit"; static constexpr const char *nifi_c2_asset_download_timeout = "nifi.c2.asset.download.timeout"; + static constexpr const char *nifi_c2_connection_timeout = "nifi.c2.connection.timeout"; // state management options static constexpr const char *nifi_state_storage_local = "nifi.state.storage.local"; diff --git a/minifi_main/MiNiFiMain.cpp b/minifi_main/MiNiFiMain.cpp index fcc2ba9d6a..8d039084a2 100644 --- a/minifi_main/MiNiFiMain.cpp +++ b/minifi_main/MiNiFiMain.cpp @@ -290,7 +290,17 @@ int main(int argc, char **argv) { process_running.test_and_set(); std::atomic restart_token{false}; - const auto request_restart = [&] { + std::atomic exit_code{0}; + const auto request_shutdown = [&] (minifi::ShutdownType shutdown_type, std::optional shutdown_exit_code) { + if (shutdown_type == minifi::ShutdownType::SHUTDOWN) { + restart_token = false; + exit_code = shutdown_exit_code.value_or(0); + flow_controller_running.clear(); + flow_controller_running.notify_all(); + logger->log_info("Initiating shutdown..."); + return; + } + if (!restart_token.exchange(true)) { // only trigger if a restart is not already in progress (the flag was unset before the exchange) flow_controller_running.clear(); @@ -408,7 +418,7 @@ int main(int argc, char **argv) { auto metrics_publisher_store = std::make_unique(configure, repo_metric_sources, flow_configuration, asset_manager.get(), bulletin_store.get()); const auto controller = std::make_unique( prov_repo, flow_repo, configure, std::move(flow_configuration), content_repo, - std::move(metrics_publisher_store), filesystem, request_restart, asset_manager.get(), bulletin_store.get()); + std::move(metrics_publisher_store), filesystem, request_shutdown, asset_manager.get(), bulletin_store.get()); const bool disk_space_watchdog_enable = configure->get(minifi::Configure::minifi_disk_space_watchdog_enable) | utils::andThen(utils::string::toBool) @@ -505,5 +515,5 @@ int main(int argc, char **argv) { process_running.clear(); process_running.notify_all(); logger->log_info("MiNiFi exit"); - return 0; + return exit_code; } diff --git a/minifi_main/MiNiFiWindowsService.cpp b/minifi_main/MiNiFiWindowsService.cpp index dd5e446f3f..e3be008946 100644 --- a/minifi_main/MiNiFiWindowsService.cpp +++ b/minifi_main/MiNiFiWindowsService.cpp @@ -26,6 +26,7 @@ #include #include "MainHelper.h" +#include "c2/C2ExitCodes.h" namespace minifi = org::apache::nifi::minifi; @@ -176,6 +177,8 @@ void RunAsServiceIfNeeded() { return; } + DWORD serviceSpecificExitCode = 0; + do { Log()->log_debug("Start exe path {}", filePath); @@ -194,6 +197,11 @@ void RunAsServiceIfNeeded() { Log()->log_info("{} started", filePath); auto res = WaitForSingleObject(processInformation.hProcess, INFINITE); + DWORD exeExitCode = 0; + const bool exeExitCodeRetrieved = GetExitCodeProcess(processInformation.hProcess, &exeExitCode); + if (!exeExitCodeRetrieved) { + Log()->log_error("!GetExitCodeProcess lastError {:#x}", GetLastError()); + } CloseHandle(processInformation.hProcess); CloseHandle(processInformation.hThread); @@ -203,6 +211,12 @@ void RunAsServiceIfNeeded() { Log()->log_error("!WaitForSingleObject hProcess return {}", res); } + if (exeExitCodeRetrieved && exeExitCode == static_cast(minifi::c2::C2_CONNECTION_TIMEOUT_EXIT_CODE)) { + Log()->log_info("Exe exited with exit code {} due to C2 connection timeout, exe won't be restarted", exeExitCode); + serviceSpecificExitCode = exeExitCode; + break; + } + Log()->log_info("Sleep {} sec before restarting exe", WAIT_TIME_EXE_RESTART/1000); res = WaitForSingleObject(s_hEvent, WAIT_TIME_EXE_RESTART); @@ -220,7 +234,8 @@ void RunAsServiceIfNeeded() { s_serviceStatus.dwControlsAccepted = 0; s_serviceStatus.dwCurrentState = SERVICE_STOPPED; - s_serviceStatus.dwWin32ExitCode = 0; + s_serviceStatus.dwWin32ExitCode = serviceSpecificExitCode == 0 ? 0 : ERROR_SERVICE_SPECIFIC_ERROR; + s_serviceStatus.dwServiceSpecificExitCode = serviceSpecificExitCode; if (!SetServiceStatus(s_statusHandle, &s_serviceStatus)) { Log()->log_error("!SetServiceStatus SERVICE_STOPPED lastError {:#x}", GetLastError()); From e376af0e6388dd80cd87d3f029dc0bb023ce639b Mon Sep 17 00:00:00 2001 From: Gabor Gyimesi Date: Tue, 27 Oct 2026 13:09:54 +0100 Subject: [PATCH 2/6] Review update --- libminifi/include/c2/C2ConnectionWatchdog.h | 4 ++++ libminifi/src/c2/C2Agent.cpp | 3 +++ libminifi/src/c2/C2ConnectionWatchdog.cpp | 3 ++- libminifi/test/unit/C2ConnectionWatchdogTests.cpp | 10 ++++++++++ minifi_main/MiNiFiMain.cpp | 8 ++++++++ 5 files changed, 27 insertions(+), 1 deletion(-) diff --git a/libminifi/include/c2/C2ConnectionWatchdog.h b/libminifi/include/c2/C2ConnectionWatchdog.h index 9c51d96ecc..3880416711 100644 --- a/libminifi/include/c2/C2ConnectionWatchdog.h +++ b/libminifi/include/c2/C2ConnectionWatchdog.h @@ -35,6 +35,10 @@ class C2ConnectionWatchdog { last_heartbeat_ = std::chrono::steady_clock::now(); } + void stop() { + watchdog_thread_.request_stop(); + } + private: std::chrono::milliseconds connection_timeout_; std::atomic last_heartbeat_ = std::chrono::steady_clock::now(); diff --git a/libminifi/src/c2/C2Agent.cpp b/libminifi/src/c2/C2Agent.cpp index 59e14a54f4..dda24fd419 100644 --- a/libminifi/src/c2/C2Agent.cpp +++ b/libminifi/src/c2/C2Agent.cpp @@ -148,6 +148,9 @@ void C2Agent::stop() { return; } controller_running_ = false; + if (connection_watchdog_) { + connection_watchdog_->stop(); + } for (const auto& id : task_ids_) { thread_pool_.stopTasks(id.to_string()); } diff --git a/libminifi/src/c2/C2ConnectionWatchdog.cpp b/libminifi/src/c2/C2ConnectionWatchdog.cpp index f82acf97be..652eba911f 100644 --- a/libminifi/src/c2/C2ConnectionWatchdog.cpp +++ b/libminifi/src/c2/C2ConnectionWatchdog.cpp @@ -38,8 +38,9 @@ C2ConnectionWatchdog::C2ConnectionWatchdog(std::chrono::milliseconds connection_ std::mutex sleep_mutex; std::unique_lock sleep_lock(sleep_mutex); while (!st.stop_requested()) { + const auto now = std::chrono::steady_clock::now(); const auto deadline = last_heartbeat_.load() + connection_timeout_; - if (std::chrono::steady_clock::now() >= deadline) { + if (now >= deadline) { logger_->log_error("C2 connection timeout exceeded, shutting down..."); request_shutdown_(minifi::ShutdownType::SHUTDOWN, C2_CONNECTION_TIMEOUT_EXIT_CODE); break; diff --git a/libminifi/test/unit/C2ConnectionWatchdogTests.cpp b/libminifi/test/unit/C2ConnectionWatchdogTests.cpp index f1b300a874..c79bd7313d 100644 --- a/libminifi/test/unit/C2ConnectionWatchdogTests.cpp +++ b/libminifi/test/unit/C2ConnectionWatchdogTests.cpp @@ -100,4 +100,14 @@ TEST_CASE("No shutdown is requested after the watchdog is destroyed", "[c2watchd REQUIRE(recorder.count() == 0); } +TEST_CASE("No shutdown is requested after the watchdog is stopped", "[c2watchdog]") { + LogTestController::getInstance().reset(); + C2ShutdownRecorder recorder; + c2::C2ConnectionWatchdog watchdog(200ms, recorder.callback()); + watchdog.stop(); + std::this_thread::sleep_for(500ms); + watchdog.heartBeatHappened(); + REQUIRE(recorder.count() == 0); +} + } // namespace org::apache::nifi::minifi::test diff --git a/minifi_main/MiNiFiMain.cpp b/minifi_main/MiNiFiMain.cpp index 8d039084a2..5349f6dcbc 100644 --- a/minifi_main/MiNiFiMain.cpp +++ b/minifi_main/MiNiFiMain.cpp @@ -43,6 +43,7 @@ #include #include #include +#include #include #include "AgentDocs.h" @@ -291,8 +292,15 @@ int main(int argc, char **argv) { std::atomic restart_token{false}; std::atomic exit_code{0}; + std::mutex shutdown_request_mutex; + bool shutdown_requested = false; const auto request_shutdown = [&] (minifi::ShutdownType shutdown_type, std::optional shutdown_exit_code) { + std::lock_guard lock(shutdown_request_mutex); + if (shutdown_requested) { + return; + } if (shutdown_type == minifi::ShutdownType::SHUTDOWN) { + shutdown_requested = true; restart_token = false; exit_code = shutdown_exit_code.value_or(0); flow_controller_running.clear(); From 7437475fda6b7fbec67ce184398bae103ab73b56 Mon Sep 17 00:00:00 2001 From: Gabor Gyimesi Date: Tue, 27 Oct 2026 14:22:28 +0100 Subject: [PATCH 3/6] Remove separate watchdog thread --- libminifi/include/c2/C2Agent.h | 4 +- libminifi/include/c2/C2ConnectionWatchdog.h | 50 -------- libminifi/src/c2/C2Agent.cpp | 20 ++-- libminifi/src/c2/C2ConnectionWatchdog.cpp | 53 -------- ....cpp => C2AgentConnectionTimeoutTests.cpp} | 60 +++++----- .../test/unit/C2ConnectionWatchdogTests.cpp | 113 ------------------ libminifi/test/unit/C2ShutdownRecorder.h | 53 -------- 7 files changed, 44 insertions(+), 309 deletions(-) delete mode 100644 libminifi/include/c2/C2ConnectionWatchdog.h delete mode 100644 libminifi/src/c2/C2ConnectionWatchdog.cpp rename libminifi/test/unit/{C2AgentWatchdogConfigurationTests.cpp => C2AgentConnectionTimeoutTests.cpp} (60%) delete mode 100644 libminifi/test/unit/C2ConnectionWatchdogTests.cpp delete mode 100644 libminifi/test/unit/C2ShutdownRecorder.h diff --git a/libminifi/include/c2/C2Agent.h b/libminifi/include/c2/C2Agent.h index 27c73a6458..1267d68cd4 100644 --- a/libminifi/include/c2/C2Agent.h +++ b/libminifi/include/c2/C2Agent.h @@ -18,6 +18,7 @@ #pragma once +#include #include #include #include @@ -43,7 +44,6 @@ #include "utils/file/FileSystem.h" #include "C2Utils.h" #include "utils/file/AssetManager.h" -#include "C2ConnectionWatchdog.h" #include "RequestShutdown.h" namespace org::apache::nifi::minifi::c2 { @@ -244,7 +244,7 @@ class C2Agent : public state::UpdateController { utils::file::AssetManager* asset_manager_; std::optional c2_connection_timeout_; - std::optional connection_watchdog_; + std::chrono::steady_clock::time_point last_successful_heartbeat_{}; }; } // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/include/c2/C2ConnectionWatchdog.h b/libminifi/include/c2/C2ConnectionWatchdog.h deleted file mode 100644 index 3880416711..0000000000 --- a/libminifi/include/c2/C2ConnectionWatchdog.h +++ /dev/null @@ -1,50 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include -#include - -#include "RequestShutdown.h" -#include "core/logging/LoggerFactory.h" -#include "minifi-cpp/core/logging/Logger.h" - -namespace org::apache::nifi::minifi::c2 { - -class C2ConnectionWatchdog { - public: - C2ConnectionWatchdog(std::chrono::milliseconds connection_timeout, minifi::RequestShutdownCallback request_shutdown); - - void heartBeatHappened() { - last_heartbeat_ = std::chrono::steady_clock::now(); - } - - void stop() { - watchdog_thread_.request_stop(); - } - - private: - std::chrono::milliseconds connection_timeout_; - std::atomic last_heartbeat_ = std::chrono::steady_clock::now(); - minifi::RequestShutdownCallback request_shutdown_; - std::shared_ptr logger_ = core::logging::LoggerFactory::getLogger(); - std::jthread watchdog_thread_; -}; - -} // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/src/c2/C2Agent.cpp b/libminifi/src/c2/C2Agent.cpp index dda24fd419..67c48259ba 100644 --- a/libminifi/src/c2/C2Agent.cpp +++ b/libminifi/src/c2/C2Agent.cpp @@ -43,6 +43,7 @@ #include "minifi-cpp/io/ArchiveStream.h" #include "io/StreamPipe.h" #include "utils/Id.h" +#include "c2/C2ExitCodes.h" #include "c2/C2Utils.h" #include "c2/protocols/RESTSender.h" #include "rapidjson/error/en.h" @@ -128,9 +129,7 @@ void C2Agent::start() { if (controller_running_) { return; } - if (c2_connection_timeout_) { - connection_watchdog_.emplace(*c2_connection_timeout_, request_shutdown_); - } + last_successful_heartbeat_ = std::chrono::steady_clock::now(); task_ids_.clear(); for (const auto& function : functions_) { utils::Identifier uuid = utils::IdGenerator::getIdGenerator()->generate(); @@ -148,14 +147,10 @@ void C2Agent::stop() { return; } controller_running_ = false; - if (connection_watchdog_) { - connection_watchdog_->stop(); - } for (const auto& id : task_ids_) { thread_pool_.stopTasks(id.to_string()); } thread_pool_.shutdown(); - connection_watchdog_.reset(); logger_->log_info("C2 agent stopped"); } @@ -271,8 +266,8 @@ void C2Agent::performHeartBeat() { } C2Payload response = protocol_->consumePayload(payload); - if (connection_watchdog_ && response.getStatus().getState() != state::UpdateState::READ_ERROR) { - connection_watchdog_->heartBeatHappened(); + if (response.getStatus().getState() != state::UpdateState::READ_ERROR) { + last_successful_heartbeat_ = std::chrono::steady_clock::now(); } enqueue_c2_server_response(std::move(response)); @@ -985,7 +980,6 @@ utils::TaskRescheduleInfo C2Agent::produce() { if (restart_needed_ && requests.empty()) { configuration_->commitChanges(); - connection_watchdog_.reset(); request_shutdown_(minifi::ShutdownType::RESTART, std::nullopt); return utils::TaskRescheduleInfo::Done(); } @@ -999,6 +993,12 @@ utils::TaskRescheduleInfo C2Agent::produce() { catch (...) { logger_->log_error("Unknown exception occurred while performing heartbeat, type: {}", getCurrentExceptionTypeName()); } + + if (c2_connection_timeout_ && controller_running_ && std::chrono::steady_clock::now() - last_successful_heartbeat_ > *c2_connection_timeout_) { + logger_->log_error("C2 connection timeout exceeded, shutting down..."); + request_shutdown_(minifi::ShutdownType::SHUTDOWN, C2_CONNECTION_TIMEOUT_EXIT_CODE); + return utils::TaskRescheduleInfo::Done(); + } } checkTriggers(); diff --git a/libminifi/src/c2/C2ConnectionWatchdog.cpp b/libminifi/src/c2/C2ConnectionWatchdog.cpp deleted file mode 100644 index 652eba911f..0000000000 --- a/libminifi/src/c2/C2ConnectionWatchdog.cpp +++ /dev/null @@ -1,53 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include "c2/C2ConnectionWatchdog.h" - -#include -#include -#include - -#include "c2/C2ExitCodes.h" - -using namespace std::literals::chrono_literals; - -namespace org::apache::nifi::minifi::c2 { - -C2ConnectionWatchdog::C2ConnectionWatchdog(std::chrono::milliseconds connection_timeout, minifi::RequestShutdownCallback request_shutdown) - : connection_timeout_(connection_timeout), - request_shutdown_(std::move(request_shutdown)) { - if (connection_timeout_ <= 0ms) { - throw std::invalid_argument("C2 connection timeout must be greater than 0"); - } - - watchdog_thread_ = std::jthread([this](const std::stop_token& st) { - std::condition_variable_any cv; - std::mutex sleep_mutex; - std::unique_lock sleep_lock(sleep_mutex); - while (!st.stop_requested()) { - const auto now = std::chrono::steady_clock::now(); - const auto deadline = last_heartbeat_.load() + connection_timeout_; - if (now >= deadline) { - logger_->log_error("C2 connection timeout exceeded, shutting down..."); - request_shutdown_(minifi::ShutdownType::SHUTDOWN, C2_CONNECTION_TIMEOUT_EXIT_CODE); - break; - } - cv.wait_until(sleep_lock, st, deadline, [] { return false; }); - } - }); -} - -} // namespace org::apache::nifi::minifi::c2 diff --git a/libminifi/test/unit/C2AgentWatchdogConfigurationTests.cpp b/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp similarity index 60% rename from libminifi/test/unit/C2AgentWatchdogConfigurationTests.cpp rename to libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp index 36af760b2e..3399994513 100644 --- a/libminifi/test/unit/C2AgentWatchdogConfigurationTests.cpp +++ b/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp @@ -14,12 +14,11 @@ * See the License for the specific language governing permissions and * limitations under the License. */ +#include #include #include #include -#include -#include "C2ShutdownRecorder.h" #include "c2/C2Agent.h" #include "c2/C2ExitCodes.h" #include "properties/Configure.h" @@ -32,20 +31,20 @@ using namespace std::literals::chrono_literals; namespace org::apache::nifi::minifi::test { -class C2AgentWatchdogConfigurationTestFixture { +class C2AgentConnectionTimeoutTestFixture { public: - C2AgentWatchdogConfigurationTestFixture() { + C2AgentConnectionTimeoutTestFixture() { LogTestController::getInstance().reset(); LogTestController::getInstance().setTrace(); configuration_->set(Configuration::nifi_c2_agent_heartbeat_period, "100 ms"); } - C2AgentWatchdogConfigurationTestFixture(C2AgentWatchdogConfigurationTestFixture&&) = delete; - C2AgentWatchdogConfigurationTestFixture(const C2AgentWatchdogConfigurationTestFixture&) = delete; - C2AgentWatchdogConfigurationTestFixture& operator=(C2AgentWatchdogConfigurationTestFixture&&) = delete; - C2AgentWatchdogConfigurationTestFixture& operator=(const C2AgentWatchdogConfigurationTestFixture&) = delete; + C2AgentConnectionTimeoutTestFixture(C2AgentConnectionTimeoutTestFixture&&) = delete; + C2AgentConnectionTimeoutTestFixture(const C2AgentConnectionTimeoutTestFixture&) = delete; + C2AgentConnectionTimeoutTestFixture& operator=(C2AgentConnectionTimeoutTestFixture&&) = delete; + C2AgentConnectionTimeoutTestFixture& operator=(const C2AgentConnectionTimeoutTestFixture&) = delete; - ~C2AgentWatchdogConfigurationTestFixture() { + ~C2AgentConnectionTimeoutTestFixture() { if (agent_) { agent_->stop(); } @@ -56,7 +55,11 @@ class C2AgentWatchdogConfigurationTestFixture { agent_ = std::make_unique(configuration_, std::weak_ptr{}, std::make_shared(), - recorder_.callback(), + [this](ShutdownType shutdown_type, std::optional exit_code) { + last_shutdown_type_ = shutdown_type; + last_exit_code_ = exit_code.value_or(0); + ++shutdown_request_count_; + }, nullptr); agent_->initialize(nullptr, nullptr, nullptr); agent_->start(); @@ -64,60 +67,61 @@ class C2AgentWatchdogConfigurationTestFixture { protected: std::shared_ptr configuration_ = std::make_shared(); - C2ShutdownRecorder recorder_; + std::atomic shutdown_request_count_{0}; + std::atomic last_shutdown_type_{ShutdownType::RESTART}; + std::atomic last_exit_code_{0}; std::unique_ptr agent_; }; -TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "No shutdown is requested when the C2 connection timeout is not set", "[c2agentwatchdog]") { +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "No shutdown is requested when the C2 connection timeout is not set", "[c2connectiontimeout]") { startAgent(); REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 1s)); std::this_thread::sleep_for(1s); - REQUIRE(recorder_.count() == 0); + REQUIRE(shutdown_request_count_ == 0); } -TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "C2 connection timeout of 0 disables the watchdog", "[c2agentwatchdog]") { +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "C2 connection timeout of 0 disables the C2 connection timeout", "[c2connectiontimeout]") { configuration_->set(Configuration::nifi_c2_connection_timeout, "0 sec"); startAgent(); REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 1s)); std::this_thread::sleep_for(1s); - REQUIRE(recorder_.count() == 0); + REQUIRE(shutdown_request_count_ == 0); REQUIRE_FALSE(LogTestController::getInstance().contains("Failed to parse C2 configuration parameter", 0ms)); } -TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "Invalid C2 connection timeout disables the watchdog", "[c2agentwatchdog]") { +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "Invalid C2 connection timeout disables the C2 connection timeout", "[c2connectiontimeout]") { configuration_->set(Configuration::nifi_c2_connection_timeout, "abc"); startAgent(); REQUIRE(LogTestController::getInstance().contains("Failed to parse C2 configuration parameter 'nifi.c2.connection.timeout': 'abc'", 1s)); std::this_thread::sleep_for(1s); - REQUIRE(recorder_.count() == 0); + REQUIRE(shutdown_request_count_ == 0); } -TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "C2 connection timeout shorter than the heartbeat period is increased", - "[c2agentwatchdog]") { +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "C2 connection timeout shorter than the heartbeat period is increased", + "[c2connectiontimeout]") { configuration_->set(Configuration::nifi_c2_connection_timeout, "50 ms"); startAgent(); REQUIRE(LogTestController::getInstance().contains("C2 connection timeout is less than the heartbeat period", 1s)); REQUIRE(LogTestController::getInstance().contains("Using 200ms as the C2 connection timeout", 1s)); } -TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "Failed heartbeats do not prevent the C2 connection timeout", "[c2agentwatchdog]") { +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "Failed heartbeats do not prevent the C2 connection timeout", "[c2connectiontimeout]") { configuration_->set(Configuration::nifi_c2_connection_timeout, "300 ms"); startAgent(); - REQUIRE(utils::verifyEventHappenedInPollTime(3s, [&] { return recorder_.count() > 0; }, 10ms)); + REQUIRE(utils::verifyEventHappenedInPollTime(3s, [&] { return shutdown_request_count_ > 0; }, 10ms)); REQUIRE(LogTestController::getInstance().contains("Performing heartbeat", 0ms)); - REQUIRE(recorder_.count() == 1); - const auto [shutdown_type, exit_code] = recorder_.firstCall(); - CHECK(shutdown_type == ShutdownType::SHUTDOWN); - REQUIRE(exit_code.has_value()); - CHECK(*exit_code == c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); + std::this_thread::sleep_for(500ms); + REQUIRE(shutdown_request_count_ == 1); + CHECK(last_shutdown_type_ == ShutdownType::SHUTDOWN); + CHECK(last_exit_code_ == c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); } -TEST_CASE_METHOD(C2AgentWatchdogConfigurationTestFixture, "Stopping the C2 agent stops the watchdog", "[c2agentwatchdog]") { +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "Stopping the C2 agent stops the C2 connection timeout check", "[c2connectiontimeout]") { configuration_->set(Configuration::nifi_c2_connection_timeout, "300 ms"); startAgent(); agent_->stop(); std::this_thread::sleep_for(600ms); - REQUIRE(recorder_.count() == 0); + REQUIRE(shutdown_request_count_ == 0); } } // namespace org::apache::nifi::minifi::test diff --git a/libminifi/test/unit/C2ConnectionWatchdogTests.cpp b/libminifi/test/unit/C2ConnectionWatchdogTests.cpp deleted file mode 100644 index c79bd7313d..0000000000 --- a/libminifi/test/unit/C2ConnectionWatchdogTests.cpp +++ /dev/null @@ -1,113 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#include -#include - -#include "C2ShutdownRecorder.h" -#include "c2/C2ConnectionWatchdog.h" -#include "c2/C2ExitCodes.h" -#include "unit/Catch.h" -#include "unit/TestBase.h" -#include "unit/TestUtils.h" - -using namespace std::literals::chrono_literals; - -namespace org::apache::nifi::minifi::test { - -namespace { -void verifyTimeoutShutdownRequested(C2ShutdownRecorder& recorder) { - REQUIRE(utils::verifyEventHappenedInPollTime(2s, [&] { return recorder.count() > 0; }, 10ms)); - REQUIRE(recorder.count() == 1); - const auto [shutdown_type, exit_code] = recorder.firstCall(); - CHECK(shutdown_type == ShutdownType::SHUTDOWN); - REQUIRE(exit_code.has_value()); - CHECK(*exit_code == c2::C2_CONNECTION_TIMEOUT_EXIT_CODE); -} -} // namespace - -TEST_CASE("Test invalid values", "[c2watchdog]") { - LogTestController::getInstance().reset(); - SECTION("Zero timeout value") { - REQUIRE_THROWS_WITH(c2::C2ConnectionWatchdog(0ms, [](ShutdownType, std::optional) {}), - "C2 connection timeout must be greater than 0"); - } - - SECTION("Negative timeout value") { - REQUIRE_THROWS_WITH(c2::C2ConnectionWatchdog(-1ms, [](ShutdownType, std::optional) {}), - "C2 connection timeout must be greater than 0"); - } -} - -TEST_CASE("Test shutdown when no C2 connection is present at all", "[c2watchdog]") { - LogTestController::getInstance().reset(); - C2ShutdownRecorder recorder; - c2::C2ConnectionWatchdog watchdog(400ms, recorder.callback()); - verifyTimeoutShutdownRequested(recorder); - REQUIRE(LogTestController::getInstance().contains("C2 connection timeout exceeded, shutting down...", 1s)); -} - -TEST_CASE("Test shutdown when C2 connection is lost", "[c2watchdog]") { - LogTestController::getInstance().reset(); - C2ShutdownRecorder recorder; - c2::C2ConnectionWatchdog watchdog(400ms, recorder.callback()); - for (int i = 0; i < 10; ++i) { - watchdog.heartBeatHappened(); - REQUIRE(recorder.count() == 0); - std::this_thread::sleep_for(50ms); - } - verifyTimeoutShutdownRequested(recorder); - REQUIRE(LogTestController::getInstance().contains("C2 connection timeout exceeded, shutting down...", 1s)); -} - -TEST_CASE("Shutdown is only requested once", "[c2watchdog]") { - LogTestController::getInstance().reset(); - C2ShutdownRecorder recorder; - c2::C2ConnectionWatchdog watchdog(100ms, recorder.callback()); - verifyTimeoutShutdownRequested(recorder); - std::this_thread::sleep_for(300ms); - REQUIRE(recorder.count() == 1); -} - -TEST_CASE("Watchdog is destroyed promptly even with a long timeout", "[c2watchdog]") { - LogTestController::getInstance().reset(); - C2ShutdownRecorder recorder; - auto watchdog = std::make_unique(1h, recorder.callback()); - const auto destruction_start = std::chrono::steady_clock::now(); - watchdog.reset(); - REQUIRE(std::chrono::steady_clock::now() - destruction_start < 1s); - REQUIRE(recorder.count() == 0); -} - -TEST_CASE("No shutdown is requested after the watchdog is destroyed", "[c2watchdog]") { - LogTestController::getInstance().reset(); - C2ShutdownRecorder recorder; - { c2::C2ConnectionWatchdog watchdog(200ms, recorder.callback()); } - std::this_thread::sleep_for(500ms); - REQUIRE(recorder.count() == 0); -} - -TEST_CASE("No shutdown is requested after the watchdog is stopped", "[c2watchdog]") { - LogTestController::getInstance().reset(); - C2ShutdownRecorder recorder; - c2::C2ConnectionWatchdog watchdog(200ms, recorder.callback()); - watchdog.stop(); - std::this_thread::sleep_for(500ms); - watchdog.heartBeatHappened(); - REQUIRE(recorder.count() == 0); -} - -} // namespace org::apache::nifi::minifi::test diff --git a/libminifi/test/unit/C2ShutdownRecorder.h b/libminifi/test/unit/C2ShutdownRecorder.h deleted file mode 100644 index da50c878df..0000000000 --- a/libminifi/test/unit/C2ShutdownRecorder.h +++ /dev/null @@ -1,53 +0,0 @@ -/** - * Licensed to the Apache Software Foundation (ASF) under one or more - * contributor license agreements. See the NOTICE file distributed with - * this work for additional information regarding copyright ownership. - * The ASF licenses this file to You under the Apache License, Version 2.0 - * (the "License"); you may not use this file except in compliance with - * the License. You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -#pragma once - -#include -#include -#include -#include -#include - -#include "RequestShutdown.h" - -namespace org::apache::nifi::minifi::test { - -class C2ShutdownRecorder { - public: - RequestShutdownCallback callback() { - return [this](ShutdownType shutdown_type, std::optional exit_code) { - std::lock_guard lock(mutex_); - calls_.emplace_back(shutdown_type, exit_code); - }; - } - - size_t count() { - std::lock_guard lock(mutex_); - return calls_.size(); - } - - std::pair> firstCall() { - std::lock_guard lock(mutex_); - return calls_.at(0); - } - - private: - std::mutex mutex_; - std::vector>> calls_; -}; - -} // namespace org::apache::nifi::minifi::test From 6c3aea9cee887cb115528897f8304dad51a9a1b0 Mon Sep 17 00:00:00 2001 From: Gabor Gyimesi Date: Thu, 5 Nov 2026 10:43:52 +0100 Subject: [PATCH 4/6] Review update --- libminifi/src/c2/C2Agent.cpp | 4 +- .../unit/C2AgentConnectionTimeoutTests.cpp | 6 +- minifi_main/MiNiFiMain.cpp | 81 ++++++++++++------- 3 files changed, 55 insertions(+), 36 deletions(-) diff --git a/libminifi/src/c2/C2Agent.cpp b/libminifi/src/c2/C2Agent.cpp index 67c48259ba..1b3a2e86f5 100644 --- a/libminifi/src/c2/C2Agent.cpp +++ b/libminifi/src/c2/C2Agent.cpp @@ -201,9 +201,9 @@ void C2Agent::configure(const std::shared_ptr &configure, bool reconf logger_->log_debug("Using {} as the heartbeat period", heart_beat_period_); c2_connection_timeout_ = parseTimeoutPeriod(configuration_, std::string{Configuration::nifi_c2_connection_timeout}, logger_); - if (c2_connection_timeout_ && *c2_connection_timeout_ <= heart_beat_period_) { + if (c2_connection_timeout_ && *c2_connection_timeout_ <= 2 * heart_beat_period_) { logger_->log_warn( - "C2 connection timeout is less than the heartbeat period: c2 connection timeout={}, heartbeat period={}. Adjusting it to be twice the " + "C2 connection timeout is less than twice the heartbeat period: C2 connection timeout={}, heartbeat period={}. Adjusting it to be twice the " "heartbeat period", *c2_connection_timeout_, heart_beat_period_); diff --git a/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp b/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp index 3399994513..9ef49654e2 100644 --- a/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp +++ b/libminifi/test/unit/C2AgentConnectionTimeoutTests.cpp @@ -97,11 +97,11 @@ TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "Invalid C2 connection tim REQUIRE(shutdown_request_count_ == 0); } -TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "C2 connection timeout shorter than the heartbeat period is increased", +TEST_CASE_METHOD(C2AgentConnectionTimeoutTestFixture, "C2 connection timeout shorter than twice the heartbeat period is increased", "[c2connectiontimeout]") { - configuration_->set(Configuration::nifi_c2_connection_timeout, "50 ms"); + configuration_->set(Configuration::nifi_c2_connection_timeout, "150 ms"); startAgent(); - REQUIRE(LogTestController::getInstance().contains("C2 connection timeout is less than the heartbeat period", 1s)); + REQUIRE(LogTestController::getInstance().contains("C2 connection timeout is less than twice the heartbeat period", 1s)); REQUIRE(LogTestController::getInstance().contains("Using 200ms as the C2 connection timeout", 1s)); } diff --git a/minifi_main/MiNiFiMain.cpp b/minifi_main/MiNiFiMain.cpp index 5349f6dcbc..4ed0142db8 100644 --- a/minifi_main/MiNiFiMain.cpp +++ b/minifi_main/MiNiFiMain.cpp @@ -52,6 +52,7 @@ #include "Fips.h" #include "FlowController.h" #include "MainHelper.h" +#include "RequestShutdown.h" #include "ResourceClaim.h" #include "agent/JsonSchema.h" #include "minifi-cpp/agent/agent_version.h" @@ -80,6 +81,41 @@ namespace utils = minifi::utils; static std::atomic_flag flow_controller_running; static std::atomic_flag process_running; +enum class LifecycleState : int32_t { + RUNNING, + RESTART_REQUESTED, + SHUTDOWN_REQUESTED +}; + +struct LifecycleRequest { + LifecycleState state = LifecycleState::RUNNING; + int32_t exit_code = 0; +}; + +static_assert(std::atomic::is_always_lock_free); +static std::atomic lifecycle_request; + +bool requestLifecycleChange(minifi::ShutdownType shutdown_type, int32_t exit_code) { + auto current_request = lifecycle_request.load(); + while (true) { + if (current_request.state == LifecycleState::SHUTDOWN_REQUESTED) { + return false; + } + if (shutdown_type == minifi::ShutdownType::RESTART && current_request.state == LifecycleState::RESTART_REQUESTED) { + return false; + } + const auto desired = shutdown_type == minifi::ShutdownType::SHUTDOWN + ? LifecycleRequest{.state = LifecycleState::SHUTDOWN_REQUESTED, .exit_code = exit_code} + : LifecycleRequest{.state = LifecycleState::RESTART_REQUESTED, .exit_code = 0}; + if (lifecycle_request.compare_exchange_weak(current_request, desired)) { + break; + } + } + flow_controller_running.clear(); + flow_controller_running.notify_all(); + return true; +} + /** * Removed the stop command from the signal handler so that we could trigger * unload after we exit the semaphore controlled critical section in main. @@ -94,8 +130,7 @@ static std::atomic_flag process_running; BOOL WINAPI consoleSignalHandler(DWORD signal) { if (!process_running.test()) { exit(0); return TRUE; } if (signal == CTRL_C_EVENT || signal == CTRL_BREAK_EVENT) { - flow_controller_running.clear(); - flow_controller_running.notify_all(); + requestLifecycleChange(minifi::ShutdownType::SHUTDOWN, 0); process_running.wait(true); return TRUE; } @@ -103,15 +138,13 @@ BOOL WINAPI consoleSignalHandler(DWORD signal) { } void SignalExitProcess() { - flow_controller_running.clear(); - flow_controller_running.notify_all(); + requestLifecycleChange(minifi::ShutdownType::SHUTDOWN, 0); } #endif void sigHandler(int signal) { if (signal == SIGINT || signal == SIGTERM) { - flow_controller_running.clear(); - flow_controller_running.notify_all(); + requestLifecycleChange(minifi::ShutdownType::SHUTDOWN, 0); } } @@ -176,6 +209,8 @@ void overridePropertiesFromCommandLine(const argparse::ArgumentParser& parser, c } int main(int argc, char **argv) { + flow_controller_running.test_and_set(); + argparse::ArgumentParser argument_parser("Apache MiNiFi C++", minifi::AgentBuild::VERSION); argument_parser.add_argument("-p", "--property") .append() @@ -290,36 +325,18 @@ int main(int argc, char **argv) { process_running.test_and_set(); - std::atomic restart_token{false}; - std::atomic exit_code{0}; - std::mutex shutdown_request_mutex; - bool shutdown_requested = false; const auto request_shutdown = [&] (minifi::ShutdownType shutdown_type, std::optional shutdown_exit_code) { - std::lock_guard lock(shutdown_request_mutex); - if (shutdown_requested) { + if (!requestLifecycleChange(shutdown_type, shutdown_exit_code.value_or(0))) { return; } if (shutdown_type == minifi::ShutdownType::SHUTDOWN) { - shutdown_requested = true; - restart_token = false; - exit_code = shutdown_exit_code.value_or(0); - flow_controller_running.clear(); - flow_controller_running.notify_all(); logger->log_info("Initiating shutdown..."); - return; - } - - if (!restart_token.exchange(true)) { - // only trigger if a restart is not already in progress (the flag was unset before the exchange) - flow_controller_running.clear(); - flow_controller_running.notify_all(); + } else { logger->log_info("Initiating restart..."); } }; do { - flow_controller_running.test_and_set(); - std::string prov_repo_class = "provenancerepository"; std::string flow_repo_class = "flowfilerepository"; std::string nifi_configuration_class_name = "adaptiveconfiguration"; @@ -513,15 +530,17 @@ int main(int argc, char **argv) { flow_repo = nullptr; prov_repo = nullptr; } while ([&] { - const auto restart_token_temp = restart_token.exchange(false); - if (restart_token_temp) { - logger->log_info("Restarting MiNiFi"); + flow_controller_running.test_and_set(); + auto expected = LifecycleRequest{.state = LifecycleState::RESTART_REQUESTED, .exit_code = 0}; + if (!lifecycle_request.compare_exchange_strong(expected, LifecycleRequest{})) { + return false; } - return restart_token_temp; + logger->log_info("Restarting MiNiFi"); + return true; }()); process_running.clear(); process_running.notify_all(); logger->log_info("MiNiFi exit"); - return exit_code; + return lifecycle_request.load().exit_code; } From 4113ab830ab42d6fd7b2d9f42280e356ae33e86b Mon Sep 17 00:00:00 2001 From: Gabor Gyimesi Date: Thu, 5 Nov 2026 15:00:02 +0100 Subject: [PATCH 5/6] Add comment for static assertion --- minifi_main/MiNiFiMain.cpp | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/minifi_main/MiNiFiMain.cpp b/minifi_main/MiNiFiMain.cpp index 4ed0142db8..1516a22f84 100644 --- a/minifi_main/MiNiFiMain.cpp +++ b/minifi_main/MiNiFiMain.cpp @@ -92,7 +92,11 @@ struct LifecycleRequest { int32_t exit_code = 0; }; -static_assert(std::atomic::is_always_lock_free); +// Ensure that LifecycleRequest can be safely used with atomic operations in signal handlers +// Extending LifecycleRequest in the future may break the lock-free guarantee +static_assert(std::atomic::is_always_lock_free, + "lifecycle_request is modified from signal handlers (see sigHandler), where only lock-free atomic operations are safe. " + "A non-lock-free std::atomic falls back to hidden mutex or a lock, which can deadlock if a signal interrupts a thread holding it."); static std::atomic lifecycle_request; bool requestLifecycleChange(minifi::ShutdownType shutdown_type, int32_t exit_code) { From 221604d4c2073a7c39f06f3a8d9dac571de18fed Mon Sep 17 00:00:00 2001 From: Gabor Gyimesi Date: Thu, 5 Nov 2026 15:35:26 +0100 Subject: [PATCH 6/6] Review update --- C2.md | 2 +- libminifi/src/c2/C2Agent.cpp | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/C2.md b/C2.md index b338a52271..1f171fcc4f 100644 --- a/C2.md +++ b/C2.md @@ -140,7 +140,7 @@ It is possible to set a persistent fallback agent id. This is needed so that the #### C2 Connection Timeout -There is an option to shut down the agent if the connection to the C2 server exceeds a specified timeout. This is controlled by the `nifi.c2.connection.timeout` property, which defines the maximum time allowed without a successful heartbeat response. If the connection is not established (or is lost) within this period, the agent will trigger a shutdown with exit code 90. If the agent is shut down with exit 90 neither the systemd service nor the Windows service will restart it. The feature is off by default, and can be enabled by setting a non-zero value. If the timeout is set to a value less than or equal to the heartbeat period, it will be automatically adjusted to twice the heartbeat period. +There is an option to shut down the agent if the connection to the C2 server exceeds a specified timeout. This is controlled by the `nifi.c2.connection.timeout` property, which defines the maximum time allowed without a successful heartbeat response. If the connection is not established (or is lost) within this period, the agent will trigger a shutdown with exit code 90. If the agent is shut down with exit 90 neither the systemd service nor the Windows service will restart it. The feature is off by default, and can be enabled by setting a non-zero value. If the timeout is set to a value less than twice the heartbeat period, it will be automatically adjusted to twice the heartbeat period. # in minifi.properties # timeout for C2 connection, setting to 0 disables the timeout diff --git a/libminifi/src/c2/C2Agent.cpp b/libminifi/src/c2/C2Agent.cpp index 1b3a2e86f5..d4b44a12ba 100644 --- a/libminifi/src/c2/C2Agent.cpp +++ b/libminifi/src/c2/C2Agent.cpp @@ -201,7 +201,7 @@ void C2Agent::configure(const std::shared_ptr &configure, bool reconf logger_->log_debug("Using {} as the heartbeat period", heart_beat_period_); c2_connection_timeout_ = parseTimeoutPeriod(configuration_, std::string{Configuration::nifi_c2_connection_timeout}, logger_); - if (c2_connection_timeout_ && *c2_connection_timeout_ <= 2 * heart_beat_period_) { + if (c2_connection_timeout_ && *c2_connection_timeout_ < 2 * heart_beat_period_) { logger_->log_warn( "C2 connection timeout is less than twice the heartbeat period: C2 connection timeout={}, heartbeat period={}. Adjusting it to be twice the " "heartbeat period",