[DL-2086] fix: improve docker image pull process (#32)

Thomas Franco · web-flow · commit fca9efe4ac4e · 2025-05-13T10:09:23.000+02:00
* [DL-2086] fix: improve  docker image pull process

* [DL-2086] fix: Alex's feedback
diff --git a/app/org/thp/cortex/services/DockerJobRunnerSrv.scala b/app/org/thp/cortex/services/DockerJobRunnerSrv.scala
@@ -1,7 +1,7 @@
 package org.thp.cortex.services
 
 import akka.actor.ActorSystem
-import org.thp.cortex.util.docker.{DockerClient => DockerJavaClient}
+import org.thp.cortex.util.docker.DockerClient
 import play.api.libs.json.Json
 import play.api.{Configuration, Logger}
 
@@ -15,7 +15,7 @@ import scala.util.Try
 
 @Singleton
 class DockerJobRunnerSrv(
-    javaClient: DockerJavaClient,
+    dockerClient: DockerClient,
     autoUpdate: Boolean,
     jobBaseDirectory: Path,
     dockerJobBaseDirectory: Path,
@@ -25,7 +25,7 @@ class DockerJobRunnerSrv(
   @Inject()
   def this(config: Configuration, system: ActorSystem) =
     this(
-      new DockerJavaClient(config),
+      new DockerClient(config),
       config.getOptional[Boolean]("docker.autoUpdate").getOrElse(true),
       Paths.get(config.get[String]("job.directory")),
       Paths.get(config.get[String]("job.dockerDirectory")),
@@ -37,7 +37,7 @@ class DockerJobRunnerSrv(
   lazy val isAvailable: Boolean =
     Try {
       logger.debug(s"Retrieve docker information ...")
-      logger.info(s"Docker is available:\n${javaClient.info}")
+      logger.info(s"Docker is available:\n${dockerClient.info}")
       true
     }.recover {
       case error =>
@@ -48,7 +48,7 @@ class DockerJobRunnerSrv(
   private def generateErrorOutput(containerId: String, f: Path) = {
     logger.warn(s"the runner didn't generate any output file $f")
     for {
-      output <- javaClient.getLogs(containerId)
+      output <- dockerClient.getLogs(containerId)
       report = Json.obj("success" -> false, "errorMessage" -> output)
       _ <- Try(Files.write(f, report.toString.getBytes(StandardCharsets.UTF_8)))
     } yield report
@@ -57,22 +57,22 @@ class DockerJobRunnerSrv(
   def run(jobDirectory: Path, dockerImage: String, timeout: Option[FiniteDuration])(implicit executionContext: ExecutionContext): Try[Unit] = {
     val to = timeout.getOrElse(FiniteDuration(5000, TimeUnit.SECONDS))
 
-    if (autoUpdate) Try(javaClient.pullImage(dockerImage))
+    if (autoUpdate) dockerClient.pullImage(dockerImage)
 
     for {
-      containerId <- javaClient.prepare(dockerImage, jobDirectory, jobBaseDirectory, dockerJobBaseDirectory, to)
+      containerId <- dockerClient.prepare(dockerImage, jobDirectory, jobBaseDirectory, dockerJobBaseDirectory, to)
       timeoutScheduled = timeout.map(to =>
         system.scheduler.scheduleOnce(to) {
           logger.info("Timeout reached, stopping the container")
-          javaClient.clean(containerId)
+          dockerClient.clean(containerId)
         }
       )
-      _ <- javaClient.execute(containerId)
+      _ <- dockerClient.execute(containerId)
       _ = timeoutScheduled.foreach(_.cancel())
       outputFile <- Try(jobDirectory.resolve("output").resolve("output.json"))
       isError = Files.notExists(outputFile) || Files.size(outputFile) == 0 || Files.isDirectory(outputFile)
       _       = if (isError) generateErrorOutput(containerId, outputFile).toOption else None
-      _ <- javaClient.clean(containerId)
+      _ <- dockerClient.clean(containerId)
     } yield ()
   }
 
diff --git a/app/org/thp/cortex/services/JobRunnerSrv.scala b/app/org/thp/cortex/services/JobRunnerSrv.scala
@@ -49,7 +49,7 @@ class JobRunnerSrv @Inject() (
     .map(_.toLowerCase)
     .collect {
       case "kubernetes" if k8sJobRunnerSrv.isAvailable => "kubernetes"
-      case "docker" if dockerJobRunnerSrv.isAvailable => "docker"
+      case "docker" if dockerJobRunnerSrv.isAvailable  => "docker"
       case "process" =>
         Seq("", "2", "3").foreach { pythonVersion =>
           val cortexUtilsVersion = processJobRunnerSrv.checkCortexUtilsVersion(pythonVersion)
@@ -229,7 +229,7 @@ class JobRunnerSrv @Inject() (
       syncStartJob(job).get
       val jobFolder = prepareJobFolder(worker, job).get
       maybeJobFolder = Some(jobFolder)
-      runners
+      val result = runners
         .foldLeft[Option[Try[Unit]]](None) {
           case (None, "kubernetes") =>
             worker
@@ -261,12 +261,15 @@ class JobRunnerSrv @Inject() (
             None
 
         }
-        .getOrElse(throw BadRequestError("Worker cannot be run"))
+        .getOrElse(Failure(BadRequestError("Worker cannot be run")))
+      Future.fromTry(result)
     }(executionContext)
+      .flatten
       .transformWith {
         case _: Success[_] =>
           extractReport(maybeJobFolder.get /* can't be none */, job)
         case Failure(error) =>
+          logger.error(s"Worker execution fails", error)
           endJob(job, JobStatus.Failure, Option(error.getMessage), maybeJobFolder.map(jf => readFile(jf.resolve("input").resolve("input.json"))))
 
       }
diff --git a/app/org/thp/cortex/util/docker/DockerClient.scala b/app/org/thp/cortex/util/docker/DockerClient.scala
@@ -8,11 +8,12 @@ import play.api.{Configuration, Logger}
 
 import java.nio.file.{Files, Path}
 import java.time.Duration
+import java.util.Collections
 import java.util.concurrent.{Executors, TimeUnit}
 import scala.concurrent.blocking
 import scala.concurrent.duration.FiniteDuration
 import scala.jdk.CollectionConverters._
-import scala.util.Try
+import scala.util.{Failure, Success, Try}
 
 class DockerClient(config: Configuration) {
   private lazy val logger: Logger           = Logger(getClass.getName)
@@ -32,25 +33,33 @@ class DockerClient(config: Configuration) {
       waitResult
     }
 
-  def prepare(image: String, jobDirectory: Path, jobBaseDirectory: Path, dockerJobBaseDirectory: Path, timeout: FiniteDuration): Try[String] = Try {
-    logger.info(s"image $image pull result: ${pullImage(image)}")
-    val containerCmd = underlyingClient
-      .createContainerCmd(image)
-      .withHostConfig(configure(jobDirectory, jobBaseDirectory, dockerJobBaseDirectory))
-    if (Files.exists(jobDirectory.resolve("input").resolve("cacerts")))
-      containerCmd.withEnv(s"REQUESTS_CA_BUNDLE=/job/input/cacerts")
-    val containerResponse = containerCmd.exec()
-    logger.info(
-      s"about to start container ${containerResponse.getId}\n" +
-        s"  timeout: ${timeout.toString}\n" +
-        s"  image  : $image\n" +
-        s"  volumes : ${jobDirectory.toAbsolutePath}"
-    )
-    if (containerResponse.getWarnings.nonEmpty) logger.warn(s"${containerResponse.getWarnings.mkString(", ")}")
-    scheduleContainerTimeout(containerResponse.getId, timeout)
-
-    containerResponse.getId
-  }
+  def pullImageIfRequired(image: String): Try[Unit] =
+    if (imageExists(image)) Success(())
+    else {
+      logger.info(s"pulling image $image ...")
+      pullImage(image)
+    }
+
+  def prepare(image: String, jobDirectory: Path, jobBaseDirectory: Path, dockerJobBaseDirectory: Path, timeout: FiniteDuration): Try[String] =
+    pullImageIfRequired(image)
+      .map { _ =>
+        val containerCmd = underlyingClient
+          .createContainerCmd(image)
+          .withHostConfig(configure(jobDirectory, jobBaseDirectory, dockerJobBaseDirectory))
+        if (Files.exists(jobDirectory.resolve("input").resolve("cacerts")))
+          containerCmd.withEnv(s"REQUESTS_CA_BUNDLE=/job/input/cacerts")
+        val containerResponse = containerCmd.exec()
+        logger.info(
+          s"about to start container ${containerResponse.getId}\n" +
+            s"  timeout: ${timeout.toString}\n" +
+            s"  image  : $image\n" +
+            s"  volumes : ${jobDirectory.toAbsolutePath}"
+        )
+        if (containerResponse.getWarnings.nonEmpty) logger.warn(s"${containerResponse.getWarnings.mkString(", ")}")
+        scheduleContainerTimeout(containerResponse.getId, timeout)
+
+        containerResponse.getId
+      }
 
   private def configure(jobDirectory: Path, jobBaseDirectory: Path, dockerJobBaseDirectory: Path): HostConfig = {
     val hostConfigMut = HostConfig
@@ -85,27 +94,67 @@ class DockerClient(config: Configuration) {
   }
 
   def info: Info = underlyingClient.infoCmd().exec()
-  def pullImage(image: String): Boolean = blocking {
-    val pullImageResultCbk = underlyingClient // Blocking
-      .pullImageCmd(image)
-      .start()
-      .awaitCompletion()
-    val timeout = config.get[FiniteDuration]("docker.pullImageTimeout")
 
-    pullImageResultCbk.awaitCompletion(timeout.toMillis, TimeUnit.MILLISECONDS)
+  def pullImage(image: String): Try[Unit] = Try {
+    blocking {
+      val pullImageResultCbk = underlyingClient // Blocking
+        .pullImageCmd(image)
+        .start()
+      val timeout = config.get[FiniteDuration]("docker.pullImageTimeout")
+
+      pullImageResultCbk.awaitCompletion(timeout.toMillis, TimeUnit.MILLISECONDS)
+      ()
+    }
   }
 
-  def clean(containerId: String): Try[Unit] = Try {
-    underlyingClient
-      .killContainerCmd(containerId)
-      .exec()
+  def imageExists(image: String): Boolean =
+    !underlyingClient.listImagesCmd().withImageNameFilter(image).exec().isEmpty
+
+  def getContainerStatus(containerId: String): Option[String] =
     underlyingClient
-      .removeContainerCmd(containerId)
-      .withForce(true)
+      .listContainersCmd()
+      .withIdFilter(Collections.singletonList(containerId))
       .exec()
-    logger.info(s"removed container $containerId")
+      .asScala
+      .headOption
+      .map(_.getStatus)
+
+  def killContainer(containerId: String): Try[Unit] = {
+    logger.info(s"Killing the container $containerId")
+    Try {
+      underlyingClient
+        .killContainerCmd(containerId)
+        .exec()
+      ()
+    }.recoverWith {
+      case error =>
+        logger.warn(s"Unable to kill the container $containerId", error)
+        Failure(error)
+    }
   }
 
+  def removeContainer(containerId: String): Try[Unit] = {
+    logger.info(s"Removing the container $containerId")
+    Try {
+      underlyingClient
+        .removeContainerCmd(containerId)
+        .withForce(true)
+        .exec()
+      ()
+    }.recoverWith {
+      case error =>
+        logger.warn(s"Unable to remove the container $containerId", error)
+        Failure(error)
+    }
+  }
+
+  def clean(containerId: String): Try[Unit] =
+    getContainerStatus(containerId).fold[Try[Unit]](Success(())) { status =>
+      if (status != "exited" && status != "dead")
+        killContainer(containerId)
+      removeContainer(containerId)
+    }
+
   def getLogs(containerId: String): Try[String] = Try {
     val stringBuilder = new StringBuilder()
     val callback      = new DockerLogsStringBuilder(stringBuilder)
diff --git a/conf/reference.conf b/conf/reference.conf
@@ -124,3 +124,4 @@ responder {
   }
 }
 
+docker.pullImageTimeout = 10 minutes

Original file line number	Diff line number	Diff line change
`@@ -124,3 +124,4 @@ responder {`
`124`	`124`	`}`
`125`	`125`	`}`
`126`	`126`
	`127`	`+docker.pullImageTimeout = 10 minutes`