[joern-slice] Fixed Parallelism

Following #4009, I have replaced the parallelism in the slicing with this utility.
joernio · Jan 9, 2024 · 92a5345 · 92a5345
1 parent 10f3cb3
commit 92a5345
Show file tree

Hide file tree

Showing 2 changed files with 42 additions and 23 deletions.
diff --git a/dataflowengineoss/src/main/scala/io/joern/dataflowengineoss/slicing/DataFlowSlicing.scala b/dataflowengineoss/src/main/scala/io/joern/dataflowengineoss/slicing/DataFlowSlicing.scala
@@ -1,27 +1,40 @@
 package io.joern.dataflowengineoss.slicing
 
 import io.joern.dataflowengineoss.language.*
+import io.joern.x2cpg.utils.ConcurrentTaskExecutionUtil
 import io.shiftleft.codepropertygraph.Cpg
 import io.shiftleft.codepropertygraph.generated.PropertyNames
 import io.shiftleft.codepropertygraph.generated.nodes.*
 import io.shiftleft.semanticcpg.language.*
+import org.slf4j.LoggerFactory
 
-import java.util.concurrent.{Callable, Executors}
+import java.util.concurrent.Callable
+import scala.concurrent.ExecutionContext
+import scala.util.{Failure, Success}
 
 object DataFlowSlicing {
 
   implicit val resolver: ICallResolver = NoResolve
+  private val logger                   = LoggerFactory.getLogger(getClass)
 
   def calculateDataFlowSlice(cpg: Cpg, config: DataFlowConfig): Option[DataFlowSlice] = {
     implicit val implicitConfig: DataFlowConfig = config
 
-    val exec = poolFromConfig(config)
-    (config.fileFilter match {
+    val tasks = (config.fileFilter match {
       case Some(fileName) => cpg.file.nameExact(fileName).method.call
       case None           => cpg.call
     }).method.withMethodNameFilter.withMethodParameterFilter.withMethodAnnotationFilter.call.withExternalCalleeFilter
-      .map(c => exec.submit(new TrackDataFlowTask(config, c)))
-      .flatMap(_.get())
+      .map(c => () => new TrackDataFlowTask(config, c).call())
+      .iterator
+
+    ConcurrentTaskExecutionUtil
+      .runInParallel(tasks)(ExecutionContext.fromExecutor(poolFromConfig(config)))
+      .flatMap {
+        case Success(slice) => slice
+        case Failure(e) =>
+          logger.warn("Exception encountered during slicing task", e)
+          None
+      }
       .reduceOption { (a, b) => DataFlowSlice(a.nodes ++ b.nodes, a.edges ++ b.edges) }
   }
 

diff --git a/dataflowengineoss/src/main/scala/io/joern/dataflowengineoss/slicing/UsageSlicing.scala b/dataflowengineoss/src/main/scala/io/joern/dataflowengineoss/slicing/UsageSlicing.scala
@@ -1,21 +1,25 @@
 package io.joern.dataflowengineoss.slicing
 
+import io.joern.x2cpg.utils.ConcurrentTaskExecutionUtil
 import io.shiftleft.codepropertygraph.Cpg
 import io.shiftleft.codepropertygraph.generated.nodes.*
 import io.shiftleft.codepropertygraph.generated.{Operators, PropertyNames}
 import io.shiftleft.semanticcpg.language.*
+import org.slf4j.LoggerFactory
 
 import java.util.concurrent.*
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.regex.Pattern
 import scala.collection.concurrent.TrieMap
-import scala.util.Try
+import scala.concurrent.{ExecutionContext, ExecutionContextExecutor}
+import scala.util.{Failure, Success, Try}
 
 /** A utility for slicing based off of usage references for identifiers and parameters. This is mainly tested around
   * JavaScript CPGs.
   */
 object UsageSlicing {
 
+  private val logger                 = LoggerFactory.getLogger(getClass)
   private val resolver               = NoResolve
   private val constructorTypeMatcher = Pattern.compile(".*new (\\w+)\\(.*")
   private val excludeOperatorCalls   = new AtomicBoolean(false)
@@ -38,30 +42,32 @@ object UsageSlicing {
 
     def typeMap = TrieMap.from(cpg.typeDecl.map(f => (f.name, f.fullName)).toMap)
 
-    val exec = poolFromConfig(config)
-    try {
-      val slices       = usageSlices(exec, cpg, declarations, typeMap)
-      val userDefTypes = userDefinedTypes(cpg)
-      ProgramUsageSlice(slices, userDefTypes)
-    } finally {
-      exec.shutdown()
-    }
+    implicit val exec: ExecutionContextExecutor = ExecutionContext.fromExecutor(poolFromConfig(config))
+    val slices                                  = usageSlices(cpg, declarations, typeMap)
+    val userDefTypes                            = userDefinedTypes(cpg)
+    ProgramUsageSlice(slices, userDefTypes)
   }
 
   import io.shiftleft.semanticcpg.codedumper.CodeDumper.dump
 
-  private def usageSlices(
-    exec: ExecutorService,
-    cpg: Cpg,
-    declarations: List[Declaration],
-    typeMap: TrieMap[String, String]
-  )(implicit config: UsagesConfig): List[MethodUsageSlice] = {
+  private def usageSlices(cpg: Cpg, declarations: List[Declaration], typeMap: TrieMap[String, String])(implicit
+    config: UsagesConfig,
+    exec: ExecutionContextExecutor
+  ): List[MethodUsageSlice] = {
     val language = cpg.metaData.language.headOption
     val root     = cpg.metaData.root.headOption
-    declarations
+    val tasks = declarations
       .filter(a => atLeastNCalls(a, config.minNumCalls) && !a.name.startsWith("_tmp_"))
-      .map(a => exec.submit(new TrackUsageTask(cpg, a, typeMap)))
-      .flatMap(_.get)
+      .map(a => () => new TrackUsageTask(cpg, a, typeMap).call())
+      .iterator
+    ConcurrentTaskExecutionUtil
+      .runInParallel(tasks)
+      .flatMap {
+        case Success(slice) => slice
+        case Failure(e) =>
+          logger.warn("Exception encountered during slicing task", e)
+          None
+      }
       .groupBy { case (scope, _) => scope }
       .view
       .sortBy(_._1.fullName)