#796 Add worker status tracking to log statuses of other running task

pramen/core/src/main/scala/za/co/absa/pramen/core/pipeline/IngestionJob.scala

@@ -28,6 +28,7 @@ import za.co.absa.pramen.core.metastore.Metastore
 import za.co.absa.pramen.core.metastore.model.{MetaTable, ReaderMode}
 import za.co.absa.pramen.core.metastore.peristence.TransientTableManager
 import za.co.absa.pramen.core.runner.splitter.{ScheduleStrategy, ScheduleStrategySourcing}
+import za.co.absa.pramen.core.state.WorkerStatusManager
 import za.co.absa.pramen.core.utils.ConfigUtils
 import za.co.absa.pramen.core.utils.Emoji.WARNING
 import za.co.absa.pramen.core.utils.SparkUtils._
@@ -167,9 +168,11 @@ class IngestionJob(operationDef: OperationDef,
                     conf: Config,
                     jobStarted: Instant,
                     inputRecordCount: Option[Long]): SaveResult = {
+    WorkerStatusManager.setStatus(s"Running '$name' for '$infoDate'. Writing data...")
     val stats = metastore.saveTable(outputTable.name, infoDate, df, inputRecordCount)
 
     try {
+      WorkerStatusManager.setStatus(s"Running '$name' for '$infoDate'. Running post-processing...")
       source.postProcess(
         sourceTable.query,
         outputTable.name,
@@ -181,6 +184,7 @@ class IngestionJob(operationDef: OperationDef,
       case _: AbstractMethodError => log.warn(s"Sources were built using old version of Pramen that does not support post processing. Ignoring...")
     }
 
+    WorkerStatusManager.setStatus(s"Running '$name' for '$infoDate'. Finalizing...")
     source.close()
 
     val jobFinished = Instant.now
@@ -203,6 +207,7 @@ class IngestionJob(operationDef: OperationDef,
       log.info(s"Getting cached record count for '${query.query}' for $from..$to...")
       getCachedDataFrame(source, query, from, to).count()
     } else {
+      WorkerStatusManager.setStatus(s"Running '$name'. Getting record count for for '${query.query}' for $from..$to")
       source.getRecordCount(sourceTable.query, from, to)
     }
   }

pramen/core/src/main/scala/za/co/absa/pramen/core/runner/orchestrator/OrchestratorImpl.scala

@@ -26,7 +26,8 @@ import za.co.absa.pramen.core.exceptions.{FatalErrorWrapper, ValidationException
 import za.co.absa.pramen.core.pipeline.{Job, JobDependency, OperationType}
 import za.co.absa.pramen.core.runner.jobrunner.ConcurrentJobRunner
 import za.co.absa.pramen.core.runner.splitter.ScheduleStrategyUtils.evaluateRunDate
-import za.co.absa.pramen.core.state.PipelineState
+import za.co.absa.pramen.core.state.{PipelineState, WorkerStatusManager}
+import za.co.absa.pramen.core.utils.Emoji
 import za.co.absa.pramen.core.utils.Emoji._
 
 import java.time.LocalDate
@@ -227,6 +228,18 @@ class OrchestratorImpl extends Orchestrator {
       log.warn(s"$FAILURE Job '${job.name}' outputting to '${outputTable.name}' has FAILED.")
       dependencyResolver.setFailedTable(outputTable.name)
     }
+
+    if (!isLazy) {
+      val statuses = WorkerStatusManager.getStatuses
+      if (statuses.nonEmpty) {
+        this.synchronized{
+          log.info(s"${Emoji.HAMMER_AND_WRENCH} Statuses of other workers:")
+          statuses.foreach { status =>
+            log.info(s"Thread ${status.threadId}: $status")
+          }
+        }
+      }
+    }
   }
 
   @throws[ValidationException]

pramen/core/src/main/scala/za/co/absa/pramen/core/runner/task/TaskRunnerBase.scala

@@ -34,7 +34,7 @@ import za.co.absa.pramen.core.metastore.model.MetaTable
 import za.co.absa.pramen.core.pipeline.JobPreRunStatus._
 import za.co.absa.pramen.core.pipeline.PipelineDef.{COUNTRY_KEY, ENVIRONMENT_NAME, PIPELINE_NAME_KEY, TENANT_KEY}
 import za.co.absa.pramen.core.pipeline._
-import za.co.absa.pramen.core.state.PipelineState
+import za.co.absa.pramen.core.state.{PipelineState, WorkerStatusManager}
 import za.co.absa.pramen.core.utils.Emoji._
 import za.co.absa.pramen.core.utils.SparkUtils._
 import za.co.absa.pramen.core.utils.hive.HiveHelper
@@ -129,25 +129,35 @@ abstract class TaskRunnerBase(conf: Config,
       spark.sparkContext.setJobDescription(description)
     }
 
-    task.job.operation.killMaxExecutionTimeSeconds match {
-      case Some(timeout) if timeout > 0 =>
-        @volatile var runStatus: RunStatus = null
-
-        try {
-          ThreadUtils.runWithTimeout(Duration(timeout, TimeUnit.SECONDS)) {
-            log.info(s"Running ${task.job.name} with the hard timeout = $timeout seconds.")
-            runStatus = doValidateOrSkipTask(task)
+    try {
+      WorkerStatusManager.setStatus(s"Running '${task.job.name}' for '${task.infoDate}'")
+
+      task.job.operation.killMaxExecutionTimeSeconds match {
+        case Some(timeout) if timeout > 0 =>
+          @volatile var runStatus: RunStatus = null
+
+          try {
+            ThreadUtils.runWithTimeout(Duration(timeout, TimeUnit.SECONDS)) {
+              log.info(s"Running '${task.job.name}' with the hard timeout = $timeout seconds.")
+              try {
+                runStatus = doValidateOrSkipTask(task)
+              } finally {
+                WorkerStatusManager.setFinished()
+              }
+            }
+            runStatus
+          } catch {
+            case NonFatal(ex) =>
+              failTask(task, started, ex)
           }
-          runStatus
-        } catch {
-          case NonFatal(ex) =>
-            failTask(task, started, ex)
-        }
-      case Some(timeout) =>
-        log.error(s"Incorrect timeout for the task: ${task.job.name}. Should be bigger than zero, got: $timeout.")
-        doValidateOrSkipTask(task)
-      case None =>
-        doValidateOrSkipTask(task)
+        case Some(timeout) =>
+          log.error(s"Incorrect timeout for the task: ${task.job.name}. Should be bigger than zero, got: $timeout.")
+          doValidateOrSkipTask(task)
+        case None =>
+          doValidateOrSkipTask(task)
+      }
+    } finally {
+      WorkerStatusManager.setFinished()
     }
   }
 

pramen/core/src/main/scala/za/co/absa/pramen/core/runner/task/TaskRunnerMultithreaded.scala

@@ -26,7 +26,7 @@ import za.co.absa.pramen.core.journal.Journal
 import za.co.absa.pramen.core.lock.TokenLockFactory
 import za.co.absa.pramen.core.pipeline.PipelineDef.{ENVIRONMENT_NAME, PIPELINE_NAME_KEY, TENANT_KEY}
 import za.co.absa.pramen.core.pipeline.{Job, Task}
-import za.co.absa.pramen.core.state.PipelineState
+import za.co.absa.pramen.core.state.{PipelineState, WorkerStatusManager}
 import za.co.absa.pramen.core.utils.{ConfigUtils, Emoji}
 
 import java.util.concurrent.Executors.newFixedThreadPool

pramen/core/src/main/scala/za/co/absa/pramen/core/state/WorkerStatus.scala

@@ -0,0 +1,22 @@
+/*
+ * Copyright 2022 ABSA Group Limited
+ *
+ * Licensed 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.
+ */
+
+package za.co.absa.pramen.core.state
+
+case class WorkerStatus(
+                         threadId: Long,
+                         status: String
+                       )

pramen/core/src/main/scala/za/co/absa/pramen/core/state/WorkerStatusManager.scala

@@ -0,0 +1,37 @@
+/*
+ * Copyright 2022 ABSA Group Limited
+ *
+ * Licensed 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.
+ */
+
+package za.co.absa.pramen.core.state
+
+import scala.collection.mutable
+
+object WorkerStatusManager {
+  private val workerStatus = new mutable.HashMap[Long, String]()
+
+  def setStatus(status: String): Unit = synchronized {
+    val threadId = Thread.currentThread().getId
+    workerStatus(threadId) = status
+  }
+
+  def setFinished(): Unit = synchronized {
+    val threadId = Thread.currentThread().getId
+    workerStatus.remove(threadId)
+  }
+
+  def getStatuses: Seq[WorkerStatus] = synchronized {
+    workerStatus.map { case (threadId, status) => WorkerStatus(threadId, status) }.toSeq
+  }
+}

pramen/core/src/main/scala/za/co/absa/pramen/core/utils/Emoji.scala

@@ -31,4 +31,5 @@ object Emoji {
   val LIGHT_BULB = "\uD83D\uDCA1"
   val STAR = "\u2B50"
   val VOLTAGE = "\u26A1"
+  val HAMMER_AND_WRENCH = "\uD83D\uDEE0\uFE0F"
 }

pramen/core/src/test/scala/za/co/absa/pramen/core/state/WorkerStatusManagerSuite.scala

@@ -0,0 +1,110 @@
+/*
+ * Copyright 2022 ABSA Group Limited
+ *
+ * Licensed 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.
+ */
+
+package za.co.absa.pramen.core.state
+
+import org.scalatest.wordspec.AnyWordSpec
+
+import java.util.concurrent.{CountDownLatch, TimeUnit}
+
+class WorkerStatusManagerSuite extends AnyWordSpec {
+  "setStatus" should {
+    "register the status of the current thread" in {
+      try {
+        WorkerStatusManager.setStatus("running task 1")
+
+        val statuses = WorkerStatusManager.getStatuses
+
+        assert(statuses.nonEmpty)
+        assert(statuses.exists(s => s.threadId == Thread.currentThread().getId && s.status == "running task 1"))
+      } finally {
+        WorkerStatusManager.setFinished()
+      }
+    }
+
+    "overwrite the previous status of the same thread" in {
+      try {
+        WorkerStatusManager.setStatus("running task 1")
+        WorkerStatusManager.setStatus("running task 2")
+
+        val statuses = WorkerStatusManager.getStatuses.filter(_.threadId == Thread.currentThread().getId)
+
+        assert(statuses.length == 1)
+        assert(statuses.head.status == "running task 2")
+      } finally {
+        WorkerStatusManager.setFinished()
+      }
+    }
+  }
+
+  "setFinished" should {
+    "remove the status of the current thread" in {
+      WorkerStatusManager.setStatus("running task 1")
+      WorkerStatusManager.setFinished()
+
+      val statuses = WorkerStatusManager.getStatuses.filter(_.threadId == Thread.currentThread().getId)
+
+      assert(statuses.isEmpty)
+    }
+
+    "do nothing if the current thread has no registered status" in {
+      WorkerStatusManager.setFinished()
+      WorkerStatusManager.setFinished()
+
+      val statuses = WorkerStatusManager.getStatuses.filter(_.threadId == Thread.currentThread().getId)
+
+      assert(statuses.isEmpty)
+    }
+  }
+
+  "getStatuses" should {
+    "return an empty collection when no statuses are registered" in {
+      WorkerStatusManager.setFinished()
+
+      assert(WorkerStatusManager.getStatuses.isEmpty)
+    }
+
+    "return statuses of all running threads" in {
+      val threadsStarted = new CountDownLatch(2)
+      val allowToFinish = new CountDownLatch(1)
+
+      val threads = (1 to 2).map { i =>
+        val thread = new Thread(new Runnable {
+          override def run(): Unit = {
+            WorkerStatusManager.setStatus(s"worker $i")
+            threadsStarted.countDown()
+            allowToFinish.await(10, TimeUnit.SECONDS)
+            WorkerStatusManager.setFinished()
+          }
+        })
+        thread.start()
+        thread
+      }
+
+      threadsStarted.await(10, TimeUnit.SECONDS)
+
+      val statuses = WorkerStatusManager.getStatuses
+
+      allowToFinish.countDown()
+      threads.foreach(_.join(10000))
+
+      assert(statuses.length == 2)
+      assert(statuses.map(_.status).sortBy(identity) == Seq("worker 1", "worker 2"))
+      assert(statuses.map(_.threadId).distinct.length == 2)
+      assert(WorkerStatusManager.getStatuses.isEmpty)
+    }
+  }
+}