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/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 d2c65996..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,7 +4,7 @@ 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.") public record DownloadExecutionResponse( @JsonProperty("message") InlineValue message, @JsonProperty("status") InlineValue status, 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..a8768d75 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionService.java @@ -0,0 +1,111 @@ +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.stereotype.Service; + +import java.util.HashMap; +import java.util.Map; +import java.util.concurrent.locks.ReentrantLock; + +/** + * 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 { + + private final RestServices restServices; + private final InFlightDownloadCounter counter; + private final DownloadLimitProperties limits; + + private final ReentrantLock lock = new ReentrantLock(); + + /** + * 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 + * 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<>(); + + @Autowired + public DownloadAdmissionService( + RestServices restServices, + InFlightDownloadCounter counter, + DownloadLimitProperties limits) { + this.restServices = restServices; + this.counter = counter; + this.limits = limits; + } + + /** + * Submit a download to AWS Batch and return its job id. + * + * @throws DownloadLimitExceededException the recipient already has {@code maxConcurrent} + * downloads running + */ + public String submit(DownloadRequest request) throws JsonProcessingException { + String key = InFlightDownloadCounter.recipientKey(request.recipient()); + + if (limits.enabled()) { + // Outside the lock: the sweep is the only expensive step. + counter.refreshIfStale(); + reserveOrReject(request, key); + } + + try { + 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 { + if (limits.enabled()) { + releaseReservation(key); + } + } + } + + private void reserveOrReject(DownloadRequest request, String key) { + lock.lock(); + try { + int inFlight = counter.countInFlight(request.recipient()) + reserved.getOrDefault(key, 0); + if (inFlight >= limits.maxConcurrent()) { + throw new DownloadLimitExceededException(limits.maxConcurrent()); + } + reserved.merge(key, 1, Integer::sum); + } finally { + lock.unlock(); + } + } + + private void releaseReservation(String key) { + lock.lock(); + try { + reserved.computeIfPresent(key, (k, count) -> count <= 1 ? null : count - 1); + } finally { + lock.unlock(); + } + } + + 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..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 @@ -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"; @@ -71,7 +75,6 @@ public DownloadJobStatusService( public DownloadJobStatusInfo getStatus(String jobId) { validateJobId(jobId); - try { JobDetail initial = describeInitialJob(jobId); @@ -83,12 +86,10 @@ public DownloadJobStatusInfo getStatus(String jobId) { return toStatusInfo(jobId, 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 + jobId, jobId, PREPARE_TYPE); JobDetail collect = findChildJob( - "collect-data-for-job-" + jobId, jobId, COLLECT_TYPE); + COLLECT_NAME_PREFIX + jobId, jobId, COLLECT_TYPE); boolean discoveryWindowExpired = discoveryWindowExpired(initial); DownloadJobStatusAggregator.WorkflowMode workflowMode = isExplicitZarr(initial.parameters()) 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..5d1e2e15 --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadLimitProperties.java @@ -0,0 +1,20 @@ +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 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, + /** 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/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/InFlightDownloadCounter.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounter.java new file mode 100644 index 00000000..2d0e9a2f --- /dev/null +++ b/server/src/main/java/au/org/aodn/ogcapi/server/processes/InFlightDownloadCounter.java @@ -0,0 +1,276 @@ +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. + * + *

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 synchronized void refreshIfStale() { + if (isStale()) { + refresh(); + } + } + + private boolean isStale() { + return Duration.between(snapshot.takenAt(), clock.instant()).compareTo(limits.refreshInterval()) >= 0; + } + + /** 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..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; @@ -43,6 +45,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 @@ -65,6 +70,17 @@ public class RestApi implements ProcessesApi, JobsApi { "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( @Parameter(in = ParameterIn.PATH, required = true, schema = @Schema()) @PathVariable("processID") @@ -89,21 +105,28 @@ 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 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 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: " + jobId); + var value = new InlineValue("Job submitted with ID: " + awsJobId); var status = new InlineValue(Integer.toString(HttpStatus.OK.value())); - var results = new DownloadExecutionResponse(value, status, jobId); + 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/java/au/org/aodn/ogcapi/server/processes/RestServices.java b/server/src/main/java/au/org/aodn/ogcapi/server/processes/RestServices.java index a481fbf0..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; @@ -91,37 +92,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; } @@ -131,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"); } } @@ -144,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); diff --git a/server/src/main/resources/application.yaml b/server/src/main/resources/application.yaml index f3dadb08..3c3e0ffc 100644 --- a/server/src/main/resources/application.yaml +++ b/server/src/main/resources/application.yaml @@ -59,6 +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. 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 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 new file mode 100644 index 00000000..aecad7b2 --- /dev/null +++ b/server/src/test/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionServiceTest.java @@ -0,0 +1,174 @@ +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.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.time.Duration; +import java.util.HashMap; +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.assertNotNull; +import static org.junit.jupiter.api.Assertions.assertThrows; +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.verify; + +@ExtendWith(MockitoExtension.class) +class DownloadAdmissionServiceTest { + + 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 DownloadAdmissionService service; + + @BeforeEach + void setUp() throws JsonProcessingException { + lenient().when(restServices.buildDownloadParameters(any())) + .thenAnswer(invocation -> new HashMap()); + lenient().when(restServices.submitDownloadJob(anyString(), any())) + .thenAnswer(invocation -> UUID.randomUUID().toString()); + lenient().when(counter.countInFlight(anyString())) + .thenAnswer(invocation -> current(invocation.getArgument(0)).get()); + + service = build(limits(true, 10)); + } + + private DownloadAdmissionService build(DownloadLimitProperties limits) { + return new DownloadAdmissionService(restServices, counter, limits); + } + + private static DownloadLimitProperties limits(boolean enabled, int maxConcurrent) { + return new DownloadLimitProperties(enabled, maxConcurrent, Duration.ofSeconds(15)); + } + + 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 underTheLimitSubmitsAndReturnsTheAwsJobId() throws Exception { + String jobId = service.submit(request(RECIPIENT)); + + assertNotNull(jobId); + verify(restServices).submitDownloadJob(eq(RECIPIENT_JOB_NAME), any()); + verify(restServices).notifyUser(any(), any(), any(), any(), any(), any(), any(), any(), any(), any()); + } + + @Test + void atTheLimitIsRejectedOutright() throws Exception { + current(RECIPIENT).set(10); + + assertThrows(DownloadLimitExceededException.class, () -> service.submit(request(RECIPIENT))); + + verify(restServices, never()).submitDownloadJob(anyString(), any()); + // 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 theRejectionMessageNamesTheConfiguredLimit() { + current(RECIPIENT).set(10); + + DownloadLimitExceededException exception = assertThrows( + DownloadLimitExceededException.class, () -> service.submit(request(RECIPIENT))); + + 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); + + assertThrows(DownloadLimitExceededException.class, () -> service.submit(request(RECIPIENT))); + String jobId = service.submit(request(OTHER_RECIPIENT)); + + assertNotNull(jobId); + } + + @Test + void aDifferentlyCasedAddressIsCountedAgainstTheSameLimit() throws Exception { + current(RECIPIENT).set(10); + + assertThrows(DownloadLimitExceededException.class, + () -> service.submit(request("Person@Example.COM"))); + } + + @Test + void theSweepHappensBeforeTheAdmissionDecision() throws Exception { + service.submit(request(RECIPIENT)); + + verify(counter).refreshIfStale(); + } + + @Test + void aSuccessfulSubmitIsRecordedAgainstTheRecipientForTheGraceWindow() throws Exception { + String jobId = service.submit(request(RECIPIENT)); + + verify(counter).recordSubmitted(jobId, RECIPIENT); + } + + @Test + void theLimitCanBeTurnedOffEntirely() throws Exception { + DownloadAdmissionService disabled = build(limits(false, 10)); + current(RECIPIENT).set(500); + + String jobId = disabled.submit(request(RECIPIENT)); + + assertNotNull(jobId); + verify(restServices).submitDownloadJob(anyString(), any()); + verify(counter, never()).refreshIfStale(); + } + + @Test + 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")); + + assertThrows(IllegalStateException.class, () -> service.submit(request(RECIPIENT))); + + // 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()); + + 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 4d6c073c..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 @@ -38,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 { 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..966dc2da --- /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)), + 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..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; @@ -38,6 +39,9 @@ class RestApiJobsTest { @Mock private DownloadJobStatusService downloadJobStatusService; + @Mock + private DownloadAdmissionService downloadAdmissionService; + private final ObjectMapper objectMapper = new ObjectMapper(); private MockMvc mockMvc; @@ -46,6 +50,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 +58,7 @@ 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.submit(any())).thenReturn(JOB_ID); String body = objectMapper.writeValueAsString(Map.of("inputs", Map.of( "uuid", "collection-id", "recipient", "person@example.com"))); @@ -69,6 +73,23 @@ void postKeepsExistingFieldsAndAddsPureJobId() throws Exception { .andExpect(jsonPath("$.jobID").value(JOB_ID)); } + @Test + 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"))); + + mockMvc.perform(post("/api/v1/ogc/processes/download/execution") + .contentType(MediaType.APPLICATION_JSON) + .accept(MediaType.APPLICATION_JSON) + .content(body)) + .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 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..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; @@ -12,7 +13,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; @@ -23,8 +24,8 @@ 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.inOrder; import static org.mockito.Mockito.never; import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; @@ -35,6 +36,9 @@ public class RestApiTest { @Mock private RestServices restServices; + @Mock + private DownloadAdmissionService downloadAdmissionService; + @InjectMocks private RestApi restApi; @@ -55,7 +59,7 @@ public void setUp() { @Test public void testExecuteDownloadDatasetSuccess() throws JsonProcessingException { - when(restServices.downloadData(any(), any(), any(), any(), any(), any(), any(), any(), any(), any())) + when(downloadAdmissionService.submit(any())) .thenReturn("test-job-id"); ResponseEntity response = restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); @@ -67,16 +71,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()); + } + + @Test + public void testExecutePassesEveryRequestInputToAdmission() throws JsonProcessingException { + when(downloadAdmissionService.submit(any())) + .thenReturn("test-job-id"); + + restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); - // 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()); + ArgumentCaptor captor = ArgumentCaptor.forClass(DownloadRequest.class); + verify(downloadAdmissionService).submit(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.submit(any())) .thenThrow(new RuntimeException("Error while getting dataset")); ResponseEntity response = restApi.execute(ProcessIdEnum.DOWNLOAD_DATASET.getValue(), executeRequest); @@ -87,10 +103,23 @@ 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()); } + @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); 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