From b2338017715c35cc3c8cfbd17806fb4f4225be11 Mon Sep 17 00:00:00 2001
From: Martin Egri <megri@users.noreply.github.com>
Date: Mon, 11 Mar 2024 23:22:24 +0100
Subject: [PATCH] Add 'jvmOptions' to docker contrib configuration (#3079)

Pull request: https://github.com/com-lihaoyi/mill/pull/3079
---
 contrib/docker/readme.adoc                    |  2 ++
 .../mill/contrib/docker/DockerModule.scala    | 15 ++++++++++++-
 .../contrib/docker/DockerModuleTest.scala     | 22 +++++++++++++++++++
 3 files changed, 38 insertions(+), 1 deletion(-)

diff --git a/contrib/docker/readme.adoc b/contrib/docker/readme.adoc
index 449546f9747..b12f8d5327e 100644
--- a/contrib/docker/readme.adoc
+++ b/contrib/docker/readme.adoc
@@ -53,6 +53,8 @@ object docker extends DockerConfig {
   def volumes = Seq("/v1", "/v2")
   // Environment variables to be set in the container (ENV instructions)
   def envVars = Map("foo" -> "bar", "foobar" -> "barfoo")
+  // JVM runtime options such as heap size settings
+  def jvmOptions = Seq("-Xmx1024M", "-XX:+HeapDumpOnOutOfMemoryError")
   // Add RUN instructions
   def run = Seq(
     "/bin/bash -c 'echo Hello World!'",
diff --git a/contrib/docker/src/mill/contrib/docker/DockerModule.scala b/contrib/docker/src/mill/contrib/docker/DockerModule.scala
index af7562f8db9..bc6da6c4f4a 100644
--- a/contrib/docker/src/mill/contrib/docker/DockerModule.scala
+++ b/contrib/docker/src/mill/contrib/docker/DockerModule.scala
@@ -19,6 +19,16 @@ trait DockerModule { outer: JavaModule =>
     def baseImage: T[String] = "gcr.io/distroless/java:latest"
     def pullBaseImage: T[Boolean] = T(baseImage().endsWith(":latest"))
 
+    /**
+     * JVM runtime options. Each item of the Seq should consist of an option and its desired value, like
+     * {{{
+     * def jvmOptions = Seq("-Xmx1024M", "-agentlib:jdwp=transport=dt_socket,server=y,address=8000", …)
+     * }}}
+     * For a full list of options consult the official documentation at
+     * [[https://docs.oracle.com/en/java/javase/21/docs/specs/man/java.html#overview-of-java-options]]
+     */
+    def jvmOptions: T[Seq[String]] = Seq.empty[String]
+
     /**
      * TCP Ports the container will listen to at runtime.
      *
@@ -113,11 +123,14 @@ trait DockerModule { outer: JavaModule =>
         if (user().isEmpty) "" else s"USER ${user()}"
       ).filter(_.nonEmpty).mkString(sys.props("line.separator"))
 
+      val quotedEntryPointArgs = (Seq("java") ++ jvmOptions() ++ Seq("-jar", s"/$jarName"))
+        .map(arg => s"\"$arg\"").mkString(", ")
+
       s"""
          |FROM ${baseImage()}
          |$lines
          |COPY $jarName /$jarName
-         |ENTRYPOINT ["java", "-jar", "/$jarName"]""".stripMargin
+         |ENTRYPOINT [$quotedEntryPointArgs]""".stripMargin
     }
 
     final def build = T {
diff --git a/contrib/docker/test/src/mill/contrib/docker/DockerModuleTest.scala b/contrib/docker/test/src/mill/contrib/docker/DockerModuleTest.scala
index b7db161c5ee..3fe71603886 100644
--- a/contrib/docker/test/src/mill/contrib/docker/DockerModuleTest.scala
+++ b/contrib/docker/test/src/mill/contrib/docker/DockerModuleTest.scala
@@ -36,6 +36,11 @@ object DockerModuleTest extends TestSuite {
       override def user = "user1"
       override def executable = testExecutable
     }
+
+    object dockerJvmOptions extends DockerConfig {
+      override def executable = testExecutable
+      override def jvmOptions = Seq("-Xmx1024M")
+    }
   }
 
   val testArtifactName = "mill-docker-contrib-test"
@@ -132,6 +137,23 @@ object DockerModuleTest extends TestSuite {
         )
         assert(dockerfileStringRefined == expected)
       }
+
+      "extra jvm options" - {
+        val eval = new TestEvaluator(Docker)
+        val Right((dockerfileString, _)) = eval(Docker.dockerJvmOptions.dockerfile)
+        val expected = multineRegex.replaceAllIn(
+          """
+            |FROM gcr.io/distroless/java:latest
+            |COPY out.jar /out.jar
+            |ENTRYPOINT ["java", "-Xmx1024M", "-jar", "/out.jar"]""".stripMargin,
+          sys.props("line.separator")
+        )
+        val dockerfileStringRefined = multineRegex.replaceAllIn(
+          dockerfileString,
+          sys.props("line.separator")
+        )
+        assert(dockerfileStringRefined == expected)
+      }
     }
   }
 }