fix: handle SIGTERM for Distributed and Standalone Mode (#907)

also add readiness probes in helm chart 

Co-authored-by: Nitish Tiwari <[email protected]>

Author: AdheipSingh

Cargo.lock

@@ -31,7 +31,6 @@ FROM gcr.io/distroless/cc-debian12:latest
 WORKDIR /parseable
 
 # Copy the static shell into base image.
-COPY --from=builder /bin/sh /bin/sh
 COPY --from=builder /parseable/target/release/parseable /usr/bin/parseable
 
-CMD ["parseable"]
+CMD ["/usr/bin/parseable"]

Dockerfile

@@ -32,4 +32,4 @@ WORKDIR /parseable
 
 COPY --from=builder /parseable/target/debug/parseable /usr/bin/parseable
 
-CMD ["parseable"]
+CMD ["/usr/bin/parseable"]

Dockerfile.debug

@@ -40,12 +40,19 @@ spec:
           {{- toYaml .Values.parseable.securityContext | nindent 8 }}
         image: {{ .Values.parseable.image.repository }}:{{ .Values.parseable.image.tag | default .Chart.AppVersion }}
         imagePullPolicy: {{ .Values.parseable.image.pullPolicy }}
-        command: ["/bin/sh", "-c"]
-        args: ["parseable s3-store --ingestor-endpoint=${HOSTNAME}.{{ include "parseable.fullname" . }}-ingestor-headless.{{ .Release.Namespace }}.svc.cluster.local:{{ .Values.parseable.highAvailability.ingestor.port }}"]
+        args:
+          - /usr/bin/parseable
+          - s3-store 
+          - --ingestor-endpoint=$(HOSTNAME).{{ include "parseable.fullname" . }}-ingestor-headless.{{ .Release.Namespace }}.svc.cluster.local:{{ .Values.parseable.highAvailability.ingestor.port }}
         env:
           {{- range $key, $value :=  .Values.parseable.highAvailability.ingestor.env }}
           - name: {{ $key }}
             value: {{ tpl $value $ | quote }}
+          - name: HOSTNAME
+            valueFrom:
+              fieldRef:
+                apiVersion: v1
+                fieldPath: metadata.name
           {{- end }}
           {{- range $secret := .Values.parseable.s3ModeSecret }}
           {{- range $key := $secret.keys }}
@@ -62,6 +69,10 @@ spec:
             value: "ingest"
         ports:
           - containerPort: {{ .Values.parseable.highAvailability.ingestor.port }}
+      {{- with .Values.readinessProbe }}
+        readinessProbe:
+          {{ toYaml . | nindent 12 }}
+      {{- end }}
         resources:
           {{- toYaml .Values.parseable.highAvailability.ingestor.resources | nindent 12 }}
         {{- if .Values.parseable.persistence.ingestor.enabled }}

helm/templates/ingestor-statefulset.yaml

@@ -41,8 +41,9 @@ spec:
           {{- toYaml .Values.parseable.securityContext | nindent 8 }}
         image: {{ .Values.parseable.image.repository }}:{{ .Values.parseable.image.tag | default .Chart.AppVersion }}
         imagePullPolicy: {{ .Values.parseable.image.pullPolicy }}
-        command: ["parseable"]
-        args: ["s3-store"]
+        args: 
+        - /usr/bin/parseable
+        - s3-store
         env:
           {{- range $key, $value :=  .Values.parseable.env }}
           - name: {{ $key }}
@@ -57,6 +58,11 @@ spec:
               secretKeyRef:
                 name: {{ $secret.name }}
                 key: {{ $key }}
+          - name: HOSTNAME
+            valueFrom:
+              fieldRef:
+                apiVersion: v1
+                fieldPath: metadata.name
           {{- end }}
           {{- end }}
           - name: P_MODE
@@ -67,6 +73,10 @@ spec:
           {{- end }}
         ports:
           - containerPort: 8000
+        {{- with .Values.readinessProbe }}
+        readinessProbe:
+          {{ toYaml . | nindent 12 }}
+        {{- end }}
         resources:
           {{- toYaml .Values.parseable.resources | nindent 12 }}
         volumeMounts:

helm/templates/querier-statefulset.yaml

@@ -37,9 +37,9 @@ spec:
           # Uncomment to debug
           # command: [ "/bin/sh", "-c", "sleep 1000000" ]
           {{- if .Values.parseable.local }}
-          args: ["parseable", "local-store"]
+          args: ["/usr/bin/parseable", "local-store"]
           {{- else }}
-          args: ["parseable", "s3-store"]
+          args: ["/usr/bin/parseable", "s3-store"]
           {{- end }}
           env:
             {{- range $key, $value :=  .Values.parseable.env }}
@@ -56,6 +56,11 @@ spec:
                 secretKeyRef:
                   name: {{ $secret.name }}
                   key: {{ $key }}
+            - name: HOSTNAME
+              valueFrom:
+                fieldRef:
+                  apiVersion: v1
+                  fieldPath: metadata.name
             {{- end }}
             {{- end }}
             {{- else}}
@@ -73,6 +78,10 @@ spec:
             {{- end }}
           ports:
             - containerPort: 8000
+          {{- with .Values.readinessProbe }}
+          readinessProbe:
+            {{ toYaml . | nindent 12 }}
+          {{- end }}
           resources:
             {{- toYaml .Values.parseable.resources | nindent 12 }}
           volumeMounts:

helm/templates/standalone-deployment.yaml

@@ -23,7 +23,11 @@ parseable:
       ## Console (UI) is available on the other service (that points to the query pod)
       service:
         type: ClusterIP
-        port: 80 
+        port: 80
+      readinessProbe:
+        httpGet:
+          path: /api/v1/readiness
+          port: 8000
       resources:
         limits:
           cpu: 500m
@@ -104,6 +108,10 @@ parseable:
   service:
     type: ClusterIP
     port: 80
+  readinessProbe:
+    httpGet:
+      path: /api/v1/readiness
+      port: 8000
   resources:
     limits:
       cpu: 500m

helm/values.yaml

@@ -36,6 +36,7 @@ anyhow = { version = "1.0", features = ["backtrace"] }
 argon2 = "0.5.0"
 async-trait = "0.1"
 base64 = "0.22.0"
+lazy_static = "1.4"
 bytes = "1.4"
 byteorder = "1.4.3"
 bzip2 = { version = "*", features = ["static"] }

server/Cargo.toml

@@ -16,19 +16,77 @@
  *
  */
 
+use crate::option::CONFIG;
 use actix_web::http::StatusCode;
 use actix_web::HttpResponse;
+use lazy_static::lazy_static;
+use std::sync::Arc;
+use tokio::signal::unix::{signal, SignalKind};
+use tokio::sync::{oneshot, Mutex};
+use tokio::time::{sleep, Duration};
 
-use crate::option::CONFIG;
+// Create a global variable to store signal status
+lazy_static! {
+    static ref SIGNAL_RECEIVED: Arc<Mutex<bool>> = Arc::new(Mutex::new(false));
+}
 
 pub async fn liveness() -> HttpResponse {
     HttpResponse::new(StatusCode::OK)
 }
 
+pub async fn handle_signals(shutdown_signal: Arc<Mutex<Option<oneshot::Sender<()>>>>) {
+    let signal_received = SIGNAL_RECEIVED.clone();
+
+    let mut sigterm =
+        signal(SignalKind::terminate()).expect("Failed to set up SIGTERM signal handler");
+    log::info!("Signal handler task started");
+
+    // Block until SIGTERM is received
+    match sigterm.recv().await {
+        Some(_) => {
+            log::info!("Received SIGTERM signal at Readiness Probe Handler");
+
+            // Set the shutdown flag to true
+            let mut shutdown_flag = signal_received.lock().await;
+            *shutdown_flag = true;
+
+            // Trigger graceful shutdown
+            if let Some(shutdown_sender) = shutdown_signal.lock().await.take() {
+                let _ = shutdown_sender.send(());
+            }
+
+            // Delay to allow readiness probe to return SERVICE_UNAVAILABLE
+            let _ = sleep(Duration::from_secs(20)).await;
+
+            // Sync to local
+            crate::event::STREAM_WRITERS.unset_all();
+
+            // Sync to S3
+            if let Err(e) = CONFIG.storage().get_object_store().sync().await {
+                log::warn!("Failed to sync local data with object store. {:?}", e);
+            }
+
+            log::info!("Local and S3 Sync done, handler SIGTERM completed.");
+        }
+        None => {
+            log::info!("Signal handler received None, indicating an error or end of stream");
+        }
+    }
+
+    log::info!("Signal handler task completed");
+}
+
 pub async fn readiness() -> HttpResponse {
-    if CONFIG.storage().get_object_store().check().await.is_ok() {
-        return HttpResponse::new(StatusCode::OK);
+    // Check if the application has received a shutdown signal
+    let shutdown_flag = SIGNAL_RECEIVED.lock().await;
+    if *shutdown_flag {
+        return HttpResponse::new(StatusCode::SERVICE_UNAVAILABLE);
     }
 
-    HttpResponse::new(StatusCode::SERVICE_UNAVAILABLE)
+    // Check the object store connection
+    if CONFIG.storage().get_object_store().check().await.is_ok() {
+        HttpResponse::new(StatusCode::OK)
+    } else {
+        HttpResponse::new(StatusCode::SERVICE_UNAVAILABLE)
+    }
 }

server/src/handlers/http/health_check.rs

@@ -18,6 +18,7 @@
 use crate::analytics;
 use crate::banner;
 use crate::handlers::airplane;
+use crate::handlers::http::health_check;
 use crate::handlers::http::ingest;
 use crate::handlers::http::logstream;
 use crate::handlers::http::middleware::RouteExt;
@@ -35,6 +36,8 @@ use crate::storage::ObjectStorageError;
 use crate::storage::PARSEABLE_ROOT_DIRECTORY;
 use crate::sync;
 
+use std::sync::Arc;
+
 use super::server::Server;
 use super::ssl_acceptor::get_ssl_acceptor;
 use super::IngestorMetadata;
@@ -56,6 +59,7 @@ use bytes::Bytes;
 use once_cell::sync::Lazy;
 use relative_path::RelativePathBuf;
 use serde_json::Value;
+use tokio::sync::{oneshot, Mutex};
 
 /// ! have to use a guard before using it
 pub static INGESTOR_META: Lazy<IngestorMetadata> =
@@ -91,17 +95,47 @@ impl ParseableServer for IngestServer {
                 .wrap(cross_origin_config())
         };
 
-        // concurrent workers equal to number of logical cores
-        let http_server = HttpServer::new(create_app_fn).workers(num_cpus::get());
+        // Create a channel to trigger server shutdown
+        let (shutdown_trigger, shutdown_rx) = oneshot::channel::<()>();
+        let server_shutdown_signal = Arc::new(Mutex::new(Some(shutdown_trigger)));
+
+        // Clone the shutdown signal for the signal handler
+        let shutdown_signal = server_shutdown_signal.clone();
+
+        // Spawn the signal handler task
+        tokio::spawn(async move {
+            health_check::handle_signals(shutdown_signal).await;
+            println!("Received shutdown signal, notifying server to shut down...");
+        });
 
-        if let Some(config) = ssl {
+        // Create the HTTP server
+        let http_server = HttpServer::new(create_app_fn)
+            .workers(num_cpus::get())
+            .shutdown_timeout(60);
+
+        // Start the server with or without TLS
+        let srv = if let Some(config) = ssl {
             http_server
                 .bind_rustls_0_22(&CONFIG.parseable.address, config)?
                 .run()
-                .await?;
         } else {
-            http_server.bind(&CONFIG.parseable.address)?.run().await?;
-        }
+            http_server.bind(&CONFIG.parseable.address)?.run()
+        };
+
+        // Graceful shutdown handling
+        let srv_handle = srv.handle();
+
+        tokio::spawn(async move {
+            // Wait for the shutdown signal
+            shutdown_rx.await.ok();
+
+            // Initiate graceful shutdown
+            log::info!("Graceful shutdown of HTTP server triggered");
+            srv_handle.stop(true).await;
+        });
+
+        // Await the server to run and handle shutdown
+        srv.await?;
 
         Ok(())
     }