Programming exercises: Stop cancelling build jobs during docker image pulls

src/main/java/de/tum/cit/aet/artemis/buildagent/service/BuildAgentDockerService.java

@@ -2,6 +2,7 @@
 
 import static de.tum.cit.aet.artemis.core.config.Constants.PROFILE_BUILDAGENT;
 
+import java.io.IOException;
 import java.nio.file.Files;
 import java.nio.file.Path;
 import java.time.Duration;
@@ -14,7 +15,9 @@
 import java.util.Map;
 import java.util.Objects;
 import java.util.Set;
+import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 import java.util.concurrent.locks.ReentrantLock;
 import java.util.stream.Collectors;
 
@@ -87,6 +90,9 @@ public class BuildAgentDockerService {
 
     private static final Logger log = LoggerFactory.getLogger(BuildAgentDockerService.class);
 
+    /** How often a running pull is re-examined while waiting, short enough to notice a stall promptly. */
+    private static final int PULL_PROGRESS_POLL_INTERVAL_SECONDS = 5;
+
     private final BuildAgentConfiguration buildAgentConfiguration;
 
     private final DistributedDataAccessService distributedDataAccessService;
@@ -125,6 +131,35 @@ public class BuildAgentDockerService {
     @Value("${artemis.continuous-integration.build-agent.short-name}")
     private String buildAgentShortName;
 
+    /**
+     * Maximum time a single Docker image pull may take before it is aborted.
+     * <p>
+     * This bounds the image pull independently of the per-exercise build timeout: how long a pull takes depends on the image size and on the registry and network,
+     * not on the exercise, so a slow registry must not eat into the time budget a student's build gets. Without this, a pull that never makes progress would block
+     * the build thread indefinitely.
+     */
+    @Value("${artemis.continuous-integration.image-pull-timeout-seconds:300}")
+    private int imagePullTimeoutSeconds;
+
+    /**
+     * Maximum time a Docker image pull may report no progress at all before it is aborted.
+     * <p>
+     * This separates the two ways a pull goes wrong. A large image over a slow link keeps emitting progress and is
+     * allowed to run until {@link #imagePullTimeoutSeconds}. A pull that is not getting through at all, because the
+     * registry is unreachable or a firewall silently drops the packets rather than refusing the connection, emits
+     * nothing, and there is no reason to hold a build thread and an agent slot for the full budget waiting for it.
+     */
+    @Value("${artemis.continuous-integration.image-pull-stall-timeout-seconds:60}")
+    private int imagePullStallTimeoutSeconds;
+
+    /**
+     * IDs of the build jobs that are currently pulling a Docker image, with the time the pull started.
+     * <p>
+     * A job is registered here for the whole time it spends in {@link #pullDockerImage}, which includes waiting for {@link #lock} while another job pulls. During
+     * that window the job legitimately has no Docker container yet, so {@link SharedQueueProcessingService} must not treat it as stale.
+     */
+    private final Map<String, Instant> ongoingImagePulls = new ConcurrentHashMap<>();
+
     private static final String AMD64_ARCHITECTURE = "amd64";
 
     private static final String ARM64_ARCHITECTURE = "arm64";
@@ -137,10 +172,31 @@ public BuildAgentDockerService(BuildAgentConfiguration buildAgentConfiguration,
         this.taskScheduler = taskScheduler;
     }
 
-    // EventListener cannot be used here, as the bean is lazy
-    // https://docs.spring.io/spring-framework/reference/core/beans/context-introduction.html#context-functionality-events-annotation
+    /**
+     * Validates the image pull configuration and schedules the periodic cleanup of dangling build containers.
+     * <p>
+     * EventListener cannot be used here, as the bean is lazy, see the
+     * <a href="https://docs.spring.io/spring-framework/reference/core/beans/context-introduction.html#context-functionality-events-annotation">Spring docs</a>.
+     *
+     * @throws IllegalArgumentException if the configured image pull timeout is not positive
+     */
     @PostConstruct
     public void applicationReady() {
+        // A non-positive timeout would make awaitCompletion return immediately, so every image pull would be reported as timed out and no build could ever run. Fail fast
+        // at startup instead of turning every build into a confusing pull failure.
+        if (imagePullTimeoutSeconds <= 0) {
+            String errorMessage = "The Docker image pull timeout must be a positive number of seconds, but was " + imagePullTimeoutSeconds
+                    + ". It should be changed in the application properties under 'artemis.continuous-integration.image-pull-timeout-seconds'.";
+            log.error(errorMessage);
+            throw new IllegalArgumentException(errorMessage);
+        }
+        if (imagePullStallTimeoutSeconds <= 0) {
+            String errorMessage = "The Docker image pull stall timeout must be a positive number of seconds, but was " + imagePullStallTimeoutSeconds
+                    + ". It should be changed in the application properties under 'artemis.continuous-integration.image-pull-stall-timeout-seconds'.";
+            log.error(errorMessage);
+            throw new IllegalArgumentException(errorMessage);
+        }
+
         // Schedule the cleanup of dangling build containers once 10 seconds after the application has started and then every containerCleanupScheduleMinutes minutes
         taskScheduler.scheduleAtFixedRate(this::cleanUpContainers, Instant.now().plusSeconds(10), Duration.ofMinutes(containerCleanupScheduleMinutes));
     }
@@ -226,8 +282,27 @@ public void cleanUpContainers() {
      */
     public static class MyPullImageResultCallback extends PullImageResultCallback {
 
+        /**
+         * How many updates the daemon has reported for this pull. Written from the docker-java callback thread and
+         * read by the waiting build thread, hence atomic.
+         * <p>
+         * A counter rather than a timestamp: the caller owns the clock, so "no progress yet" is measured from when
+         * the wait started rather than from when this object happened to be constructed.
+         */
+        private final AtomicLong progressCount = new AtomicLong();
+
+        /**
+         * How many updates the daemon has reported so far.
+         *
+         * @return the number of progress updates received for this pull
+         */
+        public long progressCount() {
+            return progressCount.get();
+        }
+
         @Override
         public void onNext(PullResponseItem item) {
+            progressCount.incrementAndGet();
             String msg = "~~~~~~~~~~~~~~~~~~~~ Pull image progress: " + item.getStatus() + " ~~~~~~~~~~~~~~~~~~~~";
             log.debug(msg);
             super.onNext(item);
@@ -262,6 +337,28 @@ public void onComplete() {
      * @throws LocalCIException if the image pull is interrupted or fails due to other exceptions.
      */
     public void pullDockerImage(BuildJobQueueItem buildJob, BuildLogsMap buildLogsMap) {
+        // Register the job for the whole pull phase, including the time spent waiting for the lock, so that the stale build job detection does not cancel a job that is
+        // simply waiting for its image.
+        ongoingImagePulls.put(buildJob.id(), Instant.now());
+        try {
+            doPullDockerImage(buildJob, buildLogsMap);
+        }
+        finally {
+            ongoingImagePulls.remove(buildJob.id());
+        }
+    }
+
+    /**
+     * Returns whether the given build job is currently pulling its Docker image.
+     *
+     * @param buildJobId the ID of the build job
+     * @return true if an image pull is in progress for this build job
+     */
+    public boolean isImagePullInProgress(String buildJobId) {
+        return ongoingImagePulls.containsKey(buildJobId);
+    }
+
+    private void doPullDockerImage(BuildJobQueueItem buildJob, BuildLogsMap buildLogsMap) {
         final String imageName = buildJob.buildConfig().dockerImage();
         if (dockerClientNotAvailable("Cannot pull Docker image.")) {
             throw new LocalCIException("Docker is not available. Cannot pull image " + imageName);
@@ -298,7 +395,7 @@ public void pullDockerImage(BuildJobQueueItem buildJob, BuildLogsMap buildLogsMa
                     // Only pull the image if the inspect command failed
                     var command = dockerClient.pullImageCmd(imageName).withPlatform(imageArchitecture);
                     var exec = command.exec(new MyPullImageResultCallback());
-                    exec.awaitCompletion();
+                    awaitPullCompletion(exec, imageName, buildJob, buildLogsMap);
 
                     // Check if the image is compatible with the current architecture
                     var inspectImageResponse = dockerClient.inspectImageCmd(imageName).exec();
@@ -317,7 +414,7 @@ public void pullDockerImage(BuildJobQueueItem buildJob, BuildLogsMap buildLogsMa
                         try {
                             var fallbackCommand = dockerClient.pullImageCmd(imageName).withPlatform(AMD64_ARCHITECTURE);
                             var fallbackExec = fallbackCommand.exec(new MyPullImageResultCallback());
-                            fallbackExec.awaitCompletion();
+                            awaitPullCompletion(fallbackExec, imageName, buildJob, buildLogsMap);
 
                             // Verify the fallback image was pulled successfully
                             var inspectImageResponse = dockerClient.inspectImageCmd(imageName).exec();
@@ -351,6 +448,73 @@ public void pullDockerImage(BuildJobQueueItem buildJob, BuildLogsMap buildLogsMa
         }
     }
 
+    /**
+     * Waits for a Docker image pull to finish, aborting it once {@code artemis.continuous-integration.image-pull-timeout-seconds} has elapsed.
+     * <p>
+     * Without a timeout a pull that stops making progress, for example because a configured registry mirror silently drops packets, blocks the build thread forever.
+     * The timed {@code awaitCompletion} keeps the error handling of the untimed one: it still calls {@code throwFirstError()}, so a pull that fails rather than stalls
+     * propagates its exception exactly as before, and it closes the callback itself, so the stalled pull is aborted rather than left running in the background.
+     *
+     * @param callback     the callback of the running pull command
+     * @param imageName    the name of the Docker image being pulled
+     * @param buildJob     the build job the pull belongs to
+     * @param buildLogsMap a map for appending log entries related to the build process
+     * @throws InterruptedException if the current thread is interrupted while waiting
+     * @throws LocalCIException     if the pull does not finish within the configured timeout
+     */
+    private void awaitPullCompletion(MyPullImageResultCallback callback, String imageName, BuildJobQueueItem buildJob, BuildLogsMap buildLogsMap) throws InterruptedException {
+        final long pollIntervalNanos = TimeUnit.SECONDS.toNanos(PULL_PROGRESS_POLL_INTERVAL_SECONDS);
+        final long deadlineNanos = System.nanoTime() + TimeUnit.SECONDS.toNanos(imagePullTimeoutSeconds);
+        final long stallNanos = TimeUnit.SECONDS.toNanos(imagePullStallTimeoutSeconds);
+        long lastProgressAtNanos = System.nanoTime();
+        long lastProgressCount = callback.progressCount();
+
+        while (true) {
+            // Wait in slices rather than one long wait, so the pull can also be judged on whether it is still moving. A slice never reaches beyond the next deadline, so a
+            // timeout shorter than the poll interval is honoured just as precisely as a longer one.
+            long nowNanos = System.nanoTime();
+            long sliceNanos = Math.max(0, Math.min(pollIntervalNanos, Math.min(lastProgressAtNanos + stallNanos - nowNanos, deadlineNanos - nowNanos)));
+            if (callback.awaitCompletion(sliceNanos, TimeUnit.NANOSECONDS)) {
+                return;
+            }
+            long progressCount = callback.progressCount();
+            if (progressCount != lastProgressCount) {
+                lastProgressCount = progressCount;
+                lastProgressAtNanos = System.nanoTime();
+            }
+            if (System.nanoTime() - lastProgressAtNanos >= stallNanos) {
+                abortPull(callback, imageName, buildJob, buildLogsMap,
+                        "reported no progress for " + imagePullStallTimeoutSeconds + " seconds. The registry is most likely unreachable from this agent, for example "
+                                + "because a firewall drops the packets instead of refusing the connection");
+            }
+            if (System.nanoTime() - deadlineNanos >= 0) {
+                abortPull(callback, imageName, buildJob, buildLogsMap, "did not finish within " + imagePullTimeoutSeconds + " seconds");
+            }
+        }
+    }
+
+    /**
+     * Closes the callback so the pull is really abandoned rather than left running in the background, then fails the build job.
+     *
+     * @param callback     the callback of the running pull command
+     * @param imageName    the name of the Docker image being pulled
+     * @param buildJob     the build job the pull belongs to
+     * @param buildLogsMap a map for appending log entries related to the build process
+     * @param reason       what went wrong, phrased to continue "Pulling docker image <name> ..."
+     */
+    private static void abortPull(MyPullImageResultCallback callback, String imageName, BuildJobQueueItem buildJob, BuildLogsMap buildLogsMap, String reason) {
+        try {
+            callback.close();
+        }
+        catch (IOException e) {
+            log.warn("Could not close the callback of the aborted pull of docker image {}", imageName, e);
+        }
+        String msg = "~~~~~~~~~~~~~~~~~~~~ Pulling docker image " + imageName + " " + reason + " ~~~~~~~~~~~~~~~~~~~~";
+        log.error(msg);
+        buildLogsMap.appendBuildLogEntry(buildJob.id(), msg);
+        throw new LocalCIException("Pulling docker image " + imageName + " " + reason);
+    }
+
     /**
      * Checks if the architecture of the Docker image is compatible with the current system.
      * On macOS ARM, amd64 images are allowed as they can run via Rosetta 2 emulation.

src/main/java/de/tum/cit/aet/artemis/buildagent/service/BuildJobManagementService.java

@@ -86,6 +86,14 @@ public class BuildJobManagementService {
      */
     private static final java.time.Duration CLUSTER_CONNECTION_RETRY_INTERVAL = java.time.Duration.ofSeconds(5);
 
+    /**
+     * How long a single wait slice for the build result lasts. The build timeout is enforced in slices of this length so
+     * that slices spent pulling the Docker image can be excluded from the build budget, see
+     * {@link #awaitBuildResult(Future, String, int)}. Short enough to report a timeout promptly, long enough to keep the
+     * polling overhead negligible over a multi-minute build.
+     */
+    private static final long BUILD_TIMEOUT_POLL_INTERVAL_MILLIS = 250;
+
     private final BuildJobExecutionService buildJobExecutionService;
 
     private final BuildAgentConfiguration buildAgentConfiguration;
@@ -98,6 +106,8 @@ public class BuildJobManagementService {
 
     private final TaskScheduler taskScheduler;
 
+    private final BuildAgentDockerService buildAgentDockerService;
+
     /**
      * Scheduled future for retrying cluster connection and initialization.
      * This is used when the build agent starts before any core node is available.
@@ -179,13 +189,15 @@ public class BuildJobManagementService {
     private final Set<String> cancelledBuildJobs = new ConcurrentSkipListSet<>();
 
     public BuildJobManagementService(DistributedDataAccessService distributedDataAccessService, BuildJobExecutionService buildJobExecutionService,
-            BuildAgentConfiguration buildAgentConfiguration, BuildJobContainerService buildJobContainerService, BuildLogsMap buildLogsMap, TaskScheduler taskScheduler) {
+            BuildAgentConfiguration buildAgentConfiguration, BuildJobContainerService buildJobContainerService, BuildLogsMap buildLogsMap, TaskScheduler taskScheduler,
+            BuildAgentDockerService buildAgentDockerService) {
         this.buildJobExecutionService = buildJobExecutionService;
         this.buildAgentConfiguration = buildAgentConfiguration;
         this.buildJobContainerService = buildJobContainerService;
         this.distributedDataAccessService = distributedDataAccessService;
         this.buildLogsMap = buildLogsMap;
         this.taskScheduler = taskScheduler;
+        this.buildAgentDockerService = buildAgentDockerService;
     }
 
     /**
@@ -357,7 +369,7 @@ public CompletableFuture<BuildResult> executeBuildJob(BuildJobQueueItem buildJob
 
         CompletableFuture<BuildResult> futureResult = createCompletableFuture(() -> {
             try {
-                return future.get(buildJobTimeoutSeconds, TimeUnit.SECONDS);
+                return awaitBuildResult(future, buildJobItem.id(), buildJobTimeoutSeconds);
             }
             catch (Exception ex) {
                 if (DockerUtil.isDockerNotAvailable(ex)) {
@@ -395,6 +407,49 @@ public CompletableFuture<BuildResult> executeBuildJob(BuildJobQueueItem buildJob
         }));
     }
 
+    /**
+     * Waits for the build result, without letting the time spent pulling the Docker image consume the build timeout.
+     * <p>
+     * Pulling the image is the first step of {@link BuildJobExecutionService#runBuildJob}, so it runs inside the window
+     * guarded by the build timeout. A cold pull of a large image can easily take longer than
+     * {@code artemis.continuous-integration.build-timeout-seconds.max} (240 seconds by default), which would cancel the
+     * job and report it as a build timeout even though the build itself never started. The pull has its own, much longer
+     * budget ({@code artemis.continuous-integration.image-pull-timeout-seconds}) and must therefore be excluded here.
+     * <p>
+     * Rather than starting the timer after image preparation, the wait is sliced: only slices during which no pull is in
+     * progress count against the build budget. That keeps the accounting correct for images that are already present
+     * locally (no pull, so the full budget applies from the start) as well as for pulls that finish mid-build.
+     *
+     * Package-private for testing.
+     *
+     * @param future                 the future of the running build job
+     * @param buildJobId             the ID of the build job, used to check whether its image pull is still running
+     * @param buildJobTimeoutSeconds the build budget in seconds, excluding any time spent pulling the image
+     * @return the build result
+     * @throws TimeoutException if the build itself, not counting image pulls, exceeded the build budget
+     */
+    BuildResult awaitBuildResult(Future<BuildResult> future, String buildJobId, int buildJobTimeoutSeconds) throws Exception {
+        final long budgetNanos = TimeUnit.SECONDS.toNanos(buildJobTimeoutSeconds);
+        long consumedNanos = 0;
+
+        while (true) {
+            final long sliceStartNanos = System.nanoTime();
+            try {
+                return future.get(BUILD_TIMEOUT_POLL_INTERVAL_MILLIS, TimeUnit.MILLISECONDS);
+            }
+            catch (TimeoutException timeout) {
+                if (buildAgentDockerService.isImagePullInProgress(buildJobId)) {
+                    // The pull is bounded by the image pull timeout, so this slice does not count against the build budget.
+                    continue;
+                }
+                consumedNanos += System.nanoTime() - sliceStartNanos;
+                if (consumedNanos >= budgetNanos) {
+                    throw timeout;
+                }
+            }
+        }
+    }
+
     private void logTimedOutBuildJob(BuildJobQueueItem buildJobItem, int buildJobTimeoutSeconds) {
         String msg = "Timed out after " + buildJobTimeoutSeconds + " seconds. "
                 + "This may be due to an infinite loop or inefficient code. Please review your code for potential issues. "