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 1/3] 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 From f71a8242c9da56a8d2149c12d344480f305f4a2a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Germ=C3=A1n=20Rodr=C3=ADguez?= <36035218+grojeda@users.noreply.github.com> Date: Tue, 1 Sep 2026 08:01:06 +1000 Subject: [PATCH 2/3] Fix admission-service concurrency issues flagged in PR review --- .../processes/DownloadAdmissionService.java | 156 +++++++++++++----- .../processes/InFlightDownloadCounter.java | 12 +- .../ogcapi/server/processes/RestServices.java | 11 +- 3 files changed, 134 insertions(+), 45 deletions(-) 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 index c313bc91..b913c097 100644 --- 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 @@ -10,10 +10,12 @@ 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.HashMap; +import java.util.HashSet; +import java.util.Iterator; import java.util.Map; +import java.util.Set; import java.util.UUID; import java.util.concurrent.ConcurrentHashMap; import java.util.concurrent.locks.ReentrantLock; @@ -47,9 +49,12 @@ public class DownloadAdmissionService { 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. + * Guards the in-memory bookkeeping only: the held queue, {@link #heldById} and + * {@link #reserved}. The actual AWS Batch submit call is a synchronous network call and is + * always made outside this lock, so a slow or throttled call from one admission decision + * or release never stalls every other user's admission decision and every status poll. + * + *

The AWS sweep is deliberately done before this is taken. */ private final ReentrantLock lock = new ReentrantLock(); @@ -59,6 +64,15 @@ public class DownloadAdmissionService { private final Map heldById = new ConcurrentHashMap<>(); private final Map released = new ConcurrentHashMap<>(); + /** + * Recipients with a submit currently in flight, not yet reflected in {@link #counter} + * because the AWS call it is waiting on has not returned. Consulted, together with the + * counter, whenever admission checks whether a recipient has room - without it, two + * decisions racing during the same in-flight submit could both see room and over-admit. + * Guarded by {@link #lock}. + */ + private final Map reserved = new HashMap<>(); + record Released(String awsJobId, Instant releasedAt) { } @@ -98,26 +112,28 @@ public DownloadAdmission submitOrHold(DownloadRequest request) throws JsonProces // Outside the lock: the sweep is the only expensive step and every user shares it. counter.refreshIfStale(); - DownloadAdmission admission; + String key = InFlightDownloadCounter.recipientKey(request.recipient()); 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); + // Never jump ahead of this user's own waiting requests. + boolean hasRoom = !hasHeldFor(request.recipient()) + && counter.countInFlight(request.recipient()) + reservedCount(key) < limits.maxConcurrent(); + if (!hasRoom) { + return hold(request, jobName, parameters); } + reserve(key); } finally { lock.unlock(); } - if (!admission.queued()) { - // Outside the lock: this is a synchronous SES call, and it is best effort anyway. + // Outside the lock: this is the synchronous AWS Batch call. + try { + DownloadAdmission admission = DownloadAdmission.submitted(submit(jobName, parameters, request.recipient())); notifyStarted(request); + return admission; + } finally { + releaseReservation(key); } - return admission; } /** @@ -167,6 +183,12 @@ public String awsJobIdOf(String jobId) { /** * 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. + * + *

Each release is claimed - removed from the held queue - under the lock, then submitted + * to AWS outside it, so this loop's AWS calls never block admission decisions or status + * polls for other downloads. When a claimed submit fails, its recipient is blocked for the + * rest of this sweep: without that, a later held download from the same recipient would be + * free to take the slot the failed one was retrying for, submitting out of order. */ @Scheduled(fixedDelayString = "${aws.batch.job.user-limit.release-interval:15s}") public void releaseHeldDownloads() { @@ -177,44 +199,77 @@ public void releaseHeldDownloads() { counter.refresh(); - List toNotify = new ArrayList<>(); lock.lock(); try { expireStale(); + } finally { + lock.unlock(); + } + + Set blockedRecipients = new HashSet<>(); + HeldDownload job; + while ((job = claimNextReleasable(blockedRecipients)) != null) { + String recipient = job.request().recipient(); + String key = InFlightDownloadCounter.recipientKey(recipient); + try { + String awsJobId = submit(job.jobName(), job.parameters(), recipient); + released.put(job.jobId(), new Released(awsJobId, clock.instant())); + log.info("Released held download {} as AWS Batch job {}", job.jobId(), awsJobId); + notifyStarted(job.request()); + } catch (Exception e) { + blockedRecipients.add(key); + HeldDownload retried = job.withAttempt(); + if (retried.attempts() >= MAX_SUBMIT_ATTEMPTS) { + 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); + requeue(retried); + } + } finally { + releaseReservation(key); + } + } + } - Deque blocked = new ArrayDeque<>(); - HeldDownload job; - while ((job = held.pollFirst()) != null) { + /** + * Remove and return the earliest held download whose recipient is not blocked this sweep + * and currently has room, reserving its slot; null when nothing in the queue qualifies. + */ + private HeldDownload claimNextReleasable(Set blockedRecipients) { + lock.lock(); + try { + Iterator it = held.iterator(); + while (it.hasNext()) { + HeldDownload job = it.next(); String recipient = job.request().recipient(); - if (counter.countInFlight(recipient) >= limits.maxConcurrent()) { - blocked.addLast(job); + String key = InFlightDownloadCounter.recipientKey(recipient); + if (blockedRecipients.contains(key)) { 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); - } + if (counter.countInFlight(recipient) + reservedCount(key) >= limits.maxConcurrent()) { + continue; } + it.remove(); + heldById.remove(job.jobId()); + reserve(key); + return job; } - held.addAll(blocked); + return null; } finally { lock.unlock(); } + } - toNotify.forEach(this::notifyStarted); + /** Put a download that failed to submit back at the end of the held queue for another try. */ + private void requeue(HeldDownload job) { + lock.lock(); + try { + held.addLast(job); + heldById.put(job.jobId(), job); + } finally { + lock.unlock(); + } } private String submit(String jobName, Map parameters, String recipient) { @@ -229,6 +284,25 @@ private boolean hasHeldFor(String recipient) { .anyMatch(job -> InFlightDownloadCounter.recipientKey(job.request().recipient()).equals(key)); } + /** Must be called with {@link #lock} held. */ + private void reserve(String key) { + reserved.merge(key, 1, Integer::sum); + } + + /** Must be called with {@link #lock} held. */ + private int reservedCount(String key) { + return reserved.getOrDefault(key, 0); + } + + private void releaseReservation(String key) { + lock.lock(); + try { + reserved.computeIfPresent(key, (k, count) -> count <= 1 ? null : count - 1); + } finally { + lock.unlock(); + } + } + 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 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 index f1eedde7..66659841 100644 --- 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 @@ -106,13 +106,21 @@ static Snapshot empty() { /** * 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. + * + *

Synchronized so that when several callers arrive after the snapshot has expired, only + * the first actually sweeps; the rest block briefly on this monitor and then see the fresh + * snapshot the first caller just took, rather than each repeating the sweep themselves. */ - public void refreshIfStale() { - if (Duration.between(snapshot.takenAt(), clock.instant()).compareTo(limits.releaseInterval()) >= 0) { + public synchronized void refreshIfStale() { + if (isStale()) { refresh(); } } + private boolean isStale() { + return Duration.between(snapshot.takenAt(), clock.instant()).compareTo(limits.releaseInterval()) >= 0; + } + /** Sweep the queues now, whatever the age of the current snapshot. */ public synchronized void refresh() { try { 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 3a3cfb0a..0e984245 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 @@ -24,6 +24,7 @@ import java.io.IOException; import java.math.BigInteger; import java.nio.charset.StandardCharsets; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.concurrent.atomic.AtomicBoolean; @@ -134,12 +135,18 @@ private String submitJob(String jobName, String jobQueue, String jobDefinition, // Filter out null or empty parameter values before submitting to AWS Batch. // AWS Batch returns "Parameter values must be provided" when the job definition // declares parameters but some submitted values are null/empty. + // + // A defensive copy, never the caller's own map: a held download keeps this exact map + // instance around so the status endpoint can describe it while it waits, and readers + // of that map run on other threads with no synchronization of their own. + Map submitParameters = parameters; if (parameters != null) { var suggestedCitation = parameters.get(DatasetDownloadEnums.Parameter.SUGGESTED_CITATION.getValue()); // empty suggested citation is acceptable since it may be from external orgs if (suggestedCitation == null || suggestedCitation.isEmpty()) { log.warn("Suggested citation is null or empty for job '{}'. Submitting with unavailable as value.", jobName); - parameters.replace(DatasetDownloadEnums.Parameter.SUGGESTED_CITATION.getValue(), "unavailable"); + submitParameters = new HashMap<>(parameters); + submitParameters.put(DatasetDownloadEnums.Parameter.SUGGESTED_CITATION.getValue(), "unavailable"); } } @@ -147,7 +154,7 @@ private String submitJob(String jobName, String jobQueue, String jobDefinition, .jobName(jobName) .jobQueue(jobQueue) .jobDefinition(jobDefinition) - .parameters(parameters) + .parameters(submitParameters) .build(); SubmitJobResponse submitJobResponse = batchClient.submitJob(submitJobRequest); From 490235aa0bb7aeef8033baff78b681741e8627c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Germ=C3=A1n=20Rodr=C3=ADguez?= <36035218+grojeda@users.noreply.github.com> Date: Thu, 3 Sep 2026 09:00:59 +1000 Subject: [PATCH 3/3] reject downloads at the per-user limit instead of queuing them --- .../DownloadLimitExceededException.java | 8 + .../exception/GlobalExceptionHandler.java | 14 + .../core/model/DownloadExecutionResponse.java | 12 +- .../core/model/DownloadJobStatusInfo.java | 29 -- .../server/processes/DownloadAdmission.java | 21 -- .../processes/DownloadAdmissionService.java | 298 ++--------------- .../processes/DownloadJobStatusService.java | 58 +--- .../processes/DownloadLimitProperties.java | 14 +- .../ogcapi/server/processes/HeldDownload.java | 25 -- .../processes/InFlightDownloadCounter.java | 2 +- .../aodn/ogcapi/server/processes/RestApi.java | 44 ++- server/src/main/resources/application.yaml | 14 +- .../DownloadAdmissionServiceTest.java | 301 +++--------------- .../DownloadJobStatusServiceTest.java | 62 ---- .../InFlightDownloadCounterTest.java | 2 +- .../server/processes/RestApiJobsTest.java | 34 +- .../ogcapi/server/processes/RestApiTest.java | 26 +- 17 files changed, 172 insertions(+), 792 deletions(-) create mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/core/exception/DownloadLimitExceededException.java delete mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmission.java delete mode 100644 server/src/main/java/au/org/aodn/ogcapi/server/processes/HeldDownload.java diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/core/exception/DownloadLimitExceededException.java b/server/src/main/java/au/org/aodn/ogcapi/server/core/exception/DownloadLimitExceededException.java new file mode 100644 index 00000000..0e1dad8f --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/core/exception/DownloadLimitExceededException.java @@ -0,0 +1,8 @@ +package au.org.aodn.ogcapi.server.core.exception; + +public class DownloadLimitExceededException extends RuntimeException { + public DownloadLimitExceededException(int maxConcurrent) { + super("You already have " + maxConcurrent + " downloads in progress. " + + "Wait for one of them to complete before starting another."); + } +} diff --git a/server/src/main/java/au/org/aodn/ogcapi/server/core/exception/GlobalExceptionHandler.java b/server/src/main/java/au/org/aodn/ogcapi/server/core/exception/GlobalExceptionHandler.java index 9aca8dca..474a183d 100644 --- a/server/src/main/java/au/org/aodn/ogcapi/server/core/exception/GlobalExceptionHandler.java +++ b/server/src/main/java/au/org/aodn/ogcapi/server/core/exception/GlobalExceptionHandler.java @@ -116,6 +116,20 @@ public ResponseEntity handleDownloadJobStatusException( return new ResponseEntity<>(errorResponse, HttpStatus.INTERNAL_SERVER_ERROR); } + @ExceptionHandler(DownloadLimitExceededException.class) + public ResponseEntity handleDownloadLimitExceededException( + DownloadLimitExceededException ex, + WebRequest request) { + ErrorResponse errorResponse = ErrorResponse + .builder() + .timestamp(LocalDateTime.now()) + .message(ex.getMessage()) + .details(request.getDescription(false)) + .build(); + + return new ResponseEntity<>(errorResponse, HttpStatus.TOO_MANY_REQUESTS); + } + @ExceptionHandler(DasUpstreamException.class) public ResponseEntity handleDasUpstreamException(DasUpstreamException ex, WebRequest request) { ErrorResponse errorResponse = ErrorResponse 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 8abf679c..81ca4cd2 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,18 +4,10 @@ import com.fasterxml.jackson.annotation.JsonProperty; import io.swagger.v3.oas.annotations.media.Schema; -@Schema(description = "Compatible download execution response with the submitted job ID, " - + "and whether the download is waiting for a free per-user slot.") +@Schema(description = "Compatible download execution response with the submitted job ID.") public record DownloadExecutionResponse( @JsonProperty("message") InlineValue message, @JsonProperty("status") InlineValue status, - @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 + @JsonProperty("jobID") String jobId ) 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 dedb970b..e05ea248 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,19 +31,6 @@ 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; } @@ -75,20 +62,4 @@ 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 deleted file mode 100644 index 8cf6ea83..00000000 --- a/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmission.java +++ /dev/null @@ -1,21 +0,0 @@ -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 index b913c097..a8768d75 100644 --- 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 @@ -1,299 +1,91 @@ package au.org.aodn.ogcapi.server.processes; +import au.org.aodn.ogcapi.server.core.exception.DownloadLimitExceededException; 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.Deque; import java.util.HashMap; -import java.util.HashSet; -import java.util.Iterator; import java.util.Map; -import java.util.Set; -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. + * Caps how many downloads one user can have running at once. A user already at the limit is + * rejected outright: the caller has to wait for one of their own downloads to finish and + * try again. */ @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 in-memory bookkeeping only: the held queue, {@link #heldById} and - * {@link #reserved}. The actual AWS Batch submit call is a synchronous network call and is - * always made outside this lock, so a slow or throttled call from one admission decision - * or release never stalls every other user's admission decision and every status poll. - * - *

The AWS sweep is deliberately done before this is taken. - */ 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<>(); - /** * Recipients with a submit currently in flight, not yet reflected in {@link #counter} - * because the AWS call it is waiting on has not returned. Consulted, together with the - * counter, whenever admission checks whether a recipient has room - without it, two - * decisions racing during the same in-flight submit could both see room and over-admit. - * Guarded by {@link #lock}. + * because the AWS call it is waiting on has not returned. Consulted together with the + * counter whenever admission checks whether a recipient has room - without it, two + * requests racing during the same in-flight submit could both see room and both go + * through, letting a user briefly exceed the limit. Guarded by {@link #lock}. */ private final Map reserved = new HashMap<>(); - 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. + * Submit a download to AWS Batch and return its job id. + * + * @throws DownloadLimitExceededException the recipient already has {@code maxConcurrent} + * downloads running */ - 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(); - + public String submit(DownloadRequest request) throws JsonProcessingException { String key = InFlightDownloadCounter.recipientKey(request.recipient()); - lock.lock(); - try { - // Never jump ahead of this user's own waiting requests. - boolean hasRoom = !hasHeldFor(request.recipient()) - && counter.countInFlight(request.recipient()) + reservedCount(key) < limits.maxConcurrent(); - if (!hasRoom) { - return hold(request, jobName, parameters); - } - reserve(key); - } finally { - lock.unlock(); - } - // Outside the lock: this is the synchronous AWS Batch call. - try { - DownloadAdmission admission = DownloadAdmission.submitted(submit(jobName, parameters, request.recipient())); - notifyStarted(request); - return admission; - } finally { - releaseReservation(key); + if (limits.enabled()) { + // Outside the lock: the sweep is the only expensive step. + counter.refreshIfStale(); + reserveOrReject(request, key); } - } - /** - * 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. - * - *

Each release is claimed - removed from the held queue - under the lock, then submitted - * to AWS outside it, so this loop's AWS calls never block admission decisions or status - * polls for other downloads. When a claimed submit fails, its recipient is blocked for the - * rest of this sweep: without that, a later held download from the same recipient would be - * free to take the slot the failed one was retrying for, submitting out of order. - */ - @Scheduled(fixedDelayString = "${aws.batch.job.user-limit.release-interval:15s}") - public void releaseHeldDownloads() { - pruneReleased(); - if (!limits.enabled() || heldById.isEmpty()) { - return; - } - - counter.refresh(); - - lock.lock(); - try { - expireStale(); + Map parameters = restServices.buildDownloadParameters(request); + String jobName = RestServices.downloadJobName(request.recipient()); + String awsJobId = restServices.submitDownloadJob(jobName, parameters); + counter.recordSubmitted(awsJobId, request.recipient()); + notifyStarted(request); + return awsJobId; } finally { - lock.unlock(); - } - - Set blockedRecipients = new HashSet<>(); - HeldDownload job; - while ((job = claimNextReleasable(blockedRecipients)) != null) { - String recipient = job.request().recipient(); - String key = InFlightDownloadCounter.recipientKey(recipient); - try { - String awsJobId = submit(job.jobName(), job.parameters(), recipient); - released.put(job.jobId(), new Released(awsJobId, clock.instant())); - log.info("Released held download {} as AWS Batch job {}", job.jobId(), awsJobId); - notifyStarted(job.request()); - } catch (Exception e) { - blockedRecipients.add(key); - HeldDownload retried = job.withAttempt(); - if (retried.attempts() >= MAX_SUBMIT_ATTEMPTS) { - 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); - requeue(retried); - } - } finally { + if (limits.enabled()) { releaseReservation(key); } } } - /** - * Remove and return the earliest held download whose recipient is not blocked this sweep - * and currently has room, reserving its slot; null when nothing in the queue qualifies. - */ - private HeldDownload claimNextReleasable(Set blockedRecipients) { + private void reserveOrReject(DownloadRequest request, String key) { lock.lock(); try { - Iterator it = held.iterator(); - while (it.hasNext()) { - HeldDownload job = it.next(); - String recipient = job.request().recipient(); - String key = InFlightDownloadCounter.recipientKey(recipient); - if (blockedRecipients.contains(key)) { - continue; - } - if (counter.countInFlight(recipient) + reservedCount(key) >= limits.maxConcurrent()) { - continue; - } - it.remove(); - heldById.remove(job.jobId()); - reserve(key); - return job; + int inFlight = counter.countInFlight(request.recipient()) + reserved.getOrDefault(key, 0); + if (inFlight >= limits.maxConcurrent()) { + throw new DownloadLimitExceededException(limits.maxConcurrent()); } - return null; - } finally { - lock.unlock(); - } - } - - /** Put a download that failed to submit back at the end of the held queue for another try. */ - private void requeue(HeldDownload job) { - lock.lock(); - try { - held.addLast(job); - heldById.put(job.jobId(), job); + reserved.merge(key, 1, Integer::sum); } finally { lock.unlock(); } } - 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)); - } - - /** Must be called with {@link #lock} held. */ - private void reserve(String key) { - reserved.merge(key, 1, Integer::sum); - } - - /** Must be called with {@link #lock} held. */ - private int reservedCount(String key) { - return reserved.getOrDefault(key, 0); - } - private void releaseReservation(String key) { lock.lock(); try { @@ -303,38 +95,6 @@ private void releaseReservation(String 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(), 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 0b94f9d8..d583de25 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 @@ -52,61 +52,44 @@ 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, - DownloadAdmissionService admissionService) { - this(batchClient, properties, aggregator, admissionService, Clock.systemUTC()); + DownloadJobStatusAggregator aggregator) { + this(batchClient, properties, aggregator, 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(awsJobId); + JobDetail initial = describeInitialJob(jobId); if (initial.status() == JobStatus.FAILED) { StatusCode status = aggregator.aggregate(new DownloadJobStatusAggregator.Snapshot( initial.status(), null, null, DownloadJobStatusAggregator.WorkflowMode.CHILD_DISCOVERY_REQUIRED, false)); - return toStatusInfo(publicJobId, status, initial, null, null); + return toStatusInfo(jobId, status, initial, null, null); } JobDetail prepare = findChildJob( - PREPARE_NAME_PREFIX + awsJobId, awsJobId, PREPARE_TYPE); + PREPARE_NAME_PREFIX + jobId, jobId, PREPARE_TYPE); JobDetail collect = findChildJob( - COLLECT_NAME_PREFIX + awsJobId, awsJobId, COLLECT_TYPE); + COLLECT_NAME_PREFIX + jobId, jobId, COLLECT_TYPE); boolean discoveryWindowExpired = discoveryWindowExpired(initial); DownloadJobStatusAggregator.WorkflowMode workflowMode = isExplicitZarr(initial.parameters()) @@ -120,40 +103,15 @@ private DownloadJobStatusInfo awsStatus(String awsJobId, String publicJobId) { workflowMode, discoveryWindowExpired)); - return toStatusInfo(publicJobId, status, initial, prepare, collect); + return toStatusInfo(jobId, status, initial, prepare, collect); } catch (DownloadJobNotFoundException e) { throw e; } catch (Exception e) { - log.error("Failed to reconstruct AWS Batch workflow for download job {}", publicJobId, e); + log.error("Failed to reconstruct AWS Batch workflow for download job {}", jobId, 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 index 4695c781..5d1e2e15 100644 --- 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 @@ -6,19 +6,15 @@ 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. + * Per-user concurrency limit for dataset downloads. A user, identified by the recipient + * email, may have at most {@code maxConcurrent} downloads in flight at once; a request past + * that is rejected outright. */ @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 + /** How long the shared in-flight snapshot is reused before the queues are swept again. */ + @DefaultValue("15s") Duration refreshInterval ) { } 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 deleted file mode 100644 index 0fbcced3..00000000 --- a/server/src/main/java/au/org/aodn/ogcapi/server/processes/HeldDownload.java +++ /dev/null @@ -1,25 +0,0 @@ -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 index 66659841..2d0e9a2f 100644 --- 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 @@ -118,7 +118,7 @@ public synchronized void refreshIfStale() { } private boolean isStale() { - return Duration.between(snapshot.takenAt(), clock.instant()).compareTo(limits.releaseInterval()) >= 0; + return Duration.between(snapshot.takenAt(), clock.instant()).compareTo(limits.refreshInterval()) >= 0; } /** Sweep the queues now, whatever the age of the current snapshot. */ 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 173bbf04..fcb5e754 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 @@ -8,8 +8,10 @@ import au.org.aodn.ogcapi.processes.model.Results; import au.org.aodn.ogcapi.processes.model.JobList; import au.org.aodn.ogcapi.processes.model.StatusInfo; +import au.org.aodn.ogcapi.server.core.exception.DownloadLimitExceededException; import au.org.aodn.ogcapi.server.core.model.DownloadExecutionResponse; import au.org.aodn.ogcapi.server.core.model.DownloadJobStatusInfo; +import au.org.aodn.ogcapi.server.core.model.ErrorResponse; import au.org.aodn.ogcapi.server.core.model.InlineValue; import au.org.aodn.ogcapi.server.core.model.enumeration.DatasetDownloadEnums; import au.org.aodn.ogcapi.server.core.model.enumeration.InlineResponseKeyEnum; @@ -65,8 +67,18 @@ 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", - "queued": false + "jobID": "123e4567-e89b-12d3-a456-426614174000" + } + """))) + @ApiResponse( + responseCode = "429", + description = "The recipient already has the maximum number of downloads in progress.", + content = @Content( + mediaType = MediaType.APPLICATION_JSON_VALUE, + schema = @Schema(implementation = ErrorResponse.class), + examples = @ExampleObject(value = """ + { + "message": "You already have 10 downloads in progress. Wait for one of them to complete before starting another." } """))) public ResponseEntity execute( @@ -96,27 +108,25 @@ public ResponseEntity execute( DownloadRequest request = new DownloadRequest(uuid, key, startDate, endDate, multiPolygon, recipient, 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 per-user limit is applied here: a recipient already at their limit is + // rejected outright, before anything is submitted to AWS Batch. // // 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); - - // 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()); + // make the first email faster. It is 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. + String awsJobId = downloadAdmissionService.submit(request); + + var value = new InlineValue("Job submitted with ID: " + awsJobId); var status = new InlineValue(Integer.toString(HttpStatus.OK.value())); - var results = new DownloadExecutionResponse( - value, status, admission.jobId(), admission.queued(), admission.queuePosition()); + var results = new DownloadExecutionResponse(value, status, awsJobId); return ResponseEntity.ok(results); + } catch (DownloadLimitExceededException e) { + // Let GlobalExceptionHandler turn this into a real 429 with a message the + // caller can act on, instead of the generic 200-wrapped error below. + throw e; } catch (Exception e) { // TODO: currently all the errors return badRequest. This should be changed to return the correct status code diff --git a/server/src/main/resources/application.yaml b/server/src/main/resources/application.yaml index cc812b9b..3c3e0ffc 100644 --- a/server/src/main/resources/application.yaml +++ b/server/src/main/resources/application.yaml @@ -59,19 +59,13 @@ 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. + # Per-user download concurrency. A recipient already running max-concurrent downloads + # is rejected outright until one of their own downloads finishes. 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 + # How long the shared in-flight snapshot is reused before the queues are swept again. + refresh-interval: 15s 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 index 381d2015..aecad7b2 100644 --- 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 @@ -1,43 +1,32 @@ package au.org.aodn.ogcapi.server.processes; +import au.org.aodn.ogcapi.server.core.exception.DownloadLimitExceededException; 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"; @@ -50,39 +39,27 @@ class DownloadAdmissionServiceTest { /** 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(restServices.submitDownloadJob(anyString(), any())) + .thenAnswer(invocation -> 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); + return new DownloadAdmissionService(restServices, counter, limits); } private static DownloadLimitProperties limits(boolean enabled, int maxConcurrent) { - return new DownloadLimitProperties( - enabled, maxConcurrent, Duration.ofSeconds(15), Duration.ofHours(24), 1000); + return new DownloadLimitProperties(enabled, maxConcurrent, Duration.ofSeconds(15)); } private AtomicInteger current(String recipient) { @@ -99,254 +76,67 @@ private DownloadRequest request(String recipient) { } @Test - void underTheLimitSubmitsStraightAwayAndReturnsTheAwsJobId() throws Exception { - String jobId = service.submitOrHold(request(RECIPIENT)).jobId(); + void underTheLimitSubmitsAndReturnsTheAwsJobId() throws Exception { + String jobId = service.submit(request(RECIPIENT)); 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 { + void atTheLimitIsRejectedOutright() throws Exception { current(RECIPIENT).set(10); - String jobId = service.submitOrHold(request(RECIPIENT)).jobId(); + assertThrows(DownloadLimitExceededException.class, () -> service.submit(request(RECIPIENT))); - 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. + // Nothing was submitted, 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 { + void theRejectionMessageNamesTheConfiguredLimit() { 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(); + DownloadLimitExceededException exception = assertThrows( + DownloadLimitExceededException.class, () -> service.submit(request(RECIPIENT))); - assertNotNull(service.findHeld(first)); - assertNotNull(service.findHeld(second)); - verify(restServices, never()).submitDownloadJob(anyString(), any()); + assertEquals("You already have 10 downloads in progress. " + + "Wait for one of them to complete before starting another.", exception.getMessage()); } @Test void oneUserAtTheLimitDoesNotBlockAnother() throws Exception { current(RECIPIENT).set(10); + current(OTHER_RECIPIENT).set(0); - 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(); + assertThrows(DownloadLimitExceededException.class, () -> service.submit(request(RECIPIENT))); + String jobId = service.submit(request(OTHER_RECIPIENT)); - 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()); + assertNotNull(jobId); } @Test - void aFailedReleaseKeepsTheDownloadHeldForAnotherAttempt() throws Exception { + void aDifferentlyCasedAddressIsCountedAgainstTheSameLimit() 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)); + assertThrows(DownloadLimitExceededException.class, + () -> service.submit(request("Person@Example.COM"))); } @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(); + void theSweepHappensBeforeTheAdmissionDecision() throws Exception { + service.submit(request(RECIPIENT)); - assertNull(service.findHeld(heldId)); + verify(counter).refreshIfStale(); } @Test - void theHoldQueueIsBounded() throws Exception { - DownloadAdmissionService bounded = build(new DownloadLimitProperties( - true, 1, Duration.ofSeconds(15), Duration.ofHours(24), 2)); - current(RECIPIENT).set(1); + void aSuccessfulSubmitIsRecordedAgainstTheRecipientForTheGraceWindow() throws Exception { + String jobId = service.submit(request(RECIPIENT)); - bounded.submitOrHold(request(RECIPIENT)); - bounded.submitOrHold(request(RECIPIENT)); - - assertThrows(IllegalStateException.class, () -> bounded.submitOrHold(request(RECIPIENT))); + verify(counter).recordSubmitted(jobId, RECIPIENT); } @Test @@ -354,32 +144,31 @@ void theLimitCanBeTurnedOffEntirely() throws Exception { DownloadAdmissionService disabled = build(limits(false, 10)); current(RECIPIENT).set(500); - String jobId = disabled.submitOrHold(request(RECIPIENT)).jobId(); + String jobId = disabled.submit(request(RECIPIENT)); - assertNull(disabled.findHeld(jobId)); + assertNotNull(jobId); verify(restServices).submitDownloadJob(anyString(), any()); verify(counter, never()).refreshIfStale(); } @Test - void theSweepHappensBeforeTheAdmissionDecision() throws Exception { - service.submitOrHold(request(RECIPIENT)); - - verify(counter).refreshIfStale(); - } + void aFailedSubmitDoesNotLeakAReservedSlot() throws Exception { + // One slot free; the counter will not move because the submit below never actually + // reaches AWS. + current(RECIPIENT).set(9); + org.mockito.Mockito.when(restServices.submitDownloadJob(anyString(), any())) + .thenThrow(new IllegalStateException("AWS Batch rejected the job")); - @Test - void aReleasedDownloadIsSubmittedWithTheSameNameAndRecipientItWasAcceptedWith() throws Exception { - current(RECIPIENT).set(10); - service.submitOrHold(request(RECIPIENT)); + assertThrows(IllegalStateException.class, () -> service.submit(request(RECIPIENT))); - current(RECIPIENT).set(0); - service.releaseHeldDownloads(); + // If the reservation taken for the failed attempt were never released, this retry + // would see 9 (counter) + 1 (leaked reservation) = 10 and be rejected even though the + // failed attempt never actually started a download. + org.mockito.Mockito.reset(restServices); + lenient().when(restServices.buildDownloadParameters(any())).thenReturn(new HashMap<>()); + lenient().when(restServices.submitDownloadJob(anyString(), any())) + .thenReturn(UUID.randomUUID().toString()); - 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()); + assertNotNull(service.submit(request(RECIPIENT))); } } 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 f325a4e3..d45d6811 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,7 +31,6 @@ 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; @@ -39,7 +38,6 @@ import static org.mockito.Mockito.lenient; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; @ExtendWith(MockitoExtension.class) class DownloadJobStatusServiceTest { @@ -56,9 +54,6 @@ 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; @@ -84,7 +79,6 @@ void setUp() { batchClient, new BatchJobProperties(QUEUE_NAME, DEFINITION_NAME, CHILD_QUEUE_NAME), new DownloadJobStatusAggregator(), - admissionService, Clock.fixed(NOW, ZoneOffset.UTC)); } @@ -142,7 +136,6 @@ 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()); @@ -286,61 +279,6 @@ 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 index 8d33f785..966dc2da 100644 --- 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 @@ -90,7 +90,7 @@ 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), + new DownloadLimitProperties(true, 10, Duration.ofSeconds(15)), clock); } 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 50639c6a..5a34d9e5 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 @@ -4,6 +4,7 @@ import au.org.aodn.ogcapi.processes.model.StatusInfo; import au.org.aodn.ogcapi.server.core.exception.DownloadJobNotFoundException; import au.org.aodn.ogcapi.server.core.exception.DownloadJobStatusException; +import au.org.aodn.ogcapi.server.core.exception.DownloadLimitExceededException; import au.org.aodn.ogcapi.server.core.exception.GlobalExceptionHandler; import au.org.aodn.ogcapi.server.core.model.DownloadJobStatusInfo; import com.fasterxml.jackson.databind.ObjectMapper; @@ -57,8 +58,7 @@ void setUp() { @Test void postKeepsExistingFieldsAndAddsPureJobId() throws Exception { - when(downloadAdmissionService.submitOrHold(any())) - .thenReturn(new DownloadAdmission(JOB_ID, false, null)); + when(downloadAdmissionService.submit(any())).thenReturn(JOB_ID); String body = objectMapper.writeValueAsString(Map.of("inputs", Map.of( "uuid", "collection-id", "recipient", "person@example.com"))); @@ -74,9 +74,9 @@ void postKeepsExistingFieldsAndAddsPureJobId() throws Exception { } @Test - void postSurfacesTheQueuedStateAndPosition() throws Exception { - when(downloadAdmissionService.submitOrHold(any())) - .thenReturn(new DownloadAdmission(JOB_ID, true, 3)); + void postAtTheDownloadLimitReturnsTooManyRequestsWithAClearMessage() throws Exception { + when(downloadAdmissionService.submit(any())) + .thenThrow(new DownloadLimitExceededException(10)); String body = objectMapper.writeValueAsString(Map.of("inputs", Map.of( "uuid", "collection-id", "recipient", "person@example.com"))); @@ -85,27 +85,9 @@ void postSurfacesTheQueuedStateAndPosition() throws Exception { .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()); + .andExpect(status().is(429)) + .andExpect(jsonPath("$.message").value("You already have 10 downloads in progress. " + + "Wait for one of them to complete before starting another.")); } @Test 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 e9abcfc6..12fe61e5 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 @@ -3,6 +3,7 @@ import au.org.aodn.ogcapi.processes.model.Execute; import au.org.aodn.ogcapi.processes.model.InlineResponse200; import au.org.aodn.ogcapi.processes.model.Results; +import au.org.aodn.ogcapi.server.core.exception.DownloadLimitExceededException; import au.org.aodn.ogcapi.server.core.model.DownloadExecutionResponse; import au.org.aodn.ogcapi.server.core.model.InlineValue; import au.org.aodn.ogcapi.server.core.model.enumeration.DatasetDownloadEnums; @@ -23,6 +24,7 @@ import static org.junit.jupiter.api.Assertions.assertEquals; import static org.junit.jupiter.api.Assertions.assertInstanceOf; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.mockito.ArgumentMatchers.any; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; @@ -57,8 +59,8 @@ public void setUp() { @Test public void testExecuteDownloadDatasetSuccess() throws JsonProcessingException { - when(downloadAdmissionService.submitOrHold(any())) - .thenReturn(new DownloadAdmission("test-job-id", false, null)); + when(downloadAdmissionService.submit(any())) + .thenReturn("test-job-id"); ResponseEntity response = restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); @@ -73,13 +75,13 @@ public void testExecuteDownloadDatasetSuccess() throws JsonProcessingException { @Test public void testExecutePassesEveryRequestInputToAdmission() throws JsonProcessingException { - when(downloadAdmissionService.submitOrHold(any())) - .thenReturn(new DownloadAdmission("test-job-id", false, null)); + when(downloadAdmissionService.submit(any())) + .thenReturn("test-job-id"); restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); ArgumentCaptor captor = ArgumentCaptor.forClass(DownloadRequest.class); - verify(downloadAdmissionService).submitOrHold(captor.capture()); + verify(downloadAdmissionService).submit(captor.capture()); DownloadRequest request = captor.getValue(); assertEquals("test-uuid", request.uuid()); assertEquals("2023-01-01", request.startDate()); @@ -90,7 +92,7 @@ public void testExecutePassesEveryRequestInputToAdmission() throws JsonProcessin @Test public void testExecuteDownloadDatasetError() throws JsonProcessingException { - when(downloadAdmissionService.submitOrHold(any())) + when(downloadAdmissionService.submit(any())) .thenThrow(new RuntimeException("Error while getting dataset")); ResponseEntity response = restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); @@ -106,6 +108,18 @@ public void testExecuteDownloadDatasetError() throws JsonProcessingException { verify(restServices, never()).notifyUser(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); } + @Test + public void testExecuteAtTheDownloadLimitLetsTheExceptionPropagate() throws JsonProcessingException { + DownloadLimitExceededException limitExceeded = new DownloadLimitExceededException(10); + when(downloadAdmissionService.submit(any())).thenThrow(limitExceeded); + + // Unlike a generic failure, this must reach GlobalExceptionHandler as a real 429 - + // not be swallowed into the 200-wrapped "Error while getting dataset" response. + DownloadLimitExceededException thrown = assertThrows(DownloadLimitExceededException.class, + () -> restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest)); + assertEquals(limitExceeded.getMessage(), thrown.getMessage()); + } + @Test public void testExecuteUnknownProcessId() { ResponseEntity response = restApi.execute("unknown-process-id", executeRequest);