From 0c54be6ddf602ed568197015507849102e4e481b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Germ=C3=A1n=20Rodr=C3=ADguez?= <36035218+grojeda@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:22:22 +1000 Subject: [PATCH] limit concurrent downloads per user and queue the rest --- .../server/core/configuration/Config.java | 2 + .../core/model/DownloadExecutionResponse.java | 12 +- .../core/model/DownloadJobStatusInfo.java | 29 ++ .../server/processes/DownloadAdmission.java | 21 + .../processes/DownloadAdmissionService.java | 277 +++++++++++++ .../processes/DownloadJobStatusService.java | 63 ++- .../processes/DownloadLimitProperties.java | 24 ++ .../server/processes/DownloadRequest.java | 20 + .../ogcapi/server/processes/HeldDownload.java | 25 ++ .../processes/InFlightDownloadCounter.java | 268 ++++++++++++ .../aodn/ogcapi/server/processes/RestApi.java | 33 +- .../ogcapi/server/processes/RestServices.java | 53 +-- server/src/main/resources/application.yaml | 13 + .../DownloadAdmissionServiceTest.java | 385 ++++++++++++++++++ .../DownloadJobStatusServiceTest.java | 61 +++ .../InFlightDownloadCounterTest.java | 298 ++++++++++++++ .../server/processes/MutableTestClock.java | 40 ++ .../server/processes/RestApiJobsTest.java | 43 +- .../ogcapi/server/processes/RestApiTest.java | 35 +- .../server/processes/RestServicesTest.java | 25 +- .../src/test/resources/application-test.yaml | 8 + 21 files changed, 1670 insertions(+), 65 deletions(-) create mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmission.java create mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionService.java create mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadLimitProperties.java create mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadRequest.java create mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/HeldDownload.java create mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounter.java create mode 100644 server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionServiceTest.java create mode 100644 server/src/test/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounterTest.java create mode 100644 server/src/test/java/au/org/aodn/ogcapi/server/processes/MutableTestClock.java diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/core/configuration/Config.java b/server/src/main/java/au/org/aodn/ogcapi/server/core/configuration/Config.java index f0c438f7..ba484234 100644 --- a/server/src/main/java/au/org/aodn/ogcapi/server/core/configuration/Config.java +++ b/server/src/main/java/au/org/aodn/ogcapi/server/core/configuration/Config.java @@ -9,6 +9,7 @@ import au.org.aodn.ogcapi.server.core.util.GeometryUtils; import au.org.aodn.ogcapi.server.core.util.RestTemplateUtils; import au.org.aodn.ogcapi.server.processes.BatchJobProperties; +import au.org.aodn.ogcapi.server.processes.DownloadLimitProperties; import com.fasterxml.jackson.annotation.JsonInclude; import com.fasterxml.jackson.databind.ObjectMapper; import jakarta.annotation.PostConstruct; @@ -35,6 +36,7 @@ GNProperties.class, DasProperties.class, BatchJobProperties.class, + DownloadLimitProperties.class, OgcApiProperties.class }) public class Config { diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadExecutionResponse.java b/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadExecutionResponse.java index d2c65996..8abf679c 100644 --- a/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadExecutionResponse.java +++ b/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadExecutionResponse.java @@ -4,10 +4,18 @@ import com.fasterxml.jackson.annotation.JsonProperty; import io.swagger.v3.oas.annotations.media.Schema; -@Schema(description = "Compatible download execution response with the submitted AWS Batch job ID.") +@Schema(description = "Compatible download execution response with the submitted job ID, " + + "and whether the download is waiting for a free per-user slot.") public record DownloadExecutionResponse( @JsonProperty("message") InlineValue message, @JsonProperty("status") InlineValue status, - @JsonProperty("jobID") String jobId + @JsonProperty("jobID") String jobId, + @Schema(description = "True when the download was accepted but is waiting for one of " + + "this user's concurrent download slots to free. It still has a job ID and " + + "will start on its own; nothing further is required from the caller.") + @JsonProperty("queued") boolean queued, + @Schema(description = "How many of this user's downloads, including this one, are " + + "waiting ahead of it; 1 means it starts next. Omitted unless queued.") + @JsonProperty("queuePosition") Integer queuePosition ) implements InlineResponse200 { } diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadJobStatusInfo.java b/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadJobStatusInfo.java index e05ea248..dedb970b 100644 --- a/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadJobStatusInfo.java +++ b/server/src/main/java/au/org/aodn/ogcapi/server/core/model/DownloadJobStatusInfo.java @@ -31,6 +31,19 @@ public class DownloadJobStatusInfo extends StatusInfo { @Schema(description = "Link to the metadata page supplied when the download was submitted.") private String metadataUrl; + @JsonProperty("queued") + @Schema(description = "True while the download is waiting for one of this user's " + + "concurrent download slots to free. Always present, so clients never have to " + + "infer the queued state from the message text.") + private boolean queued; + + @JsonProperty("queuePosition") + @JsonInclude(JsonInclude.Include.NON_ABSENT) + @Schema(description = "How many of this user's downloads, including this one, are waiting " + + "ahead of it; 1 means it starts next. Omitted unless queued. This is a position, " + + "not an estimated time: it depends on when the running downloads finish.") + private Integer queuePosition; + public String getCollection() { return collection; } @@ -62,4 +75,20 @@ public String getMetadataUrl() { public void setMetadataUrl(String metadataUrl) { this.metadataUrl = metadataUrl; } + + public boolean isQueued() { + return queued; + } + + public void setQueued(boolean queued) { + this.queued = queued; + } + + public Integer getQueuePosition() { + return queuePosition; + } + + public void setQueuePosition(Integer queuePosition) { + this.queuePosition = queuePosition; + } } diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmission.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmission.java new file mode 100644 index 00000000..8cf6ea83 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmission.java @@ -0,0 +1,21 @@ +package au.org.aodn.ogcapi.server.processes; + +/** + * The outcome of accepting a download: the job id the caller polls, and whether it went + * straight to AWS Batch or is waiting for one of that user's slots to free. + * + * @param jobId the AWS Batch job id when submitted, or a locally minted id when held + * @param queued true while the download is waiting rather than running + * @param queuePosition how many of this user's downloads, including this one, are waiting + * ahead of it. 1 means it is next. Null when not queued. + */ +public record DownloadAdmission(String jobId, boolean queued, Integer queuePosition) { + + static DownloadAdmission submitted(String awsJobId) { + return new DownloadAdmission(awsJobId, false, null); + } + + static DownloadAdmission queued(String jobId, int queuePosition) { + return new DownloadAdmission(jobId, true, queuePosition); + } +} diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionService.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionService.java new file mode 100644 index 00000000..c313bc91 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionService.java @@ -0,0 +1,277 @@ +package au.org.aodn.ogcapi.server.processes; + +import com.fasterxml.jackson.core.JsonProcessingException; +import lombok.extern.slf4j.Slf4j; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.scheduling.annotation.Scheduled; +import org.springframework.stereotype.Service; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayDeque; +import java.util.ArrayList; +import java.util.Deque; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.locks.ReentrantLock; + +/** + * Caps how many downloads one user can have running at once. A user at the limit is not + * rejected: the request is held here and submitted to AWS Batch as soon as one of that + * user's slots frees, so the caller still gets a job id it can poll. + * + *

The hold queue is in memory only. A restart therefore loses whatever was waiting, and a + * second replica would enforce its own limit rather than a shared one. + */ +@Slf4j +@Service +public class DownloadAdmissionService { + + static final String QUEUED_MESSAGE = "Download job queued, waiting for a free slot"; + + private static final int MAX_SUBMIT_ATTEMPTS = 3; + + /** + * How long the id of a released download keeps resolving to its AWS job. AWS Batch drops + * the job record itself after about a day, so a longer retention would only translate an + * id into a job that no longer exists. + */ + static final Duration RELEASED_RETENTION = Duration.ofHours(24); + + private final RestServices restServices; + private final InFlightDownloadCounter counter; + private final DownloadLimitProperties limits; + private final Clock clock; + + /** + * Guards the admit-or-hold decision and the release loop against each other. The AWS + * sweep is deliberately done before this is taken; the submit itself is inside it, so + * that reading the count and acting on it cannot interleave and over-admit. + */ + private final ReentrantLock lock = new ReentrantLock(); + + /** FIFO across all users. Guarded by {@link #lock}. */ + private final Deque held = new ArrayDeque<>(); + + private final Map heldById = new ConcurrentHashMap<>(); + private final Map released = new ConcurrentHashMap<>(); + + record Released(String awsJobId, Instant releasedAt) { + } + + @Autowired + public DownloadAdmissionService( + RestServices restServices, + InFlightDownloadCounter counter, + DownloadLimitProperties limits) { + this(restServices, counter, limits, Clock.systemUTC()); + } + + DownloadAdmissionService( + RestServices restServices, + InFlightDownloadCounter counter, + DownloadLimitProperties limits, + Clock clock) { + this.restServices = restServices; + this.counter = counter; + this.limits = limits; + this.clock = clock; + } + + /** + * Accept a download. Returns the AWS Batch job id when it was submitted straight away, or + * a locally minted id when the user was at their limit and the request is now waiting. + */ + public DownloadAdmission submitOrHold(DownloadRequest request) throws JsonProcessingException { + Map parameters = restServices.buildDownloadParameters(request); + String jobName = RestServices.downloadJobName(request.recipient()); + + if (!limits.enabled()) { + String awsJobId = restServices.submitDownloadJob(jobName, parameters); + notifyStarted(request); + return DownloadAdmission.submitted(awsJobId); + } + + // Outside the lock: the sweep is the only expensive step and every user shares it. + counter.refreshIfStale(); + + DownloadAdmission admission; + lock.lock(); + try { + if (hasHeldFor(request.recipient())) { + // Never jump ahead of this user's own waiting requests. + admission = hold(request, jobName, parameters); + } else if (counter.countInFlight(request.recipient()) < limits.maxConcurrent()) { + admission = DownloadAdmission.submitted(submit(jobName, parameters, request.recipient())); + } else { + admission = hold(request, jobName, parameters); + } + } finally { + lock.unlock(); + } + + if (!admission.queued()) { + // Outside the lock: this is a synchronous SES call, and it is best effort anyway. + notifyStarted(request); + } + return admission; + } + + /** + * The still-waiting download with this id and where it sits in its owner's queue, or null + * if this id is not one of ours. Both come from one look under the lock so the position + * cannot be taken from a queue that has already moved on. + */ + public HeldView findHeld(String jobId) { + lock.lock(); + try { + HeldDownload download = heldById.get(jobId); + return download == null ? null : new HeldView(download, positionOf(download)); + } finally { + lock.unlock(); + } + } + + /** A waiting download together with its place in its owner's queue. */ + public record HeldView(HeldDownload download, int position) { + } + + /** + * How many of this user's waiting downloads, counting this one, sit at or ahead of it. + * Only that user's own queue matters: their downloads are released as their own slots + * free, so what is waiting for other people says nothing about this one. + */ + private int positionOf(HeldDownload target) { + String key = InFlightDownloadCounter.recipientKey(target.request().recipient()); + int position = 0; + for (HeldDownload job : held) { + if (InFlightDownloadCounter.recipientKey(job.request().recipient()).equals(key)) { + position++; + if (job.jobId().equals(target.jobId())) { + return position; + } + } + } + return position; + } + + /** The AWS Batch job a released download became, or null if this id was never held. */ + public String awsJobIdOf(String jobId) { + Released record = released.get(jobId); + return record == null ? null : record.awsJobId(); + } + + /** + * Submit whatever is waiting and now fits. Nothing held means no AWS call at all, which + * is both the steady state and what keeps the scheduler away from live AWS Batch in tests. + */ + @Scheduled(fixedDelayString = "${aws.batch.job.user-limit.release-interval:15s}") + public void releaseHeldDownloads() { + pruneReleased(); + if (!limits.enabled() || heldById.isEmpty()) { + return; + } + + counter.refresh(); + + List toNotify = new ArrayList<>(); + lock.lock(); + try { + expireStale(); + + Deque blocked = new ArrayDeque<>(); + HeldDownload job; + while ((job = held.pollFirst()) != null) { + String recipient = job.request().recipient(); + if (counter.countInFlight(recipient) >= limits.maxConcurrent()) { + blocked.addLast(job); + continue; + } + try { + String awsJobId = submit(job.jobName(), job.parameters(), recipient); + heldById.remove(job.jobId()); + released.put(job.jobId(), new Released(awsJobId, clock.instant())); + toNotify.add(job.request()); + log.info("Released held download {} as AWS Batch job {}", job.jobId(), awsJobId); + } catch (Exception e) { + HeldDownload retried = job.withAttempt(); + if (retried.attempts() >= MAX_SUBMIT_ATTEMPTS) { + heldById.remove(job.jobId()); + log.error("Abandoning held download {} after {} failed submissions", + job.jobId(), retried.attempts(), e); + } else { + log.warn("Could not release held download {}, will retry", job.jobId(), e); + heldById.put(retried.jobId(), retried); + blocked.addLast(retried); + } + } + } + held.addAll(blocked); + } finally { + lock.unlock(); + } + + toNotify.forEach(this::notifyStarted); + } + + private String submit(String jobName, Map parameters, String recipient) { + String awsJobId = restServices.submitDownloadJob(jobName, parameters); + counter.recordSubmitted(awsJobId, recipient); + return awsJobId; + } + + private boolean hasHeldFor(String recipient) { + String key = InFlightDownloadCounter.recipientKey(recipient); + return held.stream() + .anyMatch(job -> InFlightDownloadCounter.recipientKey(job.request().recipient()).equals(key)); + } + + private DownloadAdmission hold(DownloadRequest request, String jobName, Map parameters) { + if (held.size() >= limits.maxHeldTotal()) { + // A safety valve, not an expected outcome: the alternative is growing the queue + // until the process runs out of memory. + throw new IllegalStateException( + "Download hold queue is full (" + limits.maxHeldTotal() + " waiting)"); + } + String jobId = UUID.randomUUID().toString(); + HeldDownload job = new HeldDownload(jobId, request, jobName, parameters, clock.instant(), 0); + held.addLast(job); + heldById.put(jobId, job); + log.info("Holding download {} for a free slot, {} now waiting", jobId, held.size()); + return DownloadAdmission.queued(jobId, positionOf(job)); + } + + private void expireStale() { + Instant cutoff = clock.instant().minus(limits.maxHoldAge()); + held.removeIf(job -> { + if (job.acceptedAt().isBefore(cutoff)) { + heldById.remove(job.jobId()); + log.warn("Abandoning download {} held since {}", job.jobId(), job.acceptedAt()); + return true; + } + return false; + }); + } + + private void pruneReleased() { + Instant cutoff = clock.instant().minus(RELEASED_RETENTION); + released.entrySet().removeIf(entry -> entry.getValue().releasedAt().isBefore(cutoff)); + } + + private void notifyStarted(DownloadRequest request) { + restServices.notifyUser( + request.recipient(), + request.uuid(), + request.key(), + request.startDate(), + request.endDate(), + request.multiPolygon(), + request.collectionTitle(), + request.fullMetadataLink(), + request.suggestedCitation(), + request.outputFormat()); + } +} diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusService.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusService.java index 69123f4c..0b94f9d8 100644 --- a/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusService.java +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusService.java @@ -37,6 +37,10 @@ public class DownloadJobStatusService { static final String PROCESS_ID = "download-dataset"; static final Duration CHILD_DISCOVERY_WINDOW = Duration.ofSeconds(30); + // These exact names are an internal contract with data-access-service. Any DAS + // naming change must be applied here at the same time. + static final String PREPARE_NAME_PREFIX = "prepare-data-for-job-"; + static final String COLLECT_NAME_PREFIX = "collect-data-for-job-"; private static final String INITIAL_TYPE = "sub-setting"; private static final String PREPARE_TYPE = "sub-setting-data-preparation"; @@ -48,47 +52,61 @@ public class DownloadJobStatusService { private final BatchClient batchClient; private final BatchJobProperties properties; private final DownloadJobStatusAggregator aggregator; + private final DownloadAdmissionService admissionService; private final Clock clock; @Autowired public DownloadJobStatusService( BatchClient batchClient, BatchJobProperties properties, - DownloadJobStatusAggregator aggregator) { - this(batchClient, properties, aggregator, Clock.systemUTC()); + DownloadJobStatusAggregator aggregator, + DownloadAdmissionService admissionService) { + this(batchClient, properties, aggregator, admissionService, Clock.systemUTC()); } DownloadJobStatusService( BatchClient batchClient, BatchJobProperties properties, DownloadJobStatusAggregator aggregator, + DownloadAdmissionService admissionService, Clock clock) { this.batchClient = batchClient; this.properties = properties; this.aggregator = aggregator; + this.admissionService = admissionService; this.clock = clock; } public DownloadJobStatusInfo getStatus(String jobId) { validateJobId(jobId); + DownloadAdmissionService.HeldView heldDownload = admissionService.findHeld(jobId); + if (heldDownload != null) { + return heldStatus(jobId, heldDownload); + } + + // A download that waited for a slot answers on the id its caller was given, not on + // the AWS job id it eventually turned into. + String awsJobId = admissionService.awsJobIdOf(jobId); + return awsStatus(awsJobId == null ? jobId : awsJobId, jobId); + } + + private DownloadJobStatusInfo awsStatus(String awsJobId, String publicJobId) { try { - JobDetail initial = describeInitialJob(jobId); + JobDetail initial = describeInitialJob(awsJobId); if (initial.status() == JobStatus.FAILED) { StatusCode status = aggregator.aggregate(new DownloadJobStatusAggregator.Snapshot( initial.status(), null, null, DownloadJobStatusAggregator.WorkflowMode.CHILD_DISCOVERY_REQUIRED, false)); - return toStatusInfo(jobId, status, initial, null, null); + return toStatusInfo(publicJobId, status, initial, null, null); } - // These exact names are an internal contract with data-access-service. Any DAS - // naming change must be applied here at the same time. JobDetail prepare = findChildJob( - "prepare-data-for-job-" + jobId, jobId, PREPARE_TYPE); + PREPARE_NAME_PREFIX + awsJobId, awsJobId, PREPARE_TYPE); JobDetail collect = findChildJob( - "collect-data-for-job-" + jobId, jobId, COLLECT_TYPE); + COLLECT_NAME_PREFIX + awsJobId, awsJobId, COLLECT_TYPE); boolean discoveryWindowExpired = discoveryWindowExpired(initial); DownloadJobStatusAggregator.WorkflowMode workflowMode = isExplicitZarr(initial.parameters()) @@ -102,15 +120,40 @@ public DownloadJobStatusInfo getStatus(String jobId) { workflowMode, discoveryWindowExpired)); - return toStatusInfo(jobId, status, initial, prepare, collect); + return toStatusInfo(publicJobId, status, initial, prepare, collect); } catch (DownloadJobNotFoundException e) { throw e; } catch (Exception e) { - log.error("Failed to reconstruct AWS Batch workflow for download job {}", jobId, e); + log.error("Failed to reconstruct AWS Batch workflow for download job {}", publicJobId, e); throw new DownloadJobStatusException(e); } } + /** + * A download still waiting for a free slot. It has no AWS job yet, so the display fields + * come from the parameters that were built when it was accepted. + */ + private DownloadJobStatusInfo heldStatus(String jobId, DownloadAdmissionService.HeldView heldView) { + HeldDownload heldDownload = heldView.download(); + + DownloadJobStatusInfo result = new DownloadJobStatusInfo(); + result.setProcessID(PROCESS_ID); + result.setType(StatusInfo.TypeEnum.PROCESS); + result.setJobID(jobId); + result.setStatus(StatusCode.ACCEPTED); + result.setMessage(DownloadAdmissionService.QUEUED_MESSAGE); + result.setQueued(true); + result.setQueuePosition(heldView.position()); + + Map parameters = heldDownload.parameters(); + result.setCollection(nonBlank(parameters.get(DatasetDownloadEnums.Parameter.COLLECTION_TITLE.getValue()))); + result.setDataSelection(nonBlank(parameters.get(DatasetDownloadEnums.Parameter.KEY.getValue()))); + result.setFormat(nonBlank(parameters.get(DatasetDownloadEnums.Parameter.OUTPUT_FORMAT.getValue()))); + result.setMetadataUrl(nonBlank(parameters.get(DatasetDownloadEnums.Parameter.FULL_METADATA_LINK.getValue()))); + result.setCreated(Date.from(heldDownload.acceptedAt())); + return result; + } + private void validateJobId(String jobId) { try { UUID parsed = UUID.fromString(jobId); diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadLimitProperties.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadLimitProperties.java new file mode 100644 index 00000000..4695c781 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadLimitProperties.java @@ -0,0 +1,24 @@ +package au.org.aodn.ogcapi.server.processes; + +import org.springframework.boot.context.properties.ConfigurationProperties; +import org.springframework.boot.context.properties.bind.DefaultValue; + +import java.time.Duration; + +/** + * Per-user admission limits for dataset downloads. A user, identified by the recipient + * email, may have at most {@code maxConcurrent} downloads in flight at once; anything + * beyond that is held in memory and released as slots free rather than rejected. + */ +@ConfigurationProperties(prefix = "aws.batch.job.user-limit") +public record DownloadLimitProperties( + @DefaultValue("true") boolean enabled, + @DefaultValue("10") int maxConcurrent, + /** Release loop period, and the time-to-live of the cached in-flight snapshot. */ + @DefaultValue("15s") Duration releaseInterval, + /** A download held longer than this is abandoned; its job id then reports as not found. */ + @DefaultValue("24h") Duration maxHoldAge, + /** Safety valve so a runaway client cannot grow the hold queue without bound. */ + @DefaultValue("1000") int maxHeldTotal +) { +} diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadRequest.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadRequest.java new file mode 100644 index 00000000..6aff7244 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadRequest.java @@ -0,0 +1,20 @@ +package au.org.aodn.ogcapi.server.processes; + +/** + * The inputs of one {@code download} execute request, as extracted from the OGC Execute + * body. Carried as a unit so a request that has to wait for a free slot can be submitted + * later exactly as it arrived. + */ +public record DownloadRequest( + String uuid, + String key, + String startDate, + String endDate, + Object multiPolygon, + String recipient, + String collectionTitle, + String fullMetadataLink, + String suggestedCitation, + String outputFormat +) { +} diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/HeldDownload.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/HeldDownload.java new file mode 100644 index 00000000..0fbcced3 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/HeldDownload.java @@ -0,0 +1,25 @@ +package au.org.aodn.ogcapi.server.processes; + +import java.time.Instant; +import java.util.Map; + +/** + * A download that was accepted but not yet submitted to AWS Batch, because its owner was at + * the per-user concurrency limit. + * + * The Batch job name and parameters are built at accept time rather than at release time so + * that releasing is a plain submit, and so the status endpoint can describe a held job - + * collection, format, metadata link - without any AWS call. + */ +record HeldDownload( + String jobId, + DownloadRequest request, + String jobName, + Map parameters, + Instant acceptedAt, + int attempts +) { + HeldDownload withAttempt() { + return new HeldDownload(jobId, request, jobName, parameters, acceptedAt, attempts + 1); + } +} diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounter.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounter.java new file mode 100644 index 00000000..f1eedde7 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounter.java @@ -0,0 +1,268 @@ +package au.org.aodn.ogcapi.server.processes; + +import au.org.aodn.ogcapi.server.core.model.enumeration.DatasetDownloadEnums; +import lombok.extern.slf4j.Slf4j; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.stereotype.Service; +import software.amazon.awssdk.services.batch.BatchClient; +import software.amazon.awssdk.services.batch.model.DescribeJobsRequest; +import software.amazon.awssdk.services.batch.model.JobDetail; +import software.amazon.awssdk.services.batch.model.JobStatus; +import software.amazon.awssdk.services.batch.model.JobSummary; +import software.amazon.awssdk.services.batch.model.ListJobsRequest; +import software.amazon.awssdk.services.batch.model.ListJobsResponse; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashSet; +import java.util.List; +import java.util.Locale; +import java.util.Map; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Counts how many downloads each user has in flight, from one sweep of the AWS Batch queues + * that every user shares rather than a query per request. + * + *

A download occupies a slot while its aggregated status is {@code accepted} or + * {@code running}. In every case {@link DownloadJobStatusAggregator} produces, that is + * equivalent to "the master job, or one of its prepare/collect children, is in a non-terminal + * Batch status", with one gap: a master that has succeeded before its children appear is in + * neither sweep yet still aggregates to running. Jobs submitted within {@link #SUBMIT_GRACE} + * are therefore counted from memory whatever the sweep saw, which also covers everything + * submitted since the last sweep. + * + *

The owning email is read back from the {@code recipient} job parameter, never from the + * job name: {@link RestServices#downloadJobName(String)} sanitises the address, so two + * different addresses can produce the same name. + */ +@Slf4j +@Service +public class InFlightDownloadCounter { + + private static final List NON_TERMINAL = List.of( + JobStatus.SUBMITTED, + JobStatus.PENDING, + JobStatus.RUNNABLE, + JobStatus.STARTING, + JobStatus.RUNNING); + + /** + * How long a freshly submitted job keeps counting from memory. It has to outlast the + * child discovery window of the status service, which is exactly the period in which a + * succeeded master with no children yet still aggregates to running. + */ + static final Duration SUBMIT_GRACE = DownloadJobStatusService.CHILD_DISCOVERY_WINDOW.plusSeconds(90); + + private static final int PAGE_SIZE = 100; + + private final BatchClient batchClient; + private final BatchJobProperties properties; + private final DownloadLimitProperties limits; + private final Clock clock; + + /** AWS job id to the submission that produced it, retained for {@link #SUBMIT_GRACE}. */ + private final Map recentSubmissions = new ConcurrentHashMap<>(); + + private volatile Snapshot snapshot = Snapshot.empty(); + + @Autowired + public InFlightDownloadCounter( + BatchClient batchClient, + BatchJobProperties properties, + DownloadLimitProperties limits) { + this(batchClient, properties, limits, Clock.systemUTC()); + } + + InFlightDownloadCounter( + BatchClient batchClient, + BatchJobProperties properties, + DownloadLimitProperties limits, + Clock clock) { + this.batchClient = batchClient; + this.properties = properties; + this.limits = limits; + this.clock = clock; + } + + private record Submission(String recipient, Instant submittedAt) { + } + + /** + * @param countsByRecipient in-flight downloads per recipient email + * @param countedMasterIds the master job ids behind those counts, so a job the sweep + * already counted is not counted again from recent submissions + */ + private record Snapshot(Map countsByRecipient, Set countedMasterIds, Instant takenAt) { + static Snapshot empty() { + return new Snapshot(Map.of(), Set.of(), Instant.EPOCH); + } + } + + /** + * Refresh the shared snapshot if it has aged past the release interval. Call this before + * taking any admission lock: it is the only part of counting that talks to AWS. + */ + public void refreshIfStale() { + if (Duration.between(snapshot.takenAt(), clock.instant()).compareTo(limits.releaseInterval()) >= 0) { + refresh(); + } + } + + /** Sweep the queues now, whatever the age of the current snapshot. */ + public synchronized void refresh() { + try { + snapshot = sweep(); + } catch (Exception e) { + // Keep serving the previous snapshot. A failed sweep must not fail the download + // request that triggered it; a stale count at worst admits a job that should have + // been held, and the next successful sweep corrects it. + log.error("Failed to sweep AWS Batch for in-flight downloads, reusing the previous snapshot", e); + } + } + + /** + * In-flight downloads for one recipient: what the last sweep saw, plus anything submitted + * too recently for that sweep to have picked it up. + */ + public int countInFlight(String recipient) { + pruneRecentSubmissions(); + String key = recipientKey(recipient); + Snapshot current = snapshot; + int count = current.countsByRecipient().getOrDefault(key, 0); + for (Map.Entry entry : recentSubmissions.entrySet()) { + if (entry.getValue().recipient().equals(key) + && !current.countedMasterIds().contains(entry.getKey())) { + count++; + } + } + return count; + } + + /** Record a job we just submitted so it counts immediately, before any sweep can see it. */ + public void recordSubmitted(String awsJobId, String recipient) { + recentSubmissions.put(awsJobId, new Submission(recipientKey(recipient), clock.instant())); + } + + /** + * The key one user is counted under. Email addresses are case-insensitive in the part + * that matters here and arrive however the user typed them, so without this a capital + * letter would silently buy a second allowance of slots. + * + *

Only ever a counting key. The address SES writes to, and the {@code recipient} job + * parameter data-access-service reads, stay exactly as the user supplied them. + */ + static String recipientKey(String recipient) { + return recipient == null ? null : recipient.trim().toLowerCase(Locale.ROOT); + } + + private void pruneRecentSubmissions() { + Instant cutoff = clock.instant().minus(SUBMIT_GRACE); + recentSubmissions.entrySet().removeIf(entry -> entry.getValue().submittedAt().isBefore(cutoff)); + } + + private Snapshot sweep() { + // The download queue and the child queue are the same by default, so sweep each + // distinct queue once rather than once per role. + Set queues = new LinkedHashSet<>(); + queues.add(properties.queue()); + queues.add(properties.childQueue()); + + Set candidateMasterIds = new LinkedHashSet<>(); + for (String queue : queues) { + for (JobSummary summary : listNonTerminal(queue)) { + if (queue.equals(properties.queue()) && summary.jobId() != null && !summary.jobId().isBlank()) { + // Anything non-terminal on the download queue is a candidate master. The + // describe below discards whatever turns out not to be one of ours. + candidateMasterIds.add(summary.jobId()); + } + String masterId = masterIdOf(summary.jobName()); + if (masterId != null) { + candidateMasterIds.add(masterId); + } + } + } + + Map counts = new HashMap<>(); + Set counted = new LinkedHashSet<>(); + for (JobDetail job : describeJobs(candidateMasterIds)) { + if (!isDownloadMaster(job)) { + continue; + } + String recipient = job.parameters().get(DatasetDownloadEnums.Parameter.RECIPIENT.getValue()); + if (recipient == null || recipient.isBlank()) { + continue; + } + counts.merge(recipientKey(recipient), 1, Integer::sum); + counted.add(job.jobId()); + } + return new Snapshot(counts, counted, clock.instant()); + } + + /** + * The master job id a prepare/collect child belongs to, or null when this is not one of + * the data-access-service child jobs. The names are the same contract + * {@link DownloadJobStatusService} relies on. + */ + static String masterIdOf(String jobName) { + if (jobName == null) { + return null; + } + if (jobName.startsWith(DownloadJobStatusService.PREPARE_NAME_PREFIX)) { + return blankToNull(jobName.substring(DownloadJobStatusService.PREPARE_NAME_PREFIX.length())); + } + if (jobName.startsWith(DownloadJobStatusService.COLLECT_NAME_PREFIX)) { + return blankToNull(jobName.substring(DownloadJobStatusService.COLLECT_NAME_PREFIX.length())); + } + return null; + } + + private static String blankToNull(String value) { + return value.isBlank() ? null : value; + } + + private boolean isDownloadMaster(JobDetail job) { + return DownloadJobStatusService.matchesQueue(properties.queue(), job.jobQueue()) + && DownloadJobStatusService.matchesJobDefinition(properties.definition(), job.jobDefinition()) + && DatasetDownloadEnums.Type.SUB_SETTING.getValue() + .equals(job.parameters().get(DatasetDownloadEnums.Parameter.TYPE.getValue())); + } + + /** + * Every non-terminal job on a queue. ListJobs returns only RUNNABLE jobs when given + * neither a filter nor a status, so the statuses are enumerated explicitly. + */ + private List listNonTerminal(String queue) { + List result = new ArrayList<>(); + for (JobStatus status : NON_TERMINAL) { + String nextToken = null; + do { + ListJobsResponse response = batchClient.listJobs(ListJobsRequest.builder() + .jobQueue(queue) + .jobStatus(status) + .maxResults(PAGE_SIZE) + .nextToken(nextToken) + .build()); + result.addAll(response.jobSummaryList()); + nextToken = response.nextToken(); + } while (nextToken != null); + } + return result; + } + + private List describeJobs(Set jobIds) { + List ids = new ArrayList<>(jobIds); + List result = new ArrayList<>(); + for (int start = 0; start < ids.size(); start += PAGE_SIZE) { + int end = Math.min(start + PAGE_SIZE, ids.size()); + result.addAll(batchClient.describeJobs(DescribeJobsRequest.builder() + .jobs(ids.subList(start, end)) + .build()).jobs()); + } + return result; + } +} diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestApi.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestApi.java index 15258db6..173bbf04 100644 --- a/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestApi.java +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestApi.java @@ -43,6 +43,9 @@ public class RestApi implements ProcessesApi, JobsApi { @Autowired private DownloadJobStatusService downloadJobStatusService; + @Autowired + private DownloadAdmissionService downloadAdmissionService; + @Override // because the produces value in the interface declaration includes "/_" which may // cause exception thrown sometimes. So i re-declared the produces value here @@ -62,7 +65,8 @@ public class RestApi implements ProcessesApi, JobsApi { { "message": {"message": "Job submitted with ID: 123e4567-e89b-12d3-a456-426614174000"}, "status": {"message": "200"}, - "jobID": "123e4567-e89b-12d3-a456-426614174000" + "jobID": "123e4567-e89b-12d3-a456-426614174000", + "queued": false } """))) public ResponseEntity execute( @@ -89,18 +93,27 @@ public ResponseEntity execute( String outputFormat = DatasetDownloadEnums.Parameter.OUTPUT_FORMAT.getStringInput(body); Object multiPolygon = DatasetDownloadEnums.Parameter.MULTI_POLYGON.getObjectInput(body); - String jobId = restServices.downloadData(uuid, key, startDate, endDate, multiPolygon, recipient, - collectionTitle, fullMetadataLink, suggestedCitation, outputFormat); + DownloadRequest request = new DownloadRequest(uuid, key, startDate, endDate, multiPolygon, + recipient, collectionTitle, fullMetadataLink, suggestedCitation, outputFormat); - // The notify user email lives here rather than in data-access-service to make the first - // email faster - // It must only be sent once AWS Batch has accepted the job and returned - // a job id, otherwise we promise the user a file that will never be produced. - restServices.notifyUser(recipient, uuid, key, startDate, endDate, multiPolygon, collectionTitle, fullMetadataLink, suggestedCitation, outputFormat); + // The per-user limit is applied here: a user already at their limit has this + // request held and released later, and gets a job id to poll either way. + // + // The notify user email lives on this side rather than in data-access-service to + // make the first email faster. It goes out with the submit wherever that happens, + // so it is still sent only once AWS Batch has accepted the job and returned a job + // id - otherwise we promise the user a file that will never be produced - and it + // is not sent at all while a download is still waiting for a slot. + DownloadAdmission admission = downloadAdmissionService.submitOrHold(request); - var value = new InlineValue("Job submitted with ID: " + jobId); + // The message keeps its historical wording for a download that went straight + // through, so existing clients that read it see no change. + var value = new InlineValue(admission.queued() + ? "Job queued with ID: " + admission.jobId() + : "Job submitted with ID: " + admission.jobId()); var status = new InlineValue(Integer.toString(HttpStatus.OK.value())); - var results = new DownloadExecutionResponse(value, status, jobId); + var results = new DownloadExecutionResponse( + value, status, admission.jobId(), admission.queued(), admission.queuePosition()); return ResponseEntity.ok(results); diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestServices.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestServices.java index a481fbf0..3a3cfb0a 100644 --- a/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestServices.java +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestServices.java @@ -91,37 +91,40 @@ public void notifyUser(String recipient, String uuid, String key, String startDa } } - public String downloadData( - String id, - String key, - String startDate, - String endDate, - Object polygons, - String recipient, - String collectionTitle, - String fullMetadataLink, - String suggestedCitation, - String outputFormat - ) throws JsonProcessingException { - - // Build the shared subset filters (uuid, key, dates, multi_polygon, output - // format) exactly as the estimate does, then add the download-only fields. + /** + * Build the AWS Batch parameters for a download: the shared subset filters (uuid, key, + * dates, multi_polygon, output format) exactly as the estimate builds them, plus the + * download-only fields. + * + *

Separate from the submit so a request that has to wait for a free slot is validated + * and rendered at accept time, and releasing it later is a plain submit. + */ + public Map buildDownloadParameters(DownloadRequest request) throws JsonProcessingException { Map parameters = SubsetParametersUtils.buildSubsetParameters( - objectMapper, id, key, startDate, endDate, polygons, outputFormat); - parameters.put(DatasetDownloadEnums.Parameter.RECIPIENT.getValue(), recipient); - parameters.put(DatasetDownloadEnums.Parameter.COLLECTION_TITLE.getValue(), collectionTitle); - parameters.put(DatasetDownloadEnums.Parameter.FULL_METADATA_LINK.getValue(), fullMetadataLink); - parameters.put(DatasetDownloadEnums.Parameter.SUGGESTED_CITATION.getValue(), suggestedCitation); + objectMapper, request.uuid(), request.key(), request.startDate(), request.endDate(), + request.multiPolygon(), request.outputFormat()); + parameters.put(DatasetDownloadEnums.Parameter.RECIPIENT.getValue(), request.recipient()); + parameters.put(DatasetDownloadEnums.Parameter.COLLECTION_TITLE.getValue(), request.collectionTitle()); + parameters.put(DatasetDownloadEnums.Parameter.FULL_METADATA_LINK.getValue(), request.fullMetadataLink()); + parameters.put(DatasetDownloadEnums.Parameter.SUGGESTED_CITATION.getValue(), request.suggestedCitation()); parameters.put( DatasetDownloadEnums.Parameter.TYPE.getValue(), DatasetDownloadEnums.Type.SUB_SETTING.getValue() ); + return parameters; + } + + /** + * The AWS Batch job name for a download. Note this sanitises the address, so it is not a + * safe key for the owning user - read the recipient job parameter instead. + */ + public static String downloadJobName(String recipient) { + return "generating-data-file-for-" + recipient.replaceAll("[^a-zA-Z0-9-_]", "-"); + } - String jobId = submitJob( - "generating-data-file-for-" + recipient.replaceAll("[^a-zA-Z0-9-_]", "-"), - this.batchJobQueue, - this.batchJobDefinition, - parameters); + /** Submit a prepared download to the configured queue and job definition. */ + public String submitDownloadJob(String jobName, Map parameters) { + String jobId = submitJob(jobName, this.batchJobQueue, this.batchJobDefinition, parameters); log.info("Job submitted with ID: {}", jobId); return jobId; } diff --git a/server/src/main/resources/application.yaml b/server/src/main/resources/application.yaml index f3dadb08..cc812b9b 100644 --- a/server/src/main/resources/application.yaml +++ b/server/src/main/resources/application.yaml @@ -59,6 +59,19 @@ aws: # Defaults to the initial queue in BatchJobProperties. Dev deployments whose DAS # uses generate-csv-data-file must override this value explicitly. child-queue: data-access-service-batch-job-queue + # Per-user download concurrency. Requests past max-concurrent are held in memory and + # submitted as that user's slots free, so they are queued rather than rejected. The + # hold queue does not survive a restart. + user-limit: + enabled: true + max-concurrent: 10 + # How often held downloads are reconsidered, and how long the shared in-flight + # snapshot is reused before the queues are swept again. + release-interval: 15s + # A download waiting longer than this is abandoned. + max-hold-age: 24h + # Safety valve on total memory held, across all users. + max-held-total: 1000 wfs-default-param: fields: diff --git a/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionServiceTest.java b/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionServiceTest.java new file mode 100644 index 00000000..381d2015 --- /dev/null +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionServiceTest.java @@ -0,0 +1,385 @@ +package au.org.aodn.ogcapi.server.processes; + +import com.fasterxml.jackson.core.JsonProcessingException; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyString; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.times; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + +@ExtendWith(MockitoExtension.class) +class DownloadAdmissionServiceTest { + + private static final Instant NOW = Instant.parse("2026-08-24T02:00:00Z"); + private static final String RECIPIENT = "person@example.com"; + private static final String RECIPIENT_JOB_NAME = "generating-data-file-for-person-example-com"; + private static final String OTHER_RECIPIENT = "someone.else@example.com"; + + @Mock + private RestServices restServices; + + @Mock + private InFlightDownloadCounter counter; + + /** In-flight downloads per recipient, standing in for what the counter would report. */ + private final Map inFlight = new HashMap<>(); + private final List submittedJobNames = new ArrayList<>(); + + private MutableTestClock clock; + private DownloadAdmissionService service; + + @BeforeEach + void setUp() throws JsonProcessingException { + lenient().when(restServices.buildDownloadParameters(any())) + .thenAnswer(invocation -> new HashMap()); + // A submit hands back a fresh AWS job id and, exactly as the real counter does, + // immediately makes that job count towards its recipient. + lenient().when(restServices.submitDownloadJob(anyString(), any())).thenAnswer(invocation -> { + submittedJobNames.add(invocation.getArgument(0)); + return UUID.randomUUID().toString(); + }); + lenient().when(counter.countInFlight(anyString())) + .thenAnswer(invocation -> current(invocation.getArgument(0)).get()); + lenient().doAnswer(invocation -> { + current(invocation.getArgument(1)).incrementAndGet(); + return null; + }).when(counter).recordSubmitted(anyString(), anyString()); + + clock = new MutableTestClock(NOW); + service = build(limits(true, 10)); + } + + private DownloadAdmissionService build(DownloadLimitProperties limits) { + return new DownloadAdmissionService(restServices, counter, limits, clock); + } + + private static DownloadLimitProperties limits(boolean enabled, int maxConcurrent) { + return new DownloadLimitProperties( + enabled, maxConcurrent, Duration.ofSeconds(15), Duration.ofHours(24), 1000); + } + + private AtomicInteger current(String recipient) { + // Key the way the real counter does, so these tests cannot accidentally rely on + // case-sensitive bookkeeping the production class does not have. + return inFlight.computeIfAbsent( + InFlightDownloadCounter.recipientKey(recipient), key -> new AtomicInteger()); + } + + private DownloadRequest request(String recipient) { + return new DownloadRequest("collection-id", "key.zarr", "2023-01-01", "2023-01-31", + "non-specified", recipient, "Test Collection", + "https://portal.example.test/details/collection-id", "Cite as", "netcdf"); + } + + @Test + void underTheLimitSubmitsStraightAwayAndReturnsTheAwsJobId() throws Exception { + String jobId = service.submitOrHold(request(RECIPIENT)).jobId(); + + assertNotNull(jobId); + assertNull(service.findHeld(jobId)); + verify(restServices).submitDownloadJob(eq(RECIPIENT_JOB_NAME), any()); + verify(restServices).notifyUser(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); + assertEquals(List.of(RECIPIENT_JOB_NAME), submittedJobNames); + } + + @Test + void atTheLimitHoldsInsteadOfRejectingAndStillReturnsAJobId() throws Exception { + current(RECIPIENT).set(10); + + String jobId = service.submitOrHold(request(RECIPIENT)).jobId(); + + assertNotNull(jobId); + // The id has to be a canonical lowercase UUID or the status endpoint will not accept it. + assertEquals(jobId, UUID.fromString(jobId).toString()); + assertNotNull(service.findHeld(jobId)); + verify(restServices, never()).submitDownloadJob(anyString(), any()); + // Nothing has started, so the user must not be told their file is being produced. + verify(restServices, never()).notifyUser(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); + } + + @Test + void anImmediateSubmitIsReportedAsNotQueued() throws Exception { + DownloadAdmission admission = service.submitOrHold(request(RECIPIENT)); + + assertFalse(admission.queued()); + assertNull(admission.queuePosition()); + } + + @Test + void aHeldDownloadReportsItsPlaceInTheUsersOwnQueue() throws Exception { + current(RECIPIENT).set(10); + + DownloadAdmission first = service.submitOrHold(request(RECIPIENT)); + DownloadAdmission second = service.submitOrHold(request(RECIPIENT)); + DownloadAdmission third = service.submitOrHold(request(RECIPIENT)); + + assertTrue(first.queued()); + assertEquals(1, first.queuePosition()); + assertEquals(2, second.queuePosition()); + assertEquals(3, third.queuePosition()); + } + + @Test + void queuePositionCountsOnlyTheSameUsersDownloads() throws Exception { + current(RECIPIENT).set(10); + current(OTHER_RECIPIENT).set(10); + + service.submitOrHold(request(OTHER_RECIPIENT)); + service.submitOrHold(request(OTHER_RECIPIENT)); + DownloadAdmission mine = service.submitOrHold(request(RECIPIENT)); + + // Two other people are waiting ahead in the shared queue, but they do not delay this + // one: it is released when its own owner's slots free. + assertEquals(1, mine.queuePosition()); + } + + @Test + void queuePositionMovesUpAsTheUsersEarlierDownloadsAreReleased() throws Exception { + current(RECIPIENT).set(10); + service.submitOrHold(request(RECIPIENT)); + String second = service.submitOrHold(request(RECIPIENT)).jobId(); + assertEquals(2, service.findHeld(second).position()); + + current(RECIPIENT).set(9); + service.releaseHeldDownloads(); + + assertEquals(1, service.findHeld(second).position(), "it should now be next"); + } + + @Test + void theEleventhDownloadIsHeldAndTheFirstTenAreNot() throws Exception { + List ids = new ArrayList<>(); + for (int i = 0; i < 11; i++) { + ids.add(service.submitOrHold(request(RECIPIENT)).jobId()); + } + + verify(restServices, times(10)).submitDownloadJob(anyString(), any()); + for (int i = 0; i < 10; i++) { + assertNull(service.findHeld(ids.get(i)), "download " + i + " should have been submitted"); + } + assertNotNull(service.findHeld(ids.get(10)), "the eleventh download should be held"); + } + + @Test + void aDifferentlyCasedAddressQueuesBehindTheSameUsersHeldDownloads() throws Exception { + current(RECIPIENT).set(10); + String first = service.submitOrHold(request(RECIPIENT)).jobId(); + + // A slot frees, but the same person writing their address differently must still + // queue behind their own earlier request rather than overtake it. + current(RECIPIENT).set(0); + String second = service.submitOrHold(request("Person@Example.COM")).jobId(); + + assertNotNull(service.findHeld(first)); + assertNotNull(service.findHeld(second)); + verify(restServices, never()).submitDownloadJob(anyString(), any()); + } + + @Test + void oneUserAtTheLimitDoesNotBlockAnother() throws Exception { + current(RECIPIENT).set(10); + + String held = service.submitOrHold(request(RECIPIENT)).jobId(); + String submitted = service.submitOrHold(request(OTHER_RECIPIENT)).jobId(); + + assertNotNull(service.findHeld(held)); + assertNull(service.findHeld(submitted)); + } + + @Test + void aNewRequestNeverJumpsAheadOfTheSameUsersOwnQueue() throws Exception { + current(RECIPIENT).set(10); + String first = service.submitOrHold(request(RECIPIENT)).jobId(); + + // A slot frees, but the second request still queues behind the first. + current(RECIPIENT).set(0); + String second = service.submitOrHold(request(RECIPIENT)).jobId(); + + assertNotNull(service.findHeld(first)); + assertNotNull(service.findHeld(second)); + verify(restServices, never()).submitDownloadJob(anyString(), any()); + } + + @Test + void theReleaseLoopSubmitsHeldDownloadsAsSlotsFree() throws Exception { + current(RECIPIENT).set(10); + String heldId = service.submitOrHold(request(RECIPIENT)).jobId(); + + // Still full: nothing is released. + service.releaseHeldDownloads(); + assertNotNull(service.findHeld(heldId)); + verify(restServices, never()).submitDownloadJob(anyString(), any()); + + // A job finished, so the held one goes. + current(RECIPIENT).set(9); + service.releaseHeldDownloads(); + + assertNull(service.findHeld(heldId)); + assertNotNull(service.awsJobIdOf(heldId)); + assertNotEquals(heldId, service.awsJobIdOf(heldId)); + verify(restServices).submitDownloadJob(anyString(), any()); + // The started email goes out now, when the download actually starts. + verify(restServices).notifyUser(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); + } + + @Test + void theReleaseLoopReleasesOnlyAsManyAsThereAreFreeSlots() throws Exception { + current(RECIPIENT).set(10); + List heldIds = new ArrayList<>(); + for (int i = 0; i < 5; i++) { + heldIds.add(service.submitOrHold(request(RECIPIENT)).jobId()); + } + + // Two slots free in one go; the other three must stay put. + current(RECIPIENT).set(8); + service.releaseHeldDownloads(); + + verify(restServices, times(2)).submitDownloadJob(anyString(), any()); + assertNull(service.findHeld(heldIds.get(0))); + assertNull(service.findHeld(heldIds.get(1))); + assertNotNull(service.findHeld(heldIds.get(2))); + assertNotNull(service.findHeld(heldIds.get(3))); + assertNotNull(service.findHeld(heldIds.get(4))); + } + + @Test + void heldDownloadsAreReleasedInTheOrderTheyArrived() throws Exception { + current(RECIPIENT).set(10); + String first = service.submitOrHold(request(RECIPIENT)).jobId(); + String second = service.submitOrHold(request(RECIPIENT)).jobId(); + + current(RECIPIENT).set(9); + service.releaseHeldDownloads(); + + assertNull(service.findHeld(first)); + assertNotNull(service.findHeld(second)); + } + + @Test + void theReleaseLoopMakesNoAwsCallWhenNothingIsHeld() { + service.releaseHeldDownloads(); + + verify(counter, never()).refresh(); + verify(restServices, never()).submitDownloadJob(anyString(), any()); + } + + @Test + void aDownloadHeldPastTheMaximumHoldAgeIsAbandoned() throws Exception { + DownloadAdmissionService shortLived = build(new DownloadLimitProperties( + true, 10, Duration.ofSeconds(15), Duration.ofMinutes(30), 1000)); + current(RECIPIENT).set(10); + String heldId = shortLived.submitOrHold(request(RECIPIENT)).jobId(); + assertNotNull(shortLived.findHeld(heldId)); + + clock.advance(Duration.ofMinutes(31)); + // Free the slot too, so the only reason it is not released is that it expired. + current(RECIPIENT).set(0); + shortLived.releaseHeldDownloads(); + + assertNull(shortLived.findHeld(heldId)); + assertNull(shortLived.awsJobIdOf(heldId)); + verify(restServices, never()).submitDownloadJob(anyString(), any()); + } + + @Test + void aFailedReleaseKeepsTheDownloadHeldForAnotherAttempt() throws Exception { + current(RECIPIENT).set(10); + String heldId = service.submitOrHold(request(RECIPIENT)).jobId(); + + when(restServices.submitDownloadJob(anyString(), any())) + .thenThrow(new IllegalStateException("AWS Batch rejected the job")); + current(RECIPIENT).set(0); + service.releaseHeldDownloads(); + + assertNotNull(service.findHeld(heldId), "a failed submit must not lose the download"); + assertNull(service.awsJobIdOf(heldId)); + } + + @Test + void aReleaseThatKeepsFailingIsEventuallyAbandoned() throws Exception { + current(RECIPIENT).set(10); + String heldId = service.submitOrHold(request(RECIPIENT)).jobId(); + + when(restServices.submitDownloadJob(anyString(), any())) + .thenThrow(new IllegalStateException("AWS Batch rejected the job")); + current(RECIPIENT).set(0); + service.releaseHeldDownloads(); + service.releaseHeldDownloads(); + service.releaseHeldDownloads(); + + assertNull(service.findHeld(heldId)); + } + + @Test + void theHoldQueueIsBounded() throws Exception { + DownloadAdmissionService bounded = build(new DownloadLimitProperties( + true, 1, Duration.ofSeconds(15), Duration.ofHours(24), 2)); + current(RECIPIENT).set(1); + + bounded.submitOrHold(request(RECIPIENT)); + bounded.submitOrHold(request(RECIPIENT)); + + assertThrows(IllegalStateException.class, () -> bounded.submitOrHold(request(RECIPIENT))); + } + + @Test + void theLimitCanBeTurnedOffEntirely() throws Exception { + DownloadAdmissionService disabled = build(limits(false, 10)); + current(RECIPIENT).set(500); + + String jobId = disabled.submitOrHold(request(RECIPIENT)).jobId(); + + assertNull(disabled.findHeld(jobId)); + verify(restServices).submitDownloadJob(anyString(), any()); + verify(counter, never()).refreshIfStale(); + } + + @Test + void theSweepHappensBeforeTheAdmissionDecision() throws Exception { + service.submitOrHold(request(RECIPIENT)); + + verify(counter).refreshIfStale(); + } + + @Test + void aReleasedDownloadIsSubmittedWithTheSameNameAndRecipientItWasAcceptedWith() throws Exception { + current(RECIPIENT).set(10); + service.submitOrHold(request(RECIPIENT)); + + current(RECIPIENT).set(0); + service.releaseHeldDownloads(); + + assertEquals(List.of(RECIPIENT_JOB_NAME), submittedJobNames); + ArgumentCaptor recipientCaptor = ArgumentCaptor.forClass(String.class); + verify(restServices).notifyUser(recipientCaptor.capture(), any(), any(), any(), any(), + any(), any(), any(), any(), any()); + assertEquals(RECIPIENT, recipientCaptor.getValue()); + } +} diff --git a/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusServiceTest.java b/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusServiceTest.java index 4d6c073c..f325a4e3 100644 --- a/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusServiceTest.java +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadJobStatusServiceTest.java @@ -31,6 +31,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.ArgumentMatchers.any; @@ -55,6 +56,9 @@ class DownloadJobStatusServiceTest { @Mock private BatchClient batchClient; + @Mock + private DownloadAdmissionService admissionService; + private final Map describedJobs = new HashMap<>(); private final Map listedJobs = new HashMap<>(); private DownloadJobStatusService service; @@ -80,6 +84,7 @@ void setUp() { batchClient, new BatchJobProperties(QUEUE_NAME, DEFINITION_NAME, CHILD_QUEUE_NAME), new DownloadJobStatusAggregator(), + admissionService, Clock.fixed(NOW, ZoneOffset.UTC)); } @@ -137,6 +142,7 @@ void configuredVersionedArnsRequireExactMatches() { batchClient, new BatchJobProperties(QUEUE_ARN, DEFINITION_ARN, CHILD_QUEUE_NAME), new DownloadJobStatusAggregator(), + admissionService, Clock.fixed(NOW, ZoneOffset.UTC)); describedJobs.put(JOB_ID, initial(JobStatus.SUCCEEDED, "one.zarr", NOW.minusSeconds(5))); assertEquals(StatusCode.SUCCESSFUL, service.getStatus(JOB_ID).getStatus()); @@ -280,6 +286,61 @@ void usesConfiguredChildQueueAndBuildsTerminalDatesWithoutSensitiveFields() { assertTrue(captor.getAllValues().stream().allMatch(request -> CHILD_QUEUE_NAME.equals(request.jobQueue()))); } + @Test + void heldDownloadReportsAcceptedWithoutCallingAws() { + String heldId = "9f1d3a2c-0000-4000-8000-000000000001"; + Map parameters = Map.of( + "collection_title", "Test Ocean Data Collection", + "key", "satellite_wind_altimeter_delayed_qc.zarr", + "output_format", "netcdf", + "full_metadata_link", "https://portal.example.test/details/collection-id"); + when(admissionService.findHeld(heldId)).thenReturn(new DownloadAdmissionService.HeldView( + new HeldDownload( + heldId, + new DownloadRequest("collection-id", null, null, null, null, + "person@example.com", null, null, null, "netcdf"), + "generating-data-file-for-person-example-com", + parameters, + NOW.minusSeconds(30), + 0), + 2)); + + DownloadJobStatusInfo status = service.getStatus(heldId); + + assertEquals(StatusCode.ACCEPTED, status.getStatus()); + assertEquals(DownloadAdmissionService.QUEUED_MESSAGE, status.getMessage()); + assertEquals(heldId, status.getJobID()); + // Machine readable, so the portal never has to match on the message text. + assertTrue(status.isQueued()); + assertEquals(2, status.getQueuePosition()); + // The display fields come from the parameters built when the download was accepted, + // so a waiting job looks the same on the status page as a running one. + assertEquals("Test Ocean Data Collection", status.getCollection()); + assertEquals("satellite_wind_altimeter_delayed_qc.zarr", status.getDataSelection()); + assertEquals("netcdf", status.getFormat()); + assertEquals("https://portal.example.test/details/collection-id", status.getMetadataUrl()); + assertEquals(Date.from(NOW.minusSeconds(30)), status.getCreated()); + assertNull(status.getStarted()); + assertNull(status.getFinished()); + verify(batchClient, never()).describeJobs(any(DescribeJobsRequest.class)); + verify(batchClient, never()).listJobs(any(ListJobsRequest.class)); + } + + @Test + void releasedDownloadResolvesToItsAwsJobButAnswersOnTheCallerJobId() { + String heldId = "9f1d3a2c-0000-4000-8000-000000000002"; + when(admissionService.awsJobIdOf(heldId)).thenReturn(JOB_ID); + describedJobs.put(JOB_ID, initial(JobStatus.RUNNING, null, null)); + + DownloadJobStatusInfo status = service.getStatus(heldId); + + assertEquals(StatusCode.RUNNING, status.getStatus()); + // The caller polls the id it was handed, never the AWS job the release created. + assertEquals(heldId, status.getJobID()); + assertFalse(status.isQueued()); + assertNull(status.getQueuePosition()); + } + @Test void queueAndDefinitionNormalizersHandleNamesAndArns() { assertTrue(DownloadJobStatusService.matchesQueue(QUEUE_NAME, QUEUE_ARN)); diff --git a/server/src/test/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounterTest.java b/server/src/test/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounterTest.java new file mode 100644 index 00000000..8d33f785 --- /dev/null +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounterTest.java @@ -0,0 +1,298 @@ +package au.org.aodn.ogcapi.server.processes; + +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import software.amazon.awssdk.services.batch.BatchClient; +import software.amazon.awssdk.services.batch.model.DescribeJobsRequest; +import software.amazon.awssdk.services.batch.model.DescribeJobsResponse; +import software.amazon.awssdk.services.batch.model.JobDetail; +import software.amazon.awssdk.services.batch.model.JobStatus; +import software.amazon.awssdk.services.batch.model.JobSummary; +import software.amazon.awssdk.services.batch.model.ListJobsRequest; +import software.amazon.awssdk.services.batch.model.ListJobsResponse; + +import java.time.Duration; +import java.time.Instant; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertNull; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.clearInvocations; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.lenient; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; + +@ExtendWith(MockitoExtension.class) +class InFlightDownloadCounterTest { + + private static final String QUEUE_NAME = "initial-queue"; + private static final String CHILD_QUEUE_NAME = "child-queue"; + private static final String QUEUE_ARN = "arn:aws:batch:ap-southeast-2:123456789012:job-queue/" + QUEUE_NAME; + private static final String DEFINITION_NAME = "download-definition"; + private static final String DEFINITION_ARN = + "arn:aws:batch:ap-southeast-2:123456789012:job-definition/" + DEFINITION_NAME + ":7"; + private static final Instant NOW = Instant.parse("2026-08-24T02:00:00Z"); + + private static final String ALICE = "alice@example.com"; + private static final String BOB = "bob@example.com"; + + @Mock + private BatchClient batchClient; + + /** Jobs DescribeJobs will answer with, by job id. */ + private final Map describedJobs = new HashMap<>(); + + /** Non-terminal job summaries per queue, as ListJobs would page them out. */ + private final Map> listedJobs = new HashMap<>(); + + private MutableTestClock clock; + private InFlightDownloadCounter counter; + + @BeforeEach + void setUp() { + lenient().when(batchClient.describeJobs(any(DescribeJobsRequest.class))).thenAnswer(invocation -> { + DescribeJobsRequest request = invocation.getArgument(0); + List jobs = request.jobs().stream() + .map(describedJobs::get) + .filter(job -> job != null) + .toList(); + return DescribeJobsResponse.builder().jobs(jobs).build(); + }); + // The counter lists by status and never by name filter, which is what distinguishes + // it from the status service. Every summary is reported under RUNNING so one entry + // per queue is enough to describe the fixture. + lenient().when(batchClient.listJobs(any(ListJobsRequest.class))).thenAnswer(invocation -> { + ListJobsRequest request = invocation.getArgument(0); + if (request.jobStatus() != JobStatus.RUNNING) { + return ListJobsResponse.builder().build(); + } + return ListJobsResponse.builder() + .jobSummaryList(listedJobs.getOrDefault(request.jobQueue(), List.of())) + .build(); + }); + + clock = new MutableTestClock(NOW); + counter = newCounter(); + } + + private InFlightDownloadCounter newCounter() { + return new InFlightDownloadCounter( + batchClient, + new BatchJobProperties(QUEUE_NAME, DEFINITION_NAME, CHILD_QUEUE_NAME), + new DownloadLimitProperties(true, 10, Duration.ofSeconds(15), Duration.ofHours(24), 1000), + clock); + } + + private void master(String jobId, String recipient) { + describedJobs.put(jobId, JobDetail.builder() + .jobId(jobId) + .jobName(RestServices.downloadJobName(recipient)) + .jobQueue(QUEUE_ARN) + .jobDefinition(DEFINITION_ARN) + .status(JobStatus.RUNNING) + .parameters(Map.of("type", "sub-setting", "recipient", recipient)) + .build()); + } + + private void onQueue(String queue, String jobId, String jobName) { + listedJobs.computeIfAbsent(queue, key -> new ArrayList<>()) + .add(JobSummary.builder().jobId(jobId).jobName(jobName).build()); + } + + @Test + void countsMasterJobsStillOnTheDownloadQueue() { + master("m1", ALICE); + master("m2", ALICE); + master("m3", BOB); + onQueue(QUEUE_NAME, "m1", RestServices.downloadJobName(ALICE)); + onQueue(QUEUE_NAME, "m2", RestServices.downloadJobName(ALICE)); + onQueue(QUEUE_NAME, "m3", RestServices.downloadJobName(BOB)); + + counter.refresh(); + + assertEquals(2, counter.countInFlight(ALICE)); + assertEquals(1, counter.countInFlight(BOB)); + } + + @Test + void countsAWorkflowWhoseMasterHasFinishedButWhoseChildrenAreStillRunning() { + // The master succeeded and left the download queue; only its prepare child is live. + master("m1", ALICE); + onQueue(CHILD_QUEUE_NAME, "c1", "prepare-data-for-job-m1"); + + counter.refresh(); + + assertEquals(1, counter.countInFlight(ALICE)); + } + + @Test + void countsAWorkflowOnlyOnceWhenBothItsChildrenAreRunning() { + master("m1", ALICE); + onQueue(CHILD_QUEUE_NAME, "c1", "prepare-data-for-job-m1"); + onQueue(CHILD_QUEUE_NAME, "c2", "collect-data-for-job-m1"); + + counter.refresh(); + + assertEquals(1, counter.countInFlight(ALICE)); + } + + @Test + void countsAWorkflowOnlyOnceWhenTheMasterAndItsChildAreBothLive() { + master("m1", ALICE); + onQueue(QUEUE_NAME, "m1", RestServices.downloadJobName(ALICE)); + onQueue(CHILD_QUEUE_NAME, "c1", "prepare-data-for-job-m1"); + + counter.refresh(); + + assertEquals(1, counter.countInFlight(ALICE)); + } + + @Test + void ignoresJobsThatAreNotDownloadMasters() { + // Something else entirely, sharing the queue. + describedJobs.put("x1", JobDetail.builder() + .jobId("x1") + .jobName("some-other-workload") + .jobQueue(QUEUE_ARN) + .jobDefinition(DEFINITION_ARN) + .status(JobStatus.RUNNING) + .parameters(Map.of("type", "something-else", "recipient", ALICE)) + .build()); + onQueue(QUEUE_NAME, "x1", "some-other-workload"); + + counter.refresh(); + + assertEquals(0, counter.countInFlight(ALICE)); + } + + @Test + void takesTheOwnerFromTheRecipientParameterNotTheSanitisedJobName() { + // Both addresses sanitise to generating-data-file-for-a-b-x-com, so counting by job + // name would merge two different users into one. + String first = "a.b@x.com"; + String second = "a-b@x-com"; + assertEquals(RestServices.downloadJobName(first), RestServices.downloadJobName(second)); + + master("m1", first); + master("m2", second); + onQueue(QUEUE_NAME, "m1", RestServices.downloadJobName(first)); + onQueue(QUEUE_NAME, "m2", RestServices.downloadJobName(second)); + + counter.refresh(); + + assertEquals(1, counter.countInFlight(first)); + assertEquals(1, counter.countInFlight(second)); + } + + @Test + void theSameAddressInDifferentCaseOrWithSpacesIsOneUser() { + // Otherwise a single capital letter silently buys a second allowance of slots. + master("m1", "Alice@Example.com"); + master("m2", "alice@example.com"); + onQueue(QUEUE_NAME, "m1", RestServices.downloadJobName("Alice@Example.com")); + onQueue(QUEUE_NAME, "m2", RestServices.downloadJobName("alice@example.com")); + + counter.refresh(); + + assertEquals(2, counter.countInFlight(ALICE)); + assertEquals(2, counter.countInFlight("ALICE@EXAMPLE.COM")); + assertEquals(2, counter.countInFlight(" alice@example.com ")); + } + + @Test + void aJustSubmittedJobCountsBeforeAnySweepCanSeeIt() { + counter.refresh(); + assertEquals(0, counter.countInFlight(ALICE)); + + counter.recordSubmitted("m-new", ALICE); + + assertEquals(1, counter.countInFlight(ALICE)); + } + + @Test + void aJustSubmittedJobIsNotCountedTwiceOnceTheSweepSeesIt() { + counter.recordSubmitted("m1", ALICE); + master("m1", ALICE); + onQueue(QUEUE_NAME, "m1", RestServices.downloadJobName(ALICE)); + + counter.refresh(); + + assertEquals(1, counter.countInFlight(ALICE)); + } + + @Test + void aRecentSubmissionStopsCountingOnceItsGraceHasPassed() { + counter.recordSubmitted("m1", ALICE); + assertEquals(1, counter.countInFlight(ALICE)); + + // Past the grace the sweep is authoritative again, so a job AWS no longer reports as + // non-terminal stops holding a slot. + clock.advance(InFlightDownloadCounter.SUBMIT_GRACE.plusSeconds(1)); + + assertEquals(0, counter.countInFlight(ALICE)); + } + + @Test + void everyNonTerminalStatusIsListed() { + counter.refresh(); + + ArgumentCaptor captor = ArgumentCaptor.forClass(ListJobsRequest.class); + verify(batchClient, atLeastOnce()).listJobs(captor.capture()); + List statuses = captor.getAllValues().stream() + .filter(request -> QUEUE_NAME.equals(request.jobQueue())) + .map(ListJobsRequest::jobStatus) + .toList(); + assertTrue(statuses.containsAll(List.of( + JobStatus.SUBMITTED, JobStatus.PENDING, JobStatus.RUNNABLE, + JobStatus.STARTING, JobStatus.RUNNING))); + // Terminal states must never be swept: a finished download frees its slot. + assertTrue(statuses.stream().noneMatch(status -> + status == JobStatus.SUCCEEDED || status == JobStatus.FAILED)); + // Counting never filters by job name; that is the status service's query. + assertTrue(captor.getAllValues().stream().allMatch(request -> request.filters().isEmpty())); + } + + @Test + void aFailedSweepKeepsServingThePreviousCount() { + master("m1", ALICE); + onQueue(QUEUE_NAME, "m1", RestServices.downloadJobName(ALICE)); + counter.refresh(); + assertEquals(1, counter.countInFlight(ALICE)); + + doThrow(new RuntimeException("AWS is having a moment")) + .when(batchClient).listJobs(any(ListJobsRequest.class)); + counter.refresh(); + + assertEquals(1, counter.countInFlight(ALICE), "a failed sweep must not zero the count"); + } + + @Test + void aFreshSnapshotIsNotSweptAgain() { + counter.refresh(); + clearInvocations(batchClient); + + counter.refreshIfStale(); + + verify(batchClient, never()).listJobs(any(ListJobsRequest.class)); + } + + @Test + void childNamesResolveToTheirMasterJob() { + assertEquals("abc", InFlightDownloadCounter.masterIdOf("prepare-data-for-job-abc")); + assertEquals("abc", InFlightDownloadCounter.masterIdOf("collect-data-for-job-abc")); + assertNull(InFlightDownloadCounter.masterIdOf("generating-data-file-for-someone")); + assertNull(InFlightDownloadCounter.masterIdOf("prepare-data-for-job-")); + assertNull(InFlightDownloadCounter.masterIdOf(null)); + } +} diff --git a/server/src/test/java/au/org/aodn/ogcapi/server/processes/MutableTestClock.java b/server/src/test/java/au/org/aodn/ogcapi/server/processes/MutableTestClock.java new file mode 100644 index 00000000..6342cb6c --- /dev/null +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/MutableTestClock.java @@ -0,0 +1,40 @@ +package au.org.aodn.ogcapi.server.processes; + +import java.time.Clock; +import java.time.Duration; +import java.time.Instant; +import java.time.ZoneId; +import java.time.ZoneOffset; + +/** + * A clock the test can wind forward. The download services take a {@link Clock} so their + * time-based behaviour - hold expiry, the grace on a just-submitted job, snapshot staleness - + * can be exercised without sleeping. + */ +final class MutableTestClock extends Clock { + + private Instant instant; + + MutableTestClock(Instant instant) { + this.instant = instant; + } + + @Override + public ZoneId getZone() { + return ZoneOffset.UTC; + } + + @Override + public Clock withZone(ZoneId zone) { + return this; + } + + @Override + public Instant instant() { + return instant; + } + + void advance(Duration amount) { + instant = instant.plus(amount); + } +} diff --git a/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiJobsTest.java b/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiJobsTest.java index 76be1f4f..50639c6a 100644 --- a/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiJobsTest.java +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiJobsTest.java @@ -38,6 +38,9 @@ class RestApiJobsTest { @Mock private DownloadJobStatusService downloadJobStatusService; + @Mock + private DownloadAdmissionService downloadAdmissionService; + private final ObjectMapper objectMapper = new ObjectMapper(); private MockMvc mockMvc; @@ -46,6 +49,7 @@ void setUp() { RestApi restApi = new RestApi(); ReflectionTestUtils.setField(restApi, "restServices", restServices); ReflectionTestUtils.setField(restApi, "downloadJobStatusService", downloadJobStatusService); + ReflectionTestUtils.setField(restApi, "downloadAdmissionService", downloadAdmissionService); mockMvc = MockMvcBuilders.standaloneSetup(restApi) .setControllerAdvice(new GlobalExceptionHandler()) .build(); @@ -53,8 +57,8 @@ void setUp() { @Test void postKeepsExistingFieldsAndAddsPureJobId() throws Exception { - when(restServices.downloadData(any(), any(), any(), any(), any(), any(), any(), any(), any(), any())) - .thenReturn(JOB_ID); + when(downloadAdmissionService.submitOrHold(any())) + .thenReturn(new DownloadAdmission(JOB_ID, false, null)); String body = objectMapper.writeValueAsString(Map.of("inputs", Map.of( "uuid", "collection-id", "recipient", "person@example.com"))); @@ -69,6 +73,41 @@ void postKeepsExistingFieldsAndAddsPureJobId() throws Exception { .andExpect(jsonPath("$.jobID").value(JOB_ID)); } + @Test + void postSurfacesTheQueuedStateAndPosition() throws Exception { + when(downloadAdmissionService.submitOrHold(any())) + .thenReturn(new DownloadAdmission(JOB_ID, true, 3)); + String body = objectMapper.writeValueAsString(Map.of("inputs", Map.of( + "uuid", "collection-id", + "recipient", "person@example.com"))); + + mockMvc.perform(post("/api/v1/ogc/processes/download/execution") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON) + .content(body)) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.jobID").value(JOB_ID)) + .andExpect(jsonPath("$.queued").value(true)) + .andExpect(jsonPath("$.queuePosition").value(3)); + } + + @Test + void postOmitsQueuePositionWhenTheDownloadWentStraightThrough() throws Exception { + when(downloadAdmissionService.submitOrHold(any())) + .thenReturn(new DownloadAdmission(JOB_ID, false, null)); + String body = objectMapper.writeValueAsString(Map.of("inputs", Map.of( + "uuid", "collection-id", + "recipient", "person@example.com"))); + + mockMvc.perform(post("/api/v1/ogc/processes/download/execution") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON) + .content(body)) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.queued").value(false)) + .andExpect(jsonPath("$.queuePosition").doesNotExist()); + } + @Test void getStatusSerializesExtendedStatusInfo() throws Exception { DownloadJobStatusInfo statusInfo = new DownloadJobStatusInfo(); diff --git a/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiTest.java b/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiTest.java index d2de112d..e9abcfc6 100644 --- a/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiTest.java +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestApiTest.java @@ -12,7 +12,7 @@ import org.junit.jupiter.api.BeforeEach; import org.junit.jupiter.api.Test; import org.junit.jupiter.api.extension.ExtendWith; -import org.mockito.InOrder; +import org.mockito.ArgumentCaptor; import org.mockito.InjectMocks; import org.mockito.Mock; import org.mockito.junit.jupiter.MockitoExtension; @@ -24,7 +24,6 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; import static org.mockito.ArgumentMatchers.any; -import static org.mockito.Mockito.inOrder; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -35,6 +34,9 @@ public class RestApiTest { @Mock private RestServices restServices; + @Mock + private DownloadAdmissionService downloadAdmissionService; + @InjectMocks private RestApi restApi; @@ -55,8 +57,8 @@ public void setUp() { @Test public void testExecuteDownloadDatasetSuccess() throws JsonProcessingException { - when(restServices.downloadData(any(), any(), any(), any(), any(), any(), any(), any(), any(), any())) - .thenReturn("test-job-id"); + when(downloadAdmissionService.submitOrHold(any())) + .thenReturn(new DownloadAdmission("test-job-id", false, null)); ResponseEntity response = restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); @@ -67,16 +69,28 @@ public void testExecuteDownloadDatasetSuccess() throws JsonProcessingException { assertEquals("Job submitted with ID: test-job-id", results.message().message()); assertEquals("200", results.status().message()); assertEquals("test-job-id", results.jobId()); + } - // The "processing started" email must go out only after AWS Batch returned a job id - InOrder inOrder = inOrder(restServices); - inOrder.verify(restServices).downloadData(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); - inOrder.verify(restServices).notifyUser(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); + @Test + public void testExecutePassesEveryRequestInputToAdmission() throws JsonProcessingException { + when(downloadAdmissionService.submitOrHold(any())) + .thenReturn(new DownloadAdmission("test-job-id", false, null)); + + restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); + + ArgumentCaptor captor = ArgumentCaptor.forClass(DownloadRequest.class); + verify(downloadAdmissionService).submitOrHold(captor.capture()); + DownloadRequest request = captor.getValue(); + assertEquals("test-uuid", request.uuid()); + assertEquals("2023-01-01", request.startDate()); + assertEquals("2023-01-31", request.endDate()); + assertEquals("test-multipolygon", request.multiPolygon()); + assertEquals("test@example.com", request.recipient()); } @Test public void testExecuteDownloadDatasetError() throws JsonProcessingException { - when(restServices.downloadData(any(), any(), any(), any(), any(), any(), any(), any(), any(), any())) + when(downloadAdmissionService.submitOrHold(any())) .thenThrow(new RuntimeException("Error while getting dataset")); ResponseEntity response = restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); @@ -87,7 +101,8 @@ public void testExecuteDownloadDatasetError() throws JsonProcessingException { InlineValue error = (InlineValue) results.get(InlineResponseKeyEnum.MESSAGE.getValue()); assertEquals("Error while getting dataset", error.message()); - // No job was submitted, so the user must not be told their data is being processed + // No job was submitted, so the user must not be told their data is being processed. + // The admission service owns that email now, so nothing here may send one either. verify(restServices, never()).notifyUser(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); } diff --git a/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestServicesTest.java b/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestServicesTest.java index 37c458be..345132a3 100644 --- a/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestServicesTest.java +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/RestServicesTest.java @@ -41,6 +41,19 @@ public void setUp() { closeableMock.close(); } + /** + * The two-step call the admission service makes: render the request at accept time, then + * submit it. Kept as one helper so these tests still read as one download request. + */ + private String downloadData(String uuid, String key, String startDate, String endDate, Object polygons, + String recipient, String collectionTitle, String fullMetadataLink, + String suggestedCitation, String outputFormat) throws JsonProcessingException { + DownloadRequest request = new DownloadRequest(uuid, key, startDate, endDate, polygons, recipient, + collectionTitle, fullMetadataLink, suggestedCitation, outputFormat); + return restServices.submitDownloadJob( + RestServices.downloadJobName(recipient), restServices.buildDownloadParameters(request)); + } + @Test public void testDownloadDataSuccess() throws JsonProcessingException { // Arrange @@ -50,7 +63,7 @@ public void testDownloadDataSuccess() throws JsonProcessingException { when(objectMapper.writeValueAsString(any())).thenReturn("test-multipolygon"); // Act - String response = restServices.downloadData( + String response = downloadData( "test-uuid", "test-dname", "2023-01-01", "2023-01-31", "test-multipolygon", "test@example.com", "Test Ocean Data Collection", "https://metadata.imas.utas.edu.au/.../test-uuid-123", "Cite data as: Mazor, T., Watermeyer, K., Hobley, T., Grinter, V., Holden, R., MacDonald, K. and Ferns, L. (2023).", "geotiff"); // Assert @@ -65,7 +78,7 @@ public void testDownloadDataJsonProcessingException() throws JsonProcessingExcep // Act & Assert try { - restServices.downloadData("test-uuid", "test-dname", "2023-01-01", "2023-01-31", "test-multipolygon", "test@example.com","Test Ocean Data Collection", "https://metadata.imas.utas.edu.au/.../test-uuid-123", "Cite data as: Mazor, T., Watermeyer, K., Hobley, T., Grinter, V., Holden, R., MacDonald, K. and Ferns, L. (2023).", "geotiff"); + downloadData("test-uuid", "test-dname", "2023-01-01", "2023-01-31", "test-multipolygon", "test@example.com","Test Ocean Data Collection", "https://metadata.imas.utas.edu.au/.../test-uuid-123", "Cite data as: Mazor, T., Watermeyer, K., Hobley, T., Grinter, V., Holden, R., MacDonald, K. and Ferns, L. (2023).", "geotiff"); } catch (JsonProcessingException e) { assertEquals("Error", e.getMessage()); } @@ -80,7 +93,7 @@ public void testDownloadDataCapturesSubmitJobRequest() throws JsonProcessingExce when(objectMapper.writeValueAsString(any())).thenReturn("test-multipolygon"); // Act - String response = restServices.downloadData( + String response = downloadData( "test-uuid", "test-dname","2023-01-01", "2023-01-31", "non-specified", "test@example.com", "Test Ocean Data Collection", "https://metadata.imas.utas.edu.au/.../test-uuid-123", "Cite data as: Mazor, T., Watermeyer, K., Hobley, T., Grinter, V., Holden, R., MacDonald, K. and Ferns, L. (2023).", "geotiff"); // Capture the submitted request @@ -108,7 +121,7 @@ public void submitJobReplacesEmptySuggestedCitationWithUnavailable() throws Json // polygons set to 'non-specified' to avoid objectMapper serialization // Act: pass empty suggestedCitation - String response = restServices.downloadData( + String response = downloadData( "test-uuid", "test-dname","2023-01-01", "2023-01-31", "non-specified", "test@example.com", "Test Ocean Data Collection", "https://metadata.imas.utas.edu.au/.../test-uuid-123", "", "geotiff"); @@ -126,10 +139,10 @@ public void submitJobReplacesEmptySuggestedCitationWithUnavailable() throws Json public void submitJobWithoutJobIdThrows() { // AWS Batch answered but gave us no job id, so the job was never really queued. // This must fail loudly: the caller sends the "processing started" email off the - // back of a successful downloadData(). + // back of a successful submit. when(batchClient.submitJob(any(SubmitJobRequest.class))).thenReturn(SubmitJobResponse.builder().build()); - assertThrows(IllegalStateException.class, () -> restServices.downloadData( + assertThrows(IllegalStateException.class, () -> downloadData( "test-uuid", "test-dname", "2023-01-01", "2023-01-31", "non-specified", "test@example.com", "Test Ocean Data Collection", "https://metadata.imas.utas.edu.au/.../test-uuid-123", "", "geotiff")); } diff --git a/server/src/test/resources/application-test.yaml b/server/src/test/resources/application-test.yaml index 80c43603..6b2db68d 100644 --- a/server/src/test/resources/application-test.yaml +++ b/server/src/test/resources/application-test.yaml @@ -21,3 +21,11 @@ elasticsearch: # no inference feature in test env so turn it off. semantic: enabled: false + +# Tests must never reach live AWS Batch. The release loop already makes no AWS call while +# nothing is held; this makes that explicit for any test that does build a Spring context. +aws: + batch: + job: + user-limit: + enabled: false