From ef4eb73087cf099da1905499e70cecc1a2cc680b Mon Sep 17 00:00:00 2001
From: Robert Gildein <robert.gildein@canonical.com>
Date: Mon, 14 Oct 2024 20:05:59 +0200
Subject: [PATCH] Fix log forwarding for driver flavour

Signed-off-by: Robert Gildein <robert.gildein@canonical.com>
---
 images/charmed-spark/bin/sparkd.sh         | 35 ++++++++++++----------
 images/charmed-spark/pebble/log-layer.yaml |  3 +-
 tests/integration/integration-tests.sh     | 13 +++++---
 3 files changed, 31 insertions(+), 20 deletions(-)

diff --git a/images/charmed-spark/bin/sparkd.sh b/images/charmed-spark/bin/sparkd.sh
index 336203fb..d143a517 100644
--- a/images/charmed-spark/bin/sparkd.sh
+++ b/images/charmed-spark/bin/sparkd.sh
@@ -1,13 +1,27 @@
 #!/bin/bash
 
 function get_log_layer {
-  LOG_LAYER_FILE="/opt/pebble/log-layer.yaml"
-  ESCAPED_LOKI_URL="$(<<< "$LOKI_URL" sed -e 's`[][\\/.*^$]`\\&`g')"
-  sed -e "s/\$LOKI_URL/$ESCAPED_LOKI_URL/g" \
+  local loki_url=$1
+  local log_layer_file=${2-"/opt/pebble/log-layer.yaml"}
+  sed -e "s/\$LOKI_URL/$loki_url/g" \
+      -e "s/\$FLAVOUR/$FLAVOUR/g" \
       -e "s/\$SPARK_APPLICATION_ID/$SPARK_APPLICATION_ID/g" \
       -e "s/\$SPARK_USER/$SPARK_USER/g" \
-      -e "s/\$SPARK_EXECUTOR_POD_NAME/$SPARK_EXECUTOR_POD_NAME/g" \
-      $LOG_LAYER_FILE
+      -e "s/\$HOSTNAME/$HOSTNAME/g" \
+      $log_layer_file
+}
+
+function log_forwarding {
+  # We need to escape special characters from URL to be able to use with template.
+  local loki_url="$(<<< "$LOKI_URL" sed -e 's`[][\\/.*^$]`\\&`g')"
+  if [ ! -z "$loki_url" ]; then
+      echo "Log-forwarding to Loki is enabled."
+      local rendered_log_layer=$(get_log_layer $loki_url)
+      echo "$rendered_log_layer" | tee /tmp/rendered_log_layer.yaml
+      pebble add logging /tmp/rendered_log_layer.yaml
+  else
+      echo "Log-forwarding to Loki is disabled."
+  fi
 }
 
 function finish {
@@ -19,22 +33,13 @@ function finish {
 }
 trap finish EXIT
 
-if [ ! -z "${LOKI_URL}" ]
-then
-    echo "Configuring log-forwarding to Loki."
-    RENDERED_LOG_LAYER=$(get_log_layer)
-    echo "$RENDERED_LOG_LAYER" | tee /tmp/rendered_log_layer.yaml
-    pebble add logging /tmp/rendered_log_layer.yaml
-else
-    echo "Log-forwarding to Loki is disabled."
-fi
-
 FLAVOUR=$1
 
 echo "Running script with ${FLAVOUR} flavour"
 
 case "${FLAVOUR}" in
   driver|executor)
+    log_forwarding
     pushd /opt/spark
     ./entrypoint.sh "$@"
     ;;
diff --git a/images/charmed-spark/pebble/log-layer.yaml b/images/charmed-spark/pebble/log-layer.yaml
index 27345ab6..231ece4a 100644
--- a/images/charmed-spark/pebble/log-layer.yaml
+++ b/images/charmed-spark/pebble/log-layer.yaml
@@ -6,7 +6,8 @@ log-targets:
     services: [all]
     labels:
       product: charmed-spark
+      flavour: $FLAVOUR
       app: spark
       app_id: $SPARK_APPLICATION_ID
       user: $SPARK_USER
-      pod: $SPARK_EXECUTOR_POD_NAME
+      hostname: $HOSTNAME
diff --git a/tests/integration/integration-tests.sh b/tests/integration/integration-tests.sh
index 3b045457..10aa6f9d 100755
--- a/tests/integration/integration-tests.sh
+++ b/tests/integration/integration-tests.sh
@@ -55,10 +55,14 @@ validate_metrics() {
 
 validate_logs() {
   log=$1
-  if [ $(grep -Ri "Configuring log-forwarding to Loki." $log | wc -l) -lt 2 ]; then
+  if [ $(grep -Ri "Log-forwarding to Loki is enabled." $log | wc -l) -lt 3 ]; then
+    echo "ERROR: Could not validate logs."
+    echo "DEBUG: Log file:\n$(cat $log)"
     exit 1
   fi
-  if [ $(grep -Ri 'Layer \\\\"logging\\\\" added successfully from \\\\"/tmp/rendered_log_layer.yaml\\\\"' $log | wc -l) -lt 2 ]; then
+  if [ $(grep -Ri 'Layer \\\\"logging\\\\" added successfully from \\\\"/tmp/rendered_log_layer.yaml\\\\"' $log | wc -l) -lt 3 ]; then
+    echo "ERROR: Could not validate logs."
+    echo "DEBUG: Log file:\n$(cat $log)"
     exit 1
   fi
 }
@@ -439,7 +443,7 @@ run_example_job_in_pod_with_log_forwarding() {
 
   PREVIOUS_JOB=$(kubectl -n $NAMESPACE get pods --sort-by=.metadata.creationTimestamp | grep driver | tail -n 1 | cut -d' ' -f1)
   # start simple http server
-  LOG_FILE="/tmp/server.log"
+  LOG_FILE="/tmp/server-loki.log"
   SERVER_PORT=9091
   python3 tests/integration/resources/test_web_server.py $SERVER_PORT > $LOG_FILE &
   HTTP_SERVER_PID=$!
@@ -453,7 +457,8 @@ run_example_job_in_pod_with_log_forwarding() {
                   --conf spark.kubernetes.driver.request.cores=100m \
                   --conf spark.kubernetes.executor.request.cores=100m \
                   --conf spark.kubernetes.container.image=$IM \
-                  --conf spark.executorEnv.LOKI_URL="$IP:$PORT" \
+                  --conf spark.executorEnv.LOKI_URL="http://$IP:$PORT" \
+                  --conf spark.kubernetes.driverEnv.LOKI_URL="http://$IP:$PORT" \
                   --class org.apache.spark.examples.SparkPi \
                   local:///opt/spark/examples/jars/$JJ 1000'