Support graceful shutdown

CHANGELOG.md

@@ -10,6 +10,7 @@ All notable changes to this project will be documented in this file.
 - Configuration overrides for the JVM security properties, such as DNS caching ([#497]).
 - Support PodDisruptionBudgets ([#509]).
 - Support for 1.23.2 ([#513]).
+- Support graceful shutdown ([#528]).
 
 ### Changed
 
@@ -31,6 +32,7 @@ All notable changes to this project will be documented in this file.
 [#505]: https://github.com/stackabletech/nifi-operator/pull/505
 [#509]: https://github.com/stackabletech/nifi-operator/pull/509
 [#513]: https://github.com/stackabletech/nifi-operator/pull/513
+[#528]: https://github.com/stackabletech/nifi-operator/pull/528
 
 ## [23.7.0] - 2023-07-14
 

Cargo.toml

@@ -16,14 +16,14 @@ fnv = "1.0"
 futures = { version = "0.3", features = ["compat"] }
 indoc = "2.0"
 pin-project = "1.1"
+product-config = { git = "https://github.com/stackabletech/product-config.git", tag = "0.6.0" }
 rand = "0.8"
 semver = "1.0"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 serde_yaml = "0.9"
 snafu = "0.7"
 stackable-operator = { git = "https://github.com/stackabletech/operator-rs.git", tag = "0.56.0" }
-product-config = { git = "https://github.com/stackabletech/product-config.git", tag = "0.6.0" }
 strum = { version = "0.25", features = ["derive"] }
 tokio = { version = "1.29", features = ["full"] }
 tracing = "0.1"

deploy/helm/nifi-operator/crds/crds.yaml

@@ -1704,6 +1704,10 @@ spec:
                                   type: array
                               type: object
                           type: object
+                        gracefulShutdownTimeout:
+                          description: Time period Pods have to gracefully shut down, e.g. `30m`, `1h` or `2d`. Consult the operator documentation for details.
+                          nullable: true
+                          type: string
                         logging:
                           default:
                             enableVectorAgent: null
@@ -5368,6 +5372,10 @@ spec:
                                         type: array
                                     type: object
                                 type: object
+                              gracefulShutdownTimeout:
+                                description: Time period Pods have to gracefully shut down, e.g. `30m`, `1h` or `2d`. Consult the operator documentation for details.
+                                nullable: true
+                                type: string
                               logging:
                                 default:
                                   enableVectorAgent: null

docs/modules/nifi/pages/usage_guide/operations/graceful-shutdown.adoc

@@ -1,7 +1,23 @@
 = Graceful shutdown
 
-Graceful shutdown of Nifi nodes is either not supported by the product itself
-or we have not implemented it yet.
+You can configure the graceful shutdown as described in xref:concepts:operations/graceful_shutdown.adoc[].
 
-Outstanding implementation work for the graceful shutdowns of all products where this functionality is relevant is tracked in
-https://github.com/stackabletech/issues/issues/357
+== Nodes
+
+As a default, NiFi nodes have `5 minutes` to shut down gracefully.
+
+The NiFi node process will receive a `SIGTERM` signal when Kubernetes wants to terminate the Pod.
+It will log the received signal as shown in the log below and initiate a graceful shutdown.
+After the graceful shutdown timeout runs out, and the process still didn't exit, Kubernetes will issue a `SIGKILL` signal.
+
+[source,text]
+----
+nifi NiFi PID [43] shutdown started
+nifi NiFi PID [43] shutdown in progress...
+[...]
+nifi 2023-11-08 10:52:14,459 INFO [pool-1-thread-1] org.apache.nifi.NiFi Application Server shutdown completed
+nifi 2023-11-08 10:52:15,139 INFO [Thread-0] org.apache.nifi.NiFi Application Server shutdown started
+nifi 2023-11-08 10:52:15,139 INFO [Thread-0] org.apache.nifi.nar.NarAutoLoader NAR Auto-Loader stopped
+nifi 2023-11-08 10:52:15,139 INFO [Thread-0] o.a.n.f.r.CompositeExternalResourceProviderService External Resource Provider Service is stopped
+nifi 2023-11-08 10:52:15,139 INFO [Thread-0] org.apache.nifi.NiFi Application Server shutdown completed
+----

rust/crd/src/lib.rs

@@ -29,6 +29,7 @@ use stackable_operator::{
     role_utils::{GenericRoleConfig, Role, RoleGroup, RoleGroupRef},
     schemars::{self, JsonSchema},
     status::condition::{ClusterCondition, HasStatusCondition},
+    time::Duration,
 };
 use std::collections::BTreeMap;
 use strum::Display;
@@ -56,6 +57,8 @@ pub const MAX_PREPARE_LOG_FILE_SIZE: MemoryQuantity = MemoryQuantity {
     unit: BinaryMultiple::Mebi,
 };
 
+const DEFAULT_NODE_GRACEFUL_SHUTDOWN_TIMEOUT: Duration = Duration::from_minutes_unchecked(5);
+
 #[derive(Snafu, Debug)]
 pub enum Error {
     #[snafu(display("object has no namespace associated"))]
@@ -264,6 +267,9 @@ pub struct NifiConfig {
     pub resources: Resources<NifiStorageConfig, NoRuntimeLimits>,
     #[fragment_attrs(serde(default))]
     pub affinity: StackableAffinity,
+    /// Time period Pods have to gracefully shut down, e.g. `30m`, `1h` or `2d`. Consult the operator documentation for details.
+    #[fragment_attrs(serde(default))]
+    pub graceful_shutdown_timeout: Option<Duration>,
 }
 
 impl NifiConfig {
@@ -310,6 +316,7 @@ impl NifiConfig {
                 },
             },
             affinity: get_affinity(cluster_name, role),
+            graceful_shutdown_timeout: Some(DEFAULT_NODE_GRACEFUL_SHUTDOWN_TIMEOUT),
         }
     }
 }

rust/operator-binary/src/config.rs

@@ -6,7 +6,7 @@ use std::{
 use product_config::{types::PropertyNameKind, ProductConfigManager};
 use snafu::{ResultExt, Snafu};
 use stackable_nifi_crd::{
-    NifiCluster, NifiConfigFragment, NifiRole, NifiSpec, NifiStorageConfig, HTTPS_PORT,
+    NifiCluster, NifiConfig, NifiConfigFragment, NifiRole, NifiSpec, NifiStorageConfig, HTTPS_PORT,
     PROTOCOL_PORT,
 };
 use stackable_operator::{
@@ -21,6 +21,7 @@ use stackable_operator::{
 use strum::{Display, EnumIter};
 
 use crate::authentication::{STACKABLE_SERVER_TLS_DIR, STACKABLE_TLS_STORE_PASSWORD};
+use crate::operations::graceful_shutdown::graceful_shutdown_config_properties;
 
 pub const NIFI_CONFIG_DIRECTORY: &str = "/stackable/nifi/conf";
 
@@ -79,7 +80,7 @@ pub enum Error {
 
 /// Create the NiFi bootstrap.conf
 pub fn build_bootstrap_conf(
-    resource_config: &Resources<NifiStorageConfig>,
+    nifi_config: &NifiConfig,
     overrides: BTreeMap<String, String>,
 ) -> Result<String, Error> {
     let mut bootstrap = BTreeMap::new();
@@ -92,15 +93,14 @@ pub fn build_bootstrap_conf(
     // Configure where NiFi's lib and conf directories live
     bootstrap.insert("lib.dir".to_string(), "./lib".to_string());
     bootstrap.insert("conf.dir".to_string(), "./conf".to_string());
-    // How long to wait after telling NiFi to shutdown before explicitly killing the Process
-    bootstrap.insert("graceful.shutdown.seconds".to_string(), "20".to_string());
+    bootstrap.extend(graceful_shutdown_config_properties(nifi_config));
 
     let mut java_args = Vec::with_capacity(18);
     // Disable JSR 199 so that we can use JSP's without running a JDK
     java_args.push("-Dorg.apache.jasper.compiler.disablejsr199=true".to_string());
 
     // Read memory limits from config
-    if let Some(heap_size_definition) = &resource_config.memory.limit {
+    if let Some(heap_size_definition) = &nifi_config.resources.memory.limit {
         tracing::debug!("Read {:?} from crd as memory limit", heap_size_definition);
 
         let heap_size = MemoryQuantity::try_from(heap_size_definition)

rust/operator-binary/src/controller.rs

@@ -1,11 +1,21 @@
 //! Ensures that `Pod`s are configured and running for each [`NifiCluster`]
-use std::{
-    borrow::Cow,
-    collections::{BTreeMap, HashMap},
-    ops::Deref,
-    sync::Arc,
+use crate::{
+    authentication::{
+        NifiAuthenticationConfig, AUTHORIZERS_XML_FILE_NAME,
+        LOGIN_IDENTITY_PROVIDERS_XML_FILE_NAME, STACKABLE_ADMIN_USER_NAME,
+        STACKABLE_SERVER_TLS_DIR, STACKABLE_TLS_STORE_PASSWORD,
+    },
+    config::{
+        self, build_bootstrap_conf, build_nifi_properties, build_state_management_xml,
+        validated_product_config, NifiRepository, JVM_SECURITY_PROPERTIES_FILE,
+        NIFI_BOOTSTRAP_CONF, NIFI_CONFIG_DIRECTORY, NIFI_PROPERTIES, NIFI_STATE_MANAGEMENT_XML,
+    },
+    operations::{graceful_shutdown::add_graceful_shutdown_config, pdb::add_pdbs},
+    product_logging::{extend_role_group_config_map, resolve_vector_aggregator_address},
+    OPERATOR_NAME,
 };
 
+use indoc::formatdoc;
 use product_config::{
     types::PropertyNameKind,
     writer::{to_java_properties_string, PropertiesWriterError},
@@ -51,6 +61,7 @@ use stackable_operator::{
     logging::controller::ReconcilerError,
     product_logging::{
         self,
+        framework::{create_vector_shutdown_file_command, remove_vector_shutdown_file_command},
         spec::{
             ConfigMapLogConfig, ContainerLogConfig, ContainerLogConfigChoice,
             CustomContainerLogConfig,
@@ -62,26 +73,17 @@ use stackable_operator::{
         statefulset::StatefulSetConditionBuilder,
     },
     time::Duration,
+    utils::COMMON_BASH_TRAP_FUNCTIONS,
+};
+use std::{
+    borrow::Cow,
+    collections::{BTreeMap, HashMap},
+    ops::Deref,
+    sync::Arc,
 };
 use strum::{EnumDiscriminants, IntoStaticStr};
 use tracing::Instrument;
 
-use crate::{
-    authentication::{
-        NifiAuthenticationConfig, AUTHORIZERS_XML_FILE_NAME,
-        LOGIN_IDENTITY_PROVIDERS_XML_FILE_NAME, STACKABLE_ADMIN_USER_NAME,
-        STACKABLE_SERVER_TLS_DIR, STACKABLE_TLS_STORE_PASSWORD,
-    },
-    config::{
-        self, build_bootstrap_conf, build_nifi_properties, build_state_management_xml,
-        validated_product_config, NifiRepository, JVM_SECURITY_PROPERTIES_FILE,
-        NIFI_BOOTSTRAP_CONF, NIFI_CONFIG_DIRECTORY, NIFI_PROPERTIES, NIFI_STATE_MANAGEMENT_XML,
-    },
-    operations::pdb::add_pdbs,
-    product_logging::{extend_role_group_config_map, resolve_vector_aggregator_address},
-    OPERATOR_NAME,
-};
-
 pub const NIFI_CONTROLLER_NAME: &str = "nificluster";
 pub const NIFI_UID: i64 = 1000;
 
@@ -277,6 +279,11 @@ pub enum Error {
     FailedToCreatePdb {
         source: crate::operations::pdb::Error,
     },
+
+    #[snafu(display("failed to configure graceful shutdown"))]
+    GracefulShutdown {
+        source: crate::operations::graceful_shutdown::Error,
+    },
 }
 
 type Result<T, E = Error> = std::result::Result<T, E>;
@@ -675,7 +682,7 @@ async fn build_node_rolegroup_config_map(
         .add_data(
             NIFI_BOOTSTRAP_CONF,
             build_bootstrap_conf(
-                &merged_config.resources,
+                merged_config,
                 rolegroup_config
                     .get(&PropertyNameKind::File(NIFI_BOOTSTRAP_CONF.to_string()))
                     .with_context(|| ProductConfigKindNotSpecifiedSnafu {
@@ -862,20 +869,20 @@ async fn build_node_rolegroup_statefulset(
     let sensitive_key_secret = &nifi.spec.cluster_config.sensitive_properties.key_secret;
 
     let prepare_container_name = Container::Prepare.to_string();
-    let mut args = vec![];
+    let mut prepare_args = vec![];
 
     if let Some(ContainerLogConfig {
         choice: Some(ContainerLogConfigChoice::Automatic(log_config)),
     }) = merged_config.logging.containers.get(&Container::Prepare)
     {
-        args.push(product_logging::framework::capture_shell_output(
+        prepare_args.push(product_logging::framework::capture_shell_output(
             STACKABLE_LOG_DIR,
             &prepare_container_name,
             log_config,
         ));
     }
 
-    args.extend(vec![
+    prepare_args.extend(vec![
         // The source directory is a secret-op mount and we do not want to write / add anything in there
         // Therefore we import all the contents to a truststore in "writeable" empty dirs.
         // Keytool is only barking if a password is not set for the destination truststore (which we set)
@@ -904,7 +911,7 @@ async fn build_node_rolegroup_statefulset(
         format!("sed -i \"s|yyyyyy|${{{}}}|g\" /stackable/nifi/conf/state-management.xml",zookeeper_chroot)
     ]);
 
-    args.extend_from_slice(nifi_auth_config.get_additional_container_args().as_slice());
+    prepare_args.extend_from_slice(nifi_auth_config.get_additional_container_args().as_slice());
 
     let mut container_prepare =
         ContainerBuilder::new(&prepare_container_name).with_context(|_| {
@@ -922,7 +929,7 @@ async fn build_node_rolegroup_statefulset(
             "pipefail".to_string(),
         ])
         .add_env_vars(env_vars.clone())
-        .args(vec![args.join(" && ")])
+        .args(vec![prepare_args.join(" && ")])
         .add_volume_mount(
             &NifiRepository::Flowfile.repository(),
             &NifiRepository::Flowfile.mount_path(),
@@ -965,10 +972,23 @@ async fn build_node_rolegroup_statefulset(
         }
     })?;
 
+    let nifi_args = vec![formatdoc! {"
+            {COMMON_BASH_TRAP_FUNCTIONS}
+            {remove_vector_shutdown_file_command}
+            prepare_signal_handlers
+            bin/nifi.sh run &
+            wait_for_termination $!
+            {create_vector_shutdown_file_command}
+            ",
+    remove_vector_shutdown_file_command =
+        remove_vector_shutdown_file_command(STACKABLE_LOG_DIR),
+    create_vector_shutdown_file_command =
+        create_vector_shutdown_file_command(STACKABLE_LOG_DIR),
+    }];
     let container_nifi = container_builder
         .image_from_product_image(resolved_product_image)
         .command(vec!["/bin/bash".to_string(), "-c".to_string()])
-        .args(vec!["bin/nifi.sh run".to_string()])
+        .args(nifi_args)
         .add_env_vars(env_vars)
         .add_volume_mount(KEYSTORE_VOLUME_NAME, KEYSTORE_NIFI_CONTAINER_MOUNT)
         .add_volume_mount(
@@ -1021,6 +1041,7 @@ async fn build_node_rolegroup_statefulset(
         .resources(merged_config.resources.clone().into());
 
     let mut pod_builder = PodBuilder::new();
+    add_graceful_shutdown_config(merged_config, &mut pod_builder).context(GracefulShutdownSnafu)?;
 
     // Add user configured extra volumes if any are specified
     for volume in &nifi.spec.cluster_config.extra_volumes {

rust/operator-binary/src/operations/graceful_shutdown.rs

@@ -0,0 +1,38 @@
+use snafu::{ResultExt, Snafu};
+use stackable_nifi_crd::NifiConfig;
+use stackable_operator::builder::PodBuilder;
+use std::collections::BTreeMap;
+
+#[derive(Debug, Snafu)]
+pub enum Error {
+    #[snafu(display("Failed to set terminationGracePeriod"))]
+    SetTerminationGracePeriod {
+        source: stackable_operator::builder::pod::Error,
+    },
+}
+
+pub fn graceful_shutdown_config_properties(config: &NifiConfig) -> BTreeMap<String, String> {
+    let mut graceful_shutdown_properties = BTreeMap::new();
+    if let Some(graceful_shutdown_timeout) = config.graceful_shutdown_timeout {
+        graceful_shutdown_properties.insert(
+            "graceful.shutdown.seconds".to_string(),
+            graceful_shutdown_timeout.as_secs().to_string(),
+        );
+    }
+    graceful_shutdown_properties
+}
+
+pub fn add_graceful_shutdown_config(
+    merged_config: &NifiConfig,
+    pod_builder: &mut PodBuilder,
+) -> Result<(), Error> {
+    // This must be always set by the merge mechanism, as we provide a default value,
+    // users can not disable graceful shutdown.
+    if let Some(graceful_shutdown_timeout) = merged_config.graceful_shutdown_timeout {
+        pod_builder
+            .termination_grace_period(&graceful_shutdown_timeout)
+            .context(SetTerminationGracePeriodSnafu)?;
+    }
+
+    Ok(())
+}

rust/operator-binary/src/operations/mod.rs

@@ -1 +1,2 @@
+pub mod graceful_shutdown;
 pub mod pdb;

tests/templates/kuttl/smoke/30-assert.yaml

@@ -7,6 +7,10 @@ apiVersion: apps/v1
 kind: StatefulSet
 metadata:
   name: test-nifi-node-default
+spec:
+  template:
+    spec:
+      terminationGracePeriodSeconds: 300
 status:
   readyReplicas: 2
   replicas: 2