From a1f632e626b59a19d2213c06e0f770362c904850 Mon Sep 17 00:00:00 2001
From: Karan Batavia <118820668+karan-batavia@users.noreply.github.com>
Date: Tue, 18 Jun 2024 14:04:38 +0530
Subject: [PATCH] Merge pull request #3 from Privado-Inc/stats-logging

Merge `stats-logging` to `master`
---
 .../scala/io/shiftleft/passes/CpgPass.scala   |  3 ++
 .../io/shiftleft/passes/ParallelCpgPass.scala |  4 ++-
 .../io/shiftleft/utils/StatsCollector.scala   | 28 +++++++++++++++++++
 3 files changed, 34 insertions(+), 1 deletion(-)
 create mode 100644 codepropertygraph/src/main/scala/io/shiftleft/utils/StatsCollector.scala

diff --git a/codepropertygraph/src/main/scala/io/shiftleft/passes/CpgPass.scala b/codepropertygraph/src/main/scala/io/shiftleft/passes/CpgPass.scala
index faeca0ba6..14903c46e 100644
--- a/codepropertygraph/src/main/scala/io/shiftleft/passes/CpgPass.scala
+++ b/codepropertygraph/src/main/scala/io/shiftleft/passes/CpgPass.scala
@@ -3,6 +3,7 @@ package io.shiftleft.passes
 import com.google.protobuf.GeneratedMessageV3
 import io.shiftleft.SerializedCpg
 import io.shiftleft.codepropertygraph.Cpg
+import io.shiftleft.utils.StatsLogger
 import org.slf4j.{Logger, LoggerFactory, MDC}
 import overflowdb.BatchedUpdate
 
@@ -62,6 +63,7 @@ abstract class ForkJoinParallelCpgPass[T <: AnyRef](
 
   override def createApplySerializeAndStore(serializedCpg: SerializedCpg, prefix: String = ""): Unit = {
     baseLogger.info(s"Start of pass: $name")
+    StatsLogger.initiateNewStage(getClass.getSimpleName, Some(name), getClass.getSuperclass.getSimpleName)
     val nanosStart = System.nanoTime()
     var nParts     = 0
     var nanosBuilt = -1L
@@ -95,6 +97,7 @@ abstract class ForkJoinParallelCpgPass[T <: AnyRef](
         baseLogger.info(
           f"Pass $name completed in ${(nanosStop - nanosStart) * 1e-6}%.0f ms (${fracRun}%.0f%% on mutations). ${nDiff}%d + ${nDiffT - nDiff}%d changes committed from ${nParts}%d parts.${serializationString}%s"
         )
+        StatsLogger.endLastStage()
       }
     }
   }
diff --git a/codepropertygraph/src/main/scala/io/shiftleft/passes/ParallelCpgPass.scala b/codepropertygraph/src/main/scala/io/shiftleft/passes/ParallelCpgPass.scala
index ff0a5b9f3..f5d37b9ed 100644
--- a/codepropertygraph/src/main/scala/io/shiftleft/passes/ParallelCpgPass.scala
+++ b/codepropertygraph/src/main/scala/io/shiftleft/passes/ParallelCpgPass.scala
@@ -1,7 +1,7 @@
 package io.shiftleft.passes
 import io.shiftleft.SerializedCpg
 import io.shiftleft.codepropertygraph.Cpg
-import io.shiftleft.utils.ExecutionContextProvider
+import io.shiftleft.utils.{ExecutionContextProvider, StatsLogger}
 import org.slf4j.MDC
 
 import java.util.concurrent.LinkedBlockingQueue
@@ -54,6 +54,7 @@ abstract class ConcurrentWriterCpgPass[T <: AnyRef](
   override def createApplySerializeAndStore(serializedCpg: SerializedCpg, prefix: String = ""): Unit = {
     import ConcurrentWriterCpgPass.producerQueueCapacity
     baseLogger.info(s"Start of enhancement: $name")
+    StatsLogger.initiateNewStage(getClass.getSimpleName, Some(name), getClass.getSuperclass.getSimpleName)
     val nanosStart = System.nanoTime()
     var nParts     = 0
     var nDiff      = 0
@@ -120,6 +121,7 @@ abstract class ConcurrentWriterCpgPass[T <: AnyRef](
       baseLogger.info(
         f"Enhancement $name completed in ${(nanosStop - nanosStart) * 1e-6}%.0f ms. ${nDiff}%d  + ${nDiffT - nDiff}%d changes committed from ${nParts}%d parts."
       )
+      StatsLogger.endLastStage()
     }
   }
 
diff --git a/codepropertygraph/src/main/scala/io/shiftleft/utils/StatsCollector.scala b/codepropertygraph/src/main/scala/io/shiftleft/utils/StatsCollector.scala
new file mode 100644
index 000000000..fb8cf1600
--- /dev/null
+++ b/codepropertygraph/src/main/scala/io/shiftleft/utils/StatsCollector.scala
@@ -0,0 +1,28 @@
+package io.shiftleft.utils
+
+object StatsLogger extends DataLogger {
+  private var logger: Option[DataLogger] = None
+  def initialise(logger: Option[DataLogger] = None): Unit = {
+    this.logger = logger
+  }
+
+  def initiateNewStage(
+    stageName: String,
+    stageFullName: Option[String] = None,
+    additionalMetaDataToLog: String = STAGE_NOT_SET
+  ): Unit = { logger.foreach(log => log.initiateNewStage(stageName, stageFullName, additionalMetaDataToLog)) }
+
+  def endLastStage(): Unit = { logger.foreach(log => log.endLastStage()) }
+}
+
+trait DataLogger {
+  // Constant used when stage or subStage is not set.
+  val STAGE_NOT_SET = "<not set>"
+  def initiateNewStage(
+    stageName: String,
+    stageFullName: Option[String] = None,
+    additionalMetaDataToLog: String = STAGE_NOT_SET
+  ): Unit
+
+  def endLastStage(): Unit
+}