try to improve machine utilization for batch runs

Author: fhackett

omnilink/.gitignore

@@ -1,3 +1,4 @@
 /*.duckdb
 /*.wal
-/*.sqlite
+/*.sqlite
+/*.sqlite-journal

omnilink/package.mill

@@ -127,6 +127,23 @@ def rr: T[PathRef] = Task:
   )()
 end rr
 
+// Weird trick: because this is all macro expansion, if someone uses Mill ops inside
+// fn, they get expanded after us, so they can't tell we skipped using Task.Anon
+private inline def inDBTransaction[T](
+    inline fn: scalasql.core.DbApi.Txn => T,
+): T =
+  // Using.resource(DBConcurrencyLimiter(concurrencyLimit = 16)): _ =>
+  Using.resource(EvalDB.dbClient()): db =>
+    db.transaction(fn)
+end inDBTransaction
+
+def vacuumDB() = Task.Command(exclusive = true):
+  Using.resource(ConcurrencyLimiter(concurrencyLimit = 1)): _ =>
+    Using.resource(EvalDB.dbClient()): client =>
+      val db = client.getAutoCommitClientConnection
+      db.updateRaw("VACUUM;")
+end vacuumDB
+
 trait EvalDBTrait:
   def dbPath: os.SubPath
   def createTablesSQL: os.SubPath
@@ -235,32 +252,28 @@ private trait ConcurrencyLimiter:
   protected def silent = false
   private var usageCount = 0
 
-  def apply(concurrencyLimit: Int): Handle =
+  def apply(concurrencyLimit: Int, weight: Int = 1): Handle =
     concurrencyLimiter.synchronized:
-      while usageCount >= concurrencyLimit
+      while (usageCount + weight) > concurrencyLimit
       do concurrencyLimiter.wait()
-      usageCount += 1
+      usageCount += weight
       if !silent
       then println("got concurrency token")
-      Handle()
+      Handle(weight)
   end apply
 
-  final class Handle private[ConcurrencyLimiter] () extends AutoCloseable:
+  final class Handle private[ConcurrencyLimiter] (weight: Int)
+      extends AutoCloseable:
     def close(): Unit =
       concurrencyLimiter.synchronized:
         if !silent
         then println("giving up concurrency token")
         concurrencyLimiter.notifyAll()
-        usageCount -= 1
+        usageCount -= weight
   end Handle
 end ConcurrencyLimiter
 
-private object ValidateConcurrencyLimiter extends ConcurrencyLimiter
-private object EnvConcurrencyLimiter extends ConcurrencyLimiter
-private object GatherTraceConcurrencyLimiter extends ConcurrencyLimiter
-private object DBConcurrencyLimiter extends ConcurrencyLimiter:
-  override protected def silent = true
-end DBConcurrencyLimiter
+private object ConcurrencyLimiter extends ConcurrencyLimiter
 
 trait TracingConfigModule extends Module:
   tracingConfig =>
@@ -289,16 +302,6 @@ trait TracingConfigModule extends Module:
           then os.copy.over(cmdsFile, dir / "compile_commands.json")
   end installCompileCommands
 
-  // Weird trick: because this is all macro expansion, if someone uses Mill ops inside
-  // fn, they get expanded after us, so they can't tell we skipped using Task.Anon
-  private inline def inDBTransaction[T](
-      inline fn: scalasql.core.DbApi.Txn => T,
-  ): T =
-    Using.resource(DBConcurrencyLimiter(concurrencyLimit = 16)): _ =>
-      Using.resource(EvalDB.dbClient()): db =>
-        db.transaction(fn)
-  end inDBTransaction
-
   trait ConfigModule extends Module:
     configModule =>
     def threadCount: Int
@@ -446,39 +449,40 @@ trait TracingConfigModule extends Module:
       end unpackTrace
 
       def unpackValidationEnv = Task[PathRef]:
-        Using.resource(EnvConcurrencyLimiter(concurrencyLimit = 4)): _ =>
-          val tracesFolder = unpackTrace().path
-          val spec = specToValidate().path
-          val mcSpec = specToValidateMC().path
-          val mcConfig = specToValidateMCConfig().path
-
-          build.omnilink.tool
-            .runner()
-            .run(
-              List[os.Shellable](
-                "gen-tla",
-                spec,
-                tracesFolder,
-                "--dest-dir",
-                Task.dest,
-              ),
-            )
+        Using.resource(ConcurrencyLimiter(concurrencyLimit = 12, weight = 1)):
+          _ =>
+            val tracesFolder = unpackTrace().path
+            val spec = specToValidate().path
+            val mcSpec = specToValidateMC().path
+            val mcConfig = specToValidateMCConfig().path
+
+            build.omnilink.tool
+              .runner()
+              .run(
+                List[os.Shellable](
+                  "gen-tla",
+                  spec,
+                  tracesFolder,
+                  "--dest-dir",
+                  Task.dest,
+                ),
+              )
 
-          os.copy(from = mcSpec, to = Task.dest / mcSpec.last)
-          os.copy(from = mcConfig, to = Task.dest / mcConfig.last)
-
-          build.omnilink.tool
-            .runner()
-            .run(
-              List[os.Shellable](
-                "showlog",
-                "porcupine",
-                tracesFolder,
-                Task.dest / "porcupine_ops.json",
-              ),
-            )
+            os.copy(from = mcSpec, to = Task.dest / mcSpec.last)
+            os.copy(from = mcConfig, to = Task.dest / mcConfig.last)
+
+            build.omnilink.tool
+              .runner()
+              .run(
+                List[os.Shellable](
+                  "showlog",
+                  "porcupine",
+                  tracesFolder,
+                  Task.dest / "porcupine_ops.json",
+                ),
+              )
 
-          PathRef(Task.dest)
+            PathRef(Task.dest)
       end unpackValidationEnv
 
       def discardTrace() = Task.Command(exclusive = true)[Unit]:
@@ -543,7 +547,7 @@ trait TracingConfigModule extends Module:
       end discardNegativeValidation
 
       def gatherTrace() = Task.Command[Unit]:
-        Using.resource(GatherTraceConcurrencyLimiter(concurrencyLimit = 6)):
+        Using.resource(ConcurrencyLimiter(concurrencyLimit = 24, weight = 1)):
           _ =>
             val exe = tracingExecutable()
 
@@ -646,10 +650,15 @@ trait TracingConfigModule extends Module:
               ()
       end gatherTrace
 
-      def concurrencyLimit = Task.Input:
-        Task.ctx().env.get("TLC_CONCURRENCY").fold(1)(_.toInt)
+      def validateTrace() = Task.Command(exclusive = true):
+        validateTraceImpl(concurrencyLimit = 1)()
+      end validateTrace
+
+      def validateTracePar(concurrencyLimit: Int) = Task.Command:
+        validateTraceImpl(concurrencyLimit)()
+      end validateTracePar
 
-      def validateTrace() = Task.Command(exclusive = true)[Unit]:
+      private def validateTraceImpl(concurrencyLimit: Int) = Task.Anon[Unit]:
         if inDBTransaction: db =>
             db.run:
               EvalDB.Validation.select
@@ -665,7 +674,13 @@ trait TracingConfigModule extends Module:
                 .size === 0
         then println(s"no traces to validate for $traceId")
         else
-          Using.resource(ValidateConcurrencyLimiter(concurrencyLimit())): _ =>
+          val weight = 16
+          Using.resource(
+            ConcurrencyLimiter(
+              concurrencyLimit = concurrencyLimit * weight,
+              weight = weight,
+            ),
+          ): _ =>
             val validationDir = unpackValidationEnv().path
             val timeFile = Task.dest / "timeinfo.txt"
             os.copy.over(
@@ -758,7 +773,7 @@ trait TracingConfigModule extends Module:
                     _.peakMemory := peakMemory,
                   )
             ()
-      end validateTrace
+      end validateTraceImpl
 
       def porcupineValidateTrace() = Task.Command(exclusive = true)[Unit]:
         require(
@@ -773,7 +788,7 @@ trait TracingConfigModule extends Module:
                 .size > 0
         then println(s"already validated trace $traceId")
         else
-          Using.resource(ValidateConcurrencyLimiter(concurrencyLimit())): _ =>
+          Using.resource(ConcurrencyLimiter(concurrencyLimit = 1)): _ =>
             val validationDir = unpackValidationEnv().path
             val timeFile = Task.dest / "timeinfo.txt"
             val startTime = LocalDateTime.now()