block large download at api level

server/src/main/java/au/org/aodn/ogcapi/server/core/configuration/Config.java

@@ -10,6 +10,7 @@
 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 au.org.aodn.ogcapi.server.processes.DownloadSizeLimitProperties;
 import com.fasterxml.jackson.annotation.JsonInclude;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import jakarta.annotation.PostConstruct;
@@ -37,6 +38,7 @@
         DasProperties.class,
         BatchJobProperties.class,
         DownloadLimitProperties.class,
+        DownloadSizeLimitProperties.class,
         OgcApiProperties.class
 })
 public class Config {

server/src/main/java/au/org/aodn/ogcapi/server/core/exception/DownloadSizeExceededException.java

@@ -0,0 +1,22 @@
+package au.org.aodn.ogcapi.server.core.exception;
+
+import java.math.BigDecimal;
+import java.math.RoundingMode;
+
+public class DownloadSizeExceededException extends RuntimeException {
+
+    private static final BigDecimal BYTES_PER_GB = BigDecimal.valueOf(1024L * 1024 * 1024);
+
+    public DownloadSizeExceededException(long estimatedBytes, long maxBytes) {
+        super("Download is unavailable because the selected dataset is too large (estimated "
+                + toGb(estimatedBytes) + " GB, limit " + toGb(maxBytes) + " GB). "
+                + "Please refine your selection to reduce the dataset size.");
+    }
+
+    private static String toGb(long bytes) {
+        return BigDecimal.valueOf(bytes)
+                .divide(BYTES_PER_GB, 1, RoundingMode.HALF_UP)
+                .stripTrailingZeros()
+                .toPlainString();
+    }
+}

server/src/main/java/au/org/aodn/ogcapi/server/core/exception/GlobalExceptionHandler.java

@@ -130,6 +130,20 @@ public ResponseEntity<ErrorResponse> handleDownloadLimitExceededException(
         return new ResponseEntity<>(errorResponse, HttpStatus.TOO_MANY_REQUESTS);
     }
 
+    @ExceptionHandler(DownloadSizeExceededException.class)
+    public ResponseEntity<ErrorResponse> handleDownloadSizeExceededException(
+            DownloadSizeExceededException ex,
+            WebRequest request) {
+        ErrorResponse errorResponse = ErrorResponse
+                .builder()
+                .timestamp(LocalDateTime.now())
+                .message(ex.getMessage())
+                .details(request.getDescription(false))
+                .build();
+
+        return new ResponseEntity<>(errorResponse, HttpStatus.UNPROCESSABLE_ENTITY);
+    }
+
     @ExceptionHandler(DasUpstreamException.class)
     public ResponseEntity<ErrorResponse> handleDasUpstreamException(DasUpstreamException ex, WebRequest request) {
         ErrorResponse errorResponse = ErrorResponse

server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadAdmissionService.java

@@ -1,6 +1,7 @@
 package au.org.aodn.ogcapi.server.processes;
 
 import au.org.aodn.ogcapi.server.core.exception.DownloadLimitExceededException;
+import au.org.aodn.ogcapi.server.core.exception.DownloadSizeExceededException;
 import com.fasterxml.jackson.core.JsonProcessingException;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -11,9 +12,11 @@
 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.
+ * Decides whether a download can be submitted. Two checks run before anything goes to AWS Batch:
+ * 1. Size: a download estimated at the size limit or more is rejected, the user has to reduce
+ * the subset.
+ * 2. Concurrency: 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
@@ -22,6 +25,7 @@ public class DownloadAdmissionService {
     private final RestServices restServices;
     private final InFlightDownloadCounter counter;
     private final DownloadLimitProperties limits;
+    private final DownloadSizeLimitProperties sizeLimit;
 
     private final ReentrantLock lock = new ReentrantLock();
 
@@ -38,21 +42,29 @@ public class DownloadAdmissionService {
     public DownloadAdmissionService(
             RestServices restServices,
             InFlightDownloadCounter counter,
-            DownloadLimitProperties limits) {
+            DownloadLimitProperties limits,
+            DownloadSizeLimitProperties sizeLimit) {
         this.restServices = restServices;
         this.counter = counter;
         this.limits = limits;
+        this.sizeLimit = sizeLimit;
     }
 
     /**
      * Submit a download to AWS Batch and return its job id.
      *
+     * @throws DownloadSizeExceededException the download is estimated at the size limit or more
      * @throws DownloadLimitExceededException the recipient already has {@code maxConcurrent}
      *                                         downloads running
      */
     public String submit(DownloadRequest request) throws JsonProcessingException {
         String key = InFlightDownloadCounter.recipientKey(request.recipient());
 
+        if (sizeLimit.enabled()) {
+            // Before the slot is reserved, so a slow estimate does not hold one.
+            rejectIfTooLarge(request);
+        }
+
         if (limits.enabled()) {
             // Outside the lock: the sweep is the only expensive step.
             counter.refreshIfStale();
@@ -73,6 +85,25 @@ public String submit(DownloadRequest request) throws JsonProcessingException {
         }
     }
 
+    private void rejectIfTooLarge(DownloadRequest request) {
+        long estimatedBytes;
+        try {
+            estimatedBytes = restServices.estimateDownloadBytes(request);
+        } catch (Exception e) {
+            // Let it through, as the portal does when its own estimate fails. Blocking every
+            // download because DAS cannot estimate would be worse than one job running out of disk.
+            log.warn("Size estimate failed for uuid {}, submitting without the size check", request.uuid(), e);
+            return;
+        }
+
+        long maxBytes = sizeLimit.maxSize().toBytes();
+        if (estimatedBytes >= maxBytes) {
+            log.info("Rejected download for uuid {}: estimated {} bytes, limit {} bytes",
+                    request.uuid(), estimatedBytes, maxBytes);
+            throw new DownloadSizeExceededException(estimatedBytes, maxBytes);
+        }
+    }
+
     private void reserveOrReject(DownloadRequest request, String key) {
         lock.lock();
         try {

server/src/main/java/au/org/aodn/ogcapi/server/processes/DownloadSizeLimitProperties.java

@@ -0,0 +1,16 @@
+package au.org.aodn.ogcapi.server.processes;
+
+import org.springframework.boot.context.properties.ConfigurationProperties;
+import org.springframework.boot.context.properties.bind.DefaultValue;
+import org.springframework.util.unit.DataSize;
+
+/**
+ * Largest download we accept. The batch job writes its output to the Fargate disk, which is
+ * 200 GB, so a download estimated at maxSize or more is rejected before it is submitted.
+ */
+@ConfigurationProperties(prefix = "aws.batch.job.size-limit")
+public record DownloadSizeLimitProperties(
+        @DefaultValue("true") boolean enabled,
+        @DefaultValue("180GB") DataSize maxSize
+) {
+}

server/src/main/java/au/org/aodn/ogcapi/server/processes/RestApi.java

@@ -9,6 +9,7 @@
 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.exception.DownloadSizeExceededException;
 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;
@@ -81,6 +82,17 @@ public class RestApi implements ProcessesApi, JobsApi {
                               "message": "You already have 10 downloads in progress. Wait for one of them to complete before starting another."
                             }
                             """)))
+    @ApiResponse(
+            responseCode = "422",
+            description = "The download is estimated at the size limit or more.",
+            content = @Content(
+                    mediaType = MediaType.APPLICATION_JSON_VALUE,
+                    schema = @Schema(implementation = ErrorResponse.class),
+                    examples = @ExampleObject(value = """
+                            {
+                              "message": "Download is unavailable because the selected dataset is too large (estimated 250.3 GB, limit 180 GB). Please refine your selection to reduce the dataset size."
+                            }
+                            """)))
     public ResponseEntity<InlineResponse200> execute(
             @Parameter(in = ParameterIn.PATH, required = true, schema = @Schema())
             @PathVariable("processID")
@@ -108,8 +120,9 @@ public ResponseEntity<InlineResponse200> execute(
                 DownloadRequest request = new DownloadRequest(uuid, key, startDate, endDate, multiPolygon,
                         recipient, 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 size limit and the per-user limit are applied here: a download that is too
+                // large, or 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
@@ -123,8 +136,8 @@ public ResponseEntity<InlineResponse200> execute(
 
                 return ResponseEntity.ok(results);
 
-            } catch (DownloadLimitExceededException e) {
-                // Let GlobalExceptionHandler turn this into a real 429 with a message the
+            } catch (DownloadLimitExceededException | DownloadSizeExceededException e) {
+                // Let GlobalExceptionHandler turn these into a real 429 or 422 with a message the
                 // caller can act on, instead of the generic 200-wrapped error below.
                 throw e;
             } catch (Exception e) {

server/src/main/java/au/org/aodn/ogcapi/server/processes/RestServices.java

@@ -7,9 +7,11 @@
 import au.org.aodn.ogcapi.server.core.service.das.DasService;
 import au.org.aodn.ogcapi.server.core.service.geoserver.wfs.DownloadWfsDataService;
 import au.org.aodn.ogcapi.server.core.service.sse.SseStreamHandler;
+import au.org.aodn.ogcapi.server.core.util.DasSseFrames;
 import au.org.aodn.ogcapi.server.core.util.EmailUtils;
 import au.org.aodn.ogcapi.server.core.util.SubsetParametersUtils;
 import com.fasterxml.jackson.core.JsonProcessingException;
+import com.fasterxml.jackson.databind.JsonNode;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import lombok.extern.slf4j.Slf4j;
 import org.springframework.beans.factory.annotation.Autowired;
@@ -32,6 +34,9 @@
 @Slf4j
 public class RestServices {
 
+    // Same field the portal reads to block a download that is too large.
+    private static final String ESTIMATED_OUTPUT_BYTES = "estimated_output_bytes";
+
     private final BatchClient batchClient;
     private final ObjectMapper objectMapper;
     private final String batchJobDefinition;
@@ -115,6 +120,27 @@ public Map<String, String> buildDownloadParameters(DownloadRequest request) thro
         return parameters;
     }
 
+    /**
+     * Estimate how many bytes a download will produce. Uses the same uuid and subset parameters
+     * as the cloud-optimised estimate, so a subset the portal has just estimated comes from the
+     * cache instead of calling DAS again.
+     */
+    public long estimateDownloadBytes(DownloadRequest request) throws JsonProcessingException {
+        Map<String, String> parameters = SubsetParametersUtils.buildSubsetParameters(
+                objectMapper, request.uuid(), request.key(), request.startDate(), request.endDate(),
+                request.multiPolygon(), request.outputFormat());
+
+        // No client of our own to keep alive, so heartbeats are ignored.
+        String estimateJson = dasService.estimateCloudOptimisedDownloadSize(
+                request.uuid(), parameters, DasSseFrames.FrameCallback.IGNORE);
+
+        JsonNode bytes = objectMapper.readTree(estimateJson).get(ESTIMATED_OUTPUT_BYTES);
+        if (bytes == null || !bytes.isNumber()) {
+            throw new IllegalStateException("Size estimate has no " + ESTIMATED_OUTPUT_BYTES + ": " + estimateJson);
+        }
+        return bytes.longValue();
+    }
+
     /**
      * 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.
@@ -123,7 +149,9 @@ public static String downloadJobName(String recipient) {
         return "generating-data-file-for-" + recipient.replaceAll("[^a-zA-Z0-9-_]", "-");
     }
 
-    /** Submit a prepared download to the configured queue and job definition. */
+    /**
+     * Submit a prepared download to the configured queue and job definition.
+     */
     public String submitDownloadJob(String jobName, Map<String, String> parameters) {
         String jobId = submitJob(jobName, this.batchJobQueue, this.batchJobDefinition, parameters);
         log.info("Job submitted with ID: {}", jobId);

server/src/main/resources/application.yaml

@@ -66,6 +66,12 @@ aws:
           max-concurrent: 10
           # How long the shared in-flight snapshot is reused before the queues are swept again.
           refresh-interval: 15s
+        # Largest download accepted. The job writes to the 200 GB Fargate disk, so a download
+        # estimated at max-size or more is rejected before it is submitted. Keep in line with
+        # EXTRA_LARGE_DOWNLOAD_BYTES in the portal.
+        size-limit:
+          enabled: true
+          max-size: 180GB
 
 wfs-default-param:
   fields:

server/src/test/java/au/org/aodn/ogcapi/server/core/service/das/DasServiceCacheTest.java

@@ -3,6 +3,9 @@
 import au.org.aodn.ogcapi.server.core.configuration.CacheConfig;
 import au.org.aodn.ogcapi.server.core.http.RecordingSseConnector;
 import au.org.aodn.ogcapi.server.core.util.DasSseFrames;
+import au.org.aodn.ogcapi.server.core.util.SubsetParametersUtils;
+import au.org.aodn.ogcapi.server.processes.DownloadRequest;
+import au.org.aodn.ogcapi.server.processes.RestServices;
 import com.fasterxml.jackson.databind.ObjectMapper;
 import org.junit.jupiter.api.AfterEach;
 import org.junit.jupiter.api.BeforeEach;
@@ -13,6 +16,7 @@
 import org.springframework.context.annotation.AnnotationConfigApplicationContext;
 import org.springframework.context.annotation.Bean;
 import org.springframework.context.annotation.Configuration;
+import org.springframework.test.util.ReflectionTestUtils;
 import org.springframework.web.client.RestTemplate;
 import org.springframework.web.reactive.function.client.WebClient;
 
@@ -188,4 +192,27 @@ public void testCacheHitSkipsTheStreamAndItsHeartbeats() {
         estimate(parameters("netcdf"), heartbeats::incrementAndGet);
         assertEquals(1, heartbeats.get(), "A cache hit never opens a stream, so nothing heartbeats");
     }
+
+    /**
+     * The size check before a download is submitted must reuse the estimate the portal has just
+     * made for the same subset. The portal estimate builds its parameters with
+     * SubsetParametersUtils, as the SSE estimate in RestServices does.
+     */
+    @Test
+    public void testDownloadSizeCheckReusesThePortalEstimate() throws Exception {
+        ObjectMapper objectMapper = new ObjectMapper();
+        RestServices restServices = new RestServices(null, objectMapper, "test-job-definition", "test-job-queue");
+        ReflectionTestUtils.setField(restServices, "dasService", dasService);
+
+        connector.respondWith(List.of(RESULT_FRAME));
+        estimate(SubsetParametersUtils.buildSubsetParameters(objectMapper, "test-uuid", "a.zarr",
+                "2020-01-01", "2020-12-31", "non-specified", "netcdf"), freshCallback());
+
+        long bytes = restServices.estimateDownloadBytes(new DownloadRequest("test-uuid", "a.zarr",
+                "2020-01-01", "2020-12-31", "non-specified", "person@example.com", "Test Collection",
+                "https://portal.example.test/details/test-uuid", "Cite as", "netcdf"));
+
+        assertEquals(123, bytes);
+        assertEquals(1, connector.requests(), "The size check is answered from the cache, not by data-access-service");
+    }
 }