serve: add --requestTimeout and --warmupRevision controls

README.md

@@ -466,13 +466,15 @@ Usage: bazel-diff serve [-hkvV] [--[no-]excludeExternalTargets]
                         [--fineGrainedHashExternalReposFile=<fineGrainedHashExte
                         rnalReposFile>] [--gitEngine=<gitEngine>]
                         [--gitPath=<gitPath>] [--port=<port>]
+                        [--requestTimeout=<requestTimeoutSeconds>]
                         [-s=<seedFilepaths>] -w=<workspacePath>
                         [-co=<bazelCommandOptions>]...
                         [--cqueryCommandOptions=<cqueryCommandOptions>]...
                         [--fineGrainedHashExternalRepos=<fineGrainedHashExternal
                         Repos>]...
                         [--ignoredRuleHashingAttributes=<ignoredRuleHashingAttri
                         butes>]... [-so=<bazelStartupOptions>]...
+                        [--warmupRevision=<warmupRevisions>]...
 Runs bazel-diff as a long-running HTTP query service that returns the impacted
 targets between two git revisions, caching generated hashes per commit SHA.
   -b, --bazelPath=<bazelPath>
@@ -522,6 +524,14 @@ targets between two git revisions, caching generated hashes per commit SHA.
       --no-initial-fetch    Skip the initial 'git fetch' before reporting
                               healthy. Useful for local/offline testing.
       --port=<port>         Port to listen on. Defaults to 8080.
+      --requestTimeout=<requestTimeoutSeconds>
+                            Maximum seconds an /impacted_targets
+                              (_with_distances) request may run before the
+                              server abandons it and responds 504. 0 (the
+                              default) means no timeout. This bounds the
+                              request the client waits on; an in-flight bazel
+                              query may keep running in the background and
+                              still populate the per-SHA cache.
   -s, --seed-filepaths=<seedFilepaths>
                             A text file with a newline separated list of
                               filepaths used as a SHA256 seed for all targets.
@@ -539,6 +549,16 @@ targets between two git revisions, caching generated hashes per commit SHA.
   -w, --workspacePath=<workspacePath>
                             Path to the Bazel workspace git clone the service
                               checks out and queries.
+      --warmupRevision=<warmupRevisions>
+                            Comma separated git revisions (branch/tag/SHA)
+                              whose hashes are generated and cached at startup,
+                              before the server reports healthy, so the first
+                              real request is warm and the Bazel analysis
+                              server is primed. Best-effort: a revision that
+                              fails to warm is logged and the server still
+                              becomes ready (serving it cold on demand).
+                              Increases time-to-healthy, so size
+                              deploy/health-check timeouts accordingly.
 ```
 <!-- END_SECTION: cli-help -->
 

cli/src/main/kotlin/com/bazel_diff/cli/ServeCommand.kt

@@ -11,6 +11,7 @@ import com.bazel_diff.hash.sha256
 import com.bazel_diff.server.BazelDiffServer
 import com.bazel_diff.server.GitClient
 import com.bazel_diff.server.HashCacheStorage
+import com.bazel_diff.server.HashProvider
 import com.bazel_diff.server.HashService
 import com.bazel_diff.server.ImpactedTargetsService
 import com.bazel_diff.server.JGitClient
@@ -87,6 +88,17 @@ class ServeCommand : Callable<Int> {
       defaultValue = "8080")
   var port: Int = 8080
 
+  @CommandLine.Option(
+      names = ["--requestTimeout"],
+      description =
+          [
+              "Maximum seconds an /impacted_targets(_with_distances) request may run before the " +
+                  "server abandons it and responds 504. 0 (the default) means no timeout. This " +
+                  "bounds the request the client waits on; an in-flight bazel query may keep " +
+                  "running in the background and still populate the per-SHA cache."],
+      defaultValue = "0")
+  var requestTimeoutSeconds: Long = 0
+
   @CommandLine.Option(
       names = ["--cacheDir"],
       description =
@@ -95,6 +107,18 @@ class ServeCommand : Callable<Int> {
       converter = [NormalisingPathConverter::class])
   lateinit var cacheDir: Path
 
+  @CommandLine.Option(
+      names = ["--warmupRevision"],
+      description =
+          [
+              "Comma separated git revisions (branch/tag/SHA) whose hashes are generated and cached " +
+                  "at startup, before the server reports healthy, so the first real request is warm " +
+                  "and the Bazel analysis server is primed. Best-effort: a revision that fails to " +
+                  "warm is logged and the server still becomes ready (serving it cold on demand). " +
+                  "Increases time-to-healthy, so size deploy/health-check timeouts accordingly."],
+      converter = [CommaSeparatedValueConverter::class])
+  var warmupRevisions: Set<String> = emptySet()
+
   @CommandLine.Option(
       names = ["--no-initial-fetch"],
       description =
@@ -264,9 +288,9 @@ class ServeCommand : Callable<Int> {
         ImpactedTargetsService(gitClient, hashService, depsTracked = trackDeps)
 
     val ready = AtomicBoolean(false)
-    val server = BazelDiffServer(port, impactedTargetsService) { ready.get() }
+    val server = BazelDiffServer(port, impactedTargetsService, requestTimeoutSeconds) { ready.get() }
     server.start()
-    performInitialFetch(gitClient, ready, server)
+    performInitialFetch(gitClient, hashService, ready, server)
     return server
   }
 
@@ -276,13 +300,22 @@ class ServeCommand : Callable<Int> {
    * so health checks report `503` and a load balancer removes the instance rather than us
    * attempting a risky in-place repair (RFC issue #29).
    */
-  fun performInitialFetch(gitClient: GitClient, ready: AtomicBoolean, server: BazelDiffServer) {
+  fun performInitialFetch(
+      gitClient: GitClient,
+      hashProvider: HashProvider,
+      ready: AtomicBoolean,
+      server: BazelDiffServer
+  ) {
     if (noInitialFetch) {
+      warmUpCache(gitClient, hashProvider)
       ready.set(true)
       return
     }
     try {
       gitClient.fetch()
+      // Warm up before flipping ready so the load balancer only routes to this instance once its
+      // configured baseline revisions are cached and the Bazel server is primed.
+      warmUpCache(gitClient, hashProvider)
       ready.set(true)
       System.err.println("[Info] initial git fetch complete; serving on port ${server.boundPort()}")
     } catch (e: Exception) {
@@ -291,6 +324,28 @@ class ServeCommand : Callable<Int> {
     }
   }
 
+  /**
+   * Best-effort priming of the hash cache (and, as a side effect, the Bazel analysis server) for
+   * each `--warmupRevision` before readiness is reported. A revision that fails to warm (bad ref,
+   * transient bazel error) is logged and skipped rather than failing the deploy: warmup is an
+   * optimization, not a correctness requirement, so the server still becomes ready and serves that
+   * revision cold on demand. Never throws, so a bad `--warmupRevision` cannot lame-duck the
+   * instance.
+   */
+  fun warmUpCache(gitClient: GitClient, hashProvider: HashProvider) {
+    for (rev in warmupRevisions) {
+      try {
+        val sha = gitClient.resolveSha(rev)
+        System.err.println("[Info] warming hash cache for revision '$rev' ($sha)")
+        hashProvider.getHashes(sha)
+        System.err.println("[Info] warmed hash cache for revision '$rev' ($sha)")
+      } catch (e: Exception) {
+        System.err.println(
+            "[Warn] warmup for revision '$rev' failed; server will serve it cold on demand: ${e.message}")
+      }
+    }
+  }
+
   /** Reads the optional seed filepaths file into a set of paths. */
   fun loadSeedFilepaths(): Set<Path> =
       seedFilepaths?.readLines()?.filter { it.isNotBlank() }?.map { File(it).toPath() }?.toSet()

cli/src/main/kotlin/com/bazel_diff/server/BazelDiffServer.kt

@@ -8,7 +8,12 @@ import java.io.IOException
 import java.net.InetSocketAddress
 import java.net.URLDecoder
 import java.nio.charset.StandardCharsets
+import java.util.concurrent.Callable
+import java.util.concurrent.ExecutionException
+import java.util.concurrent.ExecutorService
 import java.util.concurrent.Executors
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.TimeoutException
 import org.koin.core.component.KoinComponent
 import org.koin.core.component.inject
 
@@ -32,13 +37,22 @@ import org.koin.core.component.inject
 class BazelDiffServer(
     private val port: Int,
     private val impactedTargetsProvider: ImpactedTargetsProvider,
+    private val requestTimeoutSeconds: Long = 0,
     private val readiness: () -> Boolean,
 ) : KoinComponent {
   private val logger: Logger by inject()
   private val gson: Gson by inject()
 
   @Volatile private var server: HttpServer? = null
 
+  // Runs the per-request query work when a timeout is configured, so the handler thread can abandon
+  // a query that overruns [requestTimeoutSeconds]. Left idle (no threads created) when the timeout
+  // is disabled, since the query then runs inline on the handler thread.
+  private val computeExecutor: ExecutorService =
+      Executors.newCachedThreadPool { runnable ->
+        Thread(runnable, "bazel-diff-compute").apply { isDaemon = true }
+      }
+
   /** Starts the server. Returns immediately; requests are served on background threads. */
   fun start() {
     val httpServer = HttpServer.create(InetSocketAddress(port), 0)
@@ -64,6 +78,7 @@ class BazelDiffServer(
   fun stop(delaySeconds: Int = 0) {
     server?.stop(delaySeconds)
     server = null
+    computeExecutor.shutdownNow()
   }
 
   // HttpExchange exposes close() but does not implement Closeable, so we close it explicitly.
@@ -125,7 +140,11 @@ class BazelDiffServer(
             params["targetType"]?.split(",")?.map { it.trim() }?.filter { it.isNotEmpty() }?.toSet()
 
         try {
-          respondJson(exchange, 200, compute(from, to, targetTypes))
+          respondJson(exchange, 200, computeWithTimeout { compute(from, to, targetTypes) })
+        } catch (e: TimeoutException) {
+          logger.w { "request exceeded ${requestTimeoutSeconds}s timeout, abandoning" }
+          respondJson(
+              exchange, 504, mapOf("error" to "request timed out after ${requestTimeoutSeconds}s"))
         } catch (e: DistancesUnavailableException) {
           respondJson(exchange, 400, mapOf("error" to (e.message ?: "distances unavailable")))
         } catch (e: GitClientException) {
@@ -137,6 +156,29 @@ class BazelDiffServer(
         }
       }
 
+  /**
+   * Runs [compute] bounded by [requestTimeoutSeconds]. When the budget is exceeded the in-flight
+   * task is interrupted and a [TimeoutException] is thrown so the caller can respond `504`. Note the
+   * underlying `bazel query` may not honor the interrupt and can keep running in the background —
+   * abandoning it frees the client and the handler thread, and the query will still populate the
+   * per-SHA cache so a retry is fast. A non-positive timeout (the default) runs [compute] inline on
+   * the handler thread with no bound, preserving the original behavior.
+   */
+  private fun <T> computeWithTimeout(compute: () -> T): T {
+    if (requestTimeoutSeconds <= 0) return compute()
+    val future = computeExecutor.submit(Callable { compute() })
+    try {
+      return future.get(requestTimeoutSeconds, TimeUnit.SECONDS)
+    } catch (e: TimeoutException) {
+      future.cancel(true)
+      throw e
+    } catch (e: ExecutionException) {
+      // Unwrap so the handler's per-type catch blocks (git error, distances unavailable, ...) see
+      // the original exception rather than an ExecutionException wrapper.
+      throw e.cause ?: e
+    }
+  }
+
   private fun respondJson(exchange: HttpExchange, status: Int, body: Any) {
     exchange.responseHeaders.add("Content-Type", "application/json")
     respond(exchange, status, gson.toJson(body))

cli/src/test/kotlin/com/bazel_diff/cli/ServeCommandTest.kt

@@ -66,6 +66,34 @@ class ServeCommandTest : KoinTest {
     }
   }
 
+  /** Records which SHAs warmup asked for and optionally fails for a configured subset. */
+  private class RecordingHashProvider(private val failFor: Set<String> = emptySet()) :
+      com.bazel_diff.server.HashProvider {
+    val warmed = mutableListOf<String>()
+
+    override fun getHashes(sha: String): com.bazel_diff.interactor.HashFileData {
+      warmed += sha
+      if (sha in failFor) throw RuntimeException("boom for $sha")
+      return com.bazel_diff.interactor.HashFileData(emptyMap(), null)
+    }
+
+    override fun <T> withWorkspaceAt(sha: String, block: () -> T): T = block()
+  }
+
+  private object NoopImpactedTargets : com.bazel_diff.server.ImpactedTargetsProvider {
+    override fun getImpactedTargets(
+        fromRev: String,
+        toRev: String,
+        targetTypes: Set<String>?
+    ) = com.bazel_diff.server.ImpactedTargetsResult(fromRev, toRev, emptyList())
+
+    override fun getImpactedTargetsWithDistances(
+        fromRev: String,
+        toRev: String,
+        targetTypes: Set<String>?
+    ) = com.bazel_diff.server.ImpactedTargetsWithDistancesResult(fromRev, toRev, emptyList())
+  }
+
   /**
    * A serve command configured to bind an ephemeral port with a cache dir under the temp folder.
    */
@@ -194,6 +222,35 @@ class ServeCommandTest : KoinTest {
     assertThat(healthCode(server)).isEqualTo(200)
   }
 
+  @Test
+  fun warmUpCacheGeneratesEachRevisionAndIsBestEffort() {
+    val cmd = command().apply { warmupRevisions = linkedSetOf("main", "release") }
+    val git = FakeGitClient() // resolveSha is identity, so warmup SHAs == revision names
+    val provider = RecordingHashProvider(failFor = setOf("release"))
+
+    // A failing revision must not throw or stop the others from warming.
+    cmd.warmUpCache(git, provider)
+
+    assertThat(provider.warmed).isEqualTo(listOf("main", "release"))
+  }
+
+  @Test
+  fun warmupFailureStillBecomesReady() {
+    val cmd = command().apply { warmupRevisions = linkedSetOf("bogus") }
+    val git = FakeGitClient()
+    val provider = RecordingHashProvider(failFor = setOf("bogus"))
+    val ready = java.util.concurrent.atomic.AtomicBoolean(false)
+    val server =
+        com.bazel_diff.server.BazelDiffServer(0, NoopImpactedTargets) { ready.get() }
+            .also { startedServers += it }
+    server.start()
+
+    cmd.performInitialFetch(git, provider, ready, server)
+
+    assertThat(ready.get()).isEqualTo(true)
+    assertThat(healthCode(server)).isEqualTo(200)
+  }
+
   @Test
   fun buildAndStartServerLameDucksOnFetchFailure() {
     val cmd = command()

cli/src/test/kotlin/com/bazel_diff/server/BazelDiffServerTest.kt

@@ -77,7 +77,7 @@ class BazelDiffServerTest : KoinTest {
   @Before
   fun setUp() {
     // Port 0 binds an ephemeral port so parallel test runs never collide.
-    server = BazelDiffServer(0, providerProxy(), { ready.get() })
+    server = BazelDiffServer(0, providerProxy()) { ready.get() }
     server.start()
   }
 
@@ -220,6 +220,46 @@ class BazelDiffServerTest : KoinTest {
     assertThat(get("/impacted_targets_with_distances?from=a&to=b").code).isEqualTo(503)
   }
 
+  @Test
+  fun slowRequestTimesOutWith504() {
+    // A provider that blocks longer than the server's 1s request budget must be abandoned with 504
+    // rather than making the client wait for the full compute.
+    val slowProvider =
+        object : ImpactedTargetsProvider {
+          override fun getImpactedTargets(
+              fromRev: String,
+              toRev: String,
+              targetTypes: Set<String>?
+          ): ImpactedTargetsResult {
+            Thread.sleep(10_000)
+            return ImpactedTargetsResult(fromRev, toRev, emptyList())
+          }
+
+          override fun getImpactedTargetsWithDistances(
+              fromRev: String,
+              toRev: String,
+              targetTypes: Set<String>?
+          ) = throw UnsupportedOperationException()
+        }
+    val slowServer = BazelDiffServer(0, slowProvider, requestTimeoutSeconds = 1) { true }
+    slowServer.start()
+    try {
+      val conn =
+          URL("http://localhost:${slowServer.boundPort()}/impacted_targets?from=a&to=b")
+              .openConnection() as HttpURLConnection
+      val code = conn.responseCode
+      val body =
+          (conn.errorStream ?: conn.inputStream)?.let {
+            BufferedReader(InputStreamReader(it, StandardCharsets.UTF_8)).readText()
+          } ?: ""
+      conn.disconnect()
+      assertThat(code).isEqualTo(504)
+      assertThat(body).contains("timed out")
+    } finally {
+      slowServer.stop()
+    }
+  }
+
   @Test
   fun distancesUnavailableReturns400() {
     provider =