fix(backend): spool released exports safely

backend/src/main/kotlin/org/loculus/backend/config/BackendSpringConfig.kt

@@ -39,6 +39,12 @@ object BackendSpringProperty {
     const val PIPELINE_VERSION_UPGRADE_CHECK_INTERVAL_SECONDS =
         "loculus.pipeline-version-upgrade-check.interval-seconds"
     const val STREAM_BATCH_SIZE = "loculus.stream.batch-size"
+    const val STREAM_SPOOL_DIR = "loculus.stream.spool-dir"
+    const val STREAM_MAX_CONCURRENT_SPOOLS = "loculus.stream.max-concurrent-spools"
+    const val STREAM_SPOOL_MAX_TOTAL_BYTES = "loculus.stream.spool-max-total-bytes"
+    const val STREAM_SPOOL_MIN_FREE_BYTES = "loculus.stream.spool-min-free-bytes"
+    const val STREAM_SPOOL_FILE_TTL_MINUTES = "loculus.stream.spool-file-ttl-minutes"
+    const val STREAM_SPOOL_SWEEP_EVERY_MINUTES = "loculus.stream.spool-sweep-every-minutes"
     const val DEBUG_MODE = "loculus.debug-mode"
     const val ENABLE_SEQSETS = "loculus.enable-seqsets"
     const val SEQSET_CITATIONS_RUN_EVERY_MINUTES = "loculus.seqset-citations.run-every-minutes"

backend/src/main/kotlin/org/loculus/backend/config/ReleasedDataSpoolCleanupInterceptor.kt

@@ -0,0 +1,47 @@
+package org.loculus.backend.config
+
+import jakarta.servlet.http.HttpServletRequest
+import jakarta.servlet.http.HttpServletResponse
+import org.loculus.backend.service.submission.SpooledStream
+import org.springframework.web.context.request.NativeWebRequest
+import org.springframework.web.context.request.RequestAttributes
+import org.springframework.web.context.request.async.CallableProcessingInterceptor
+import org.springframework.web.servlet.HandlerInterceptor
+import java.util.concurrent.Callable
+
+internal const val RELEASED_DATA_SPOOL_ATTRIBUTE = "org.loculus.backend.releasedDataSpool"
+
+internal class ReleasedDataSpoolCleanupInterceptor :
+    HandlerInterceptor,
+    CallableProcessingInterceptor {
+    override fun afterCompletion(
+        request: HttpServletRequest,
+        response: HttpServletResponse,
+        handler: Any,
+        ex: Exception?,
+    ) {
+        closeReadySpool(request.getAttribute(RELEASED_DATA_SPOOL_ATTRIBUTE))
+    }
+
+    override fun <T> postProcess(request: NativeWebRequest, task: Callable<T>, concurrentResult: Any?) {
+        closeReadySpool(request.getAttribute(RELEASED_DATA_SPOOL_ATTRIBUTE, RequestAttributes.SCOPE_REQUEST))
+    }
+
+    override fun <T> handleTimeout(request: NativeWebRequest, task: Callable<T>): Any {
+        closeReadySpool(request.getAttribute(RELEASED_DATA_SPOOL_ATTRIBUTE, RequestAttributes.SCOPE_REQUEST))
+        return CallableProcessingInterceptor.RESULT_NONE
+    }
+
+    override fun <T> handleError(request: NativeWebRequest, task: Callable<T>, throwable: Throwable): Any {
+        closeReadySpool(request.getAttribute(RELEASED_DATA_SPOOL_ATTRIBUTE, RequestAttributes.SCOPE_REQUEST))
+        return CallableProcessingInterceptor.RESULT_NONE
+    }
+
+    override fun <T> afterCompletion(request: NativeWebRequest, task: Callable<T>) {
+        closeReadySpool(request.getAttribute(RELEASED_DATA_SPOOL_ATTRIBUTE, RequestAttributes.SCOPE_REQUEST))
+    }
+
+    private fun closeReadySpool(attribute: Any?) {
+        (attribute as? SpooledStream)?.closeIfTransferNotStarted()
+    }
+}

backend/src/main/kotlin/org/loculus/backend/config/SecurityConfig.kt

@@ -13,6 +13,7 @@ import org.springframework.context.annotation.Configuration
 import org.springframework.core.convert.converter.Converter
 import org.springframework.http.HttpMethod
 import org.springframework.security.access.AccessDeniedException
+import org.springframework.security.config.ObjectPostProcessor
 import org.springframework.security.config.annotation.web.builders.HttpSecurity
 import org.springframework.security.config.annotation.web.configuration.EnableWebSecurity
 import org.springframework.security.core.AuthenticationException
@@ -28,6 +29,7 @@ import org.springframework.security.web.access.AccessDeniedHandler
 import org.springframework.security.web.access.AccessDeniedHandlerImpl
 import org.springframework.security.web.access.DelegatingAccessDeniedHandler
 import org.springframework.security.web.csrf.CsrfException
+import org.springframework.security.web.header.HeaderWriterFilter
 import org.springframework.stereotype.Component
 
 private val log = KotlinLogging.logger { }
@@ -80,6 +82,15 @@ class SecurityConfig {
         httpSecurity: HttpSecurity,
         keycloakAuthoritiesConverter: KeycloakAuthenticationConverter,
     ): SecurityFilterChain = httpSecurity
+        .headers { headers ->
+            headers.addObjectPostProcessor(
+                object : ObjectPostProcessor<HeaderWriterFilter> {
+                    override fun <O : HeaderWriterFilter> postProcess(filter: O): O = filter.apply {
+                        setShouldWriteHeadersEagerly(true)
+                    }
+                },
+            )
+        }
         .authorizeHttpRequests { auth ->
             auth.requestMatchers(
                 "/",

backend/src/main/kotlin/org/loculus/backend/config/WebConfig.kt

@@ -4,12 +4,15 @@ import org.loculus.backend.auth.UserConverter
 import org.loculus.backend.log.OrganismMdcInterceptor
 import org.springframework.context.annotation.Configuration
 import org.springframework.web.method.support.HandlerMethodArgumentResolver
+import org.springframework.web.servlet.config.annotation.AsyncSupportConfigurer
 import org.springframework.web.servlet.config.annotation.CorsRegistry
 import org.springframework.web.servlet.config.annotation.InterceptorRegistry
 import org.springframework.web.servlet.config.annotation.WebMvcConfigurer
 
 @Configuration
 class WebConfig(private val backendConfig: BackendConfig) : WebMvcConfigurer {
+    private val releasedDataSpoolCleanup = ReleasedDataSpoolCleanupInterceptor()
+
     override fun addCorsMappings(registry: CorsRegistry) {
         registry.addMapping("/**")
             .allowedOrigins("*") // Allow requests from any origin
@@ -21,6 +24,11 @@ class WebConfig(private val backendConfig: BackendConfig) : WebMvcConfigurer {
     override fun addInterceptors(registry: InterceptorRegistry) {
         registry.addInterceptor(ReadOnlyModeInterceptor(backendConfig))
         registry.addInterceptor(OrganismMdcInterceptor())
+        registry.addInterceptor(releasedDataSpoolCleanup)
+    }
+
+    override fun configureAsyncSupport(configurer: AsyncSupportConfigurer) {
+        configurer.registerCallableInterceptors(releasedDataSpoolCleanup)
     }
 
     override fun addArgumentResolvers(resolvers: MutableList<HandlerMethodArgumentResolver>) {

backend/src/main/kotlin/org/loculus/backend/controller/ExceptionHandler.kt

@@ -116,10 +116,16 @@ class ExceptionHandler : ResponseEntityExceptionHandler() {
     fun handleServiceUnavailableException(e: ServiceUnavailableException): ResponseEntity<ProblemDetail> {
         log.info { "Caught service unavailable exception: ${e.message}" }
 
-        return responseEntity(
-            HttpStatus.SERVICE_UNAVAILABLE,
-            e.message,
-        )
+        val response = ResponseEntity.status(HttpStatus.SERVICE_UNAVAILABLE)
+        e.retryAfterSeconds?.let { response.header(HttpHeaders.RETRY_AFTER, it.toString()) }
+        return response
+            .contentType(MediaType.APPLICATION_PROBLEM_JSON)
+            .body(
+                ProblemDetail.forStatusAndDetail(
+                    HttpStatus.SERVICE_UNAVAILABLE,
+                    e.message ?: HttpStatus.SERVICE_UNAVAILABLE.reasonPhrase,
+                ),
+            )
     }
 
     private fun responseEntity(httpStatus: HttpStatus, detail: String?): ResponseEntity<ProblemDetail> =
@@ -171,4 +177,4 @@ class NotFoundException(message: String) : RuntimeException(message)
 class ProcessingValidationException(message: String) : RuntimeException(message)
 class DuplicateKeyException(message: String) : RuntimeException(message)
 class ConflictException(message: String) : RuntimeException(message)
-class ServiceUnavailableException(message: String) : RuntimeException(message)
+class ServiceUnavailableException(message: String, val retryAfterSeconds: Long? = null) : RuntimeException(message)

backend/src/main/kotlin/org/loculus/backend/controller/SubmissionController.kt

@@ -38,6 +38,7 @@ import org.loculus.backend.api.UnprocessedData
 import org.loculus.backend.auth.AuthenticatedUser
 import org.loculus.backend.auth.HiddenParam
 import org.loculus.backend.config.BackendConfig
+import org.loculus.backend.config.RELEASED_DATA_SPOOL_ATTRIBUTE
 import org.loculus.backend.controller.LoculusCustomHeaders.X_TOTAL_RECORDS
 import org.loculus.backend.log.ORGANISM_MDC_KEY
 import org.loculus.backend.log.REQUEST_ID_MDC_KEY
@@ -52,6 +53,8 @@ import org.loculus.backend.model.SubmissionParams
 import org.loculus.backend.model.SubmitModel
 import org.loculus.backend.service.datauseterms.DataUseTermsPreconditionValidator
 import org.loculus.backend.service.groupmanagement.GroupManagementPreconditionValidator
+import org.loculus.backend.service.submission.SpooledStream
+import org.loculus.backend.service.submission.StreamSpoolService
 import org.loculus.backend.service.submission.SubmissionDatabaseService
 import org.loculus.backend.utils.Accession
 import org.loculus.backend.utils.FastaEntry
@@ -99,6 +102,7 @@ open class SubmissionController(
     private val releasedDataModel: ReleasedDataModel,
     private val submissionDatabaseService: SubmissionDatabaseService,
     private val iteratorStreamer: IteratorStreamer,
+    private val streamSpoolService: StreamSpoolService,
     private val requestIdContext: RequestIdContext,
     private val backendConfig: BackendConfig,
     private val objectMapper: ObjectMapper,
@@ -321,13 +325,25 @@ open class SubmissionController(
         "No database changes since last request " +
             "(Etag in HttpHeaders.IF_NONE_MATCH matches lastDatabaseWriteETag)",
     )
+    @ApiResponse(
+        responseCode = "503",
+        description = "The server cannot create another export yet",
+        headers = [
+            Header(
+                name = "Retry-After",
+                description = "Seconds to wait before retrying",
+                schema = Schema(type = "integer"),
+            ),
+        ],
+    )
     @GetMapping("/get-released-data", produces = [MediaType.APPLICATION_NDJSON_VALUE])
     fun getReleasedData(
         @PathVariable @Valid organism: Organism,
         @RequestParam compression: CompressionFormat?,
         @Parameter(
             description = "(Optional) Only retrieve all released data if Etag has changed.",
         ) @RequestHeader(value = HttpHeaders.IF_NONE_MATCH, required = false) ifNoneMatch: String?,
+        request: HttpServletRequest,
     ): ResponseEntity<StreamingResponseBody> {
         val lastDatabaseWriteETag = releasedDataModel.getLastDatabaseWriteETag(
             tableNames = RELEASED_DATA_RELATED_TABLES,
@@ -342,18 +358,21 @@ open class SubmissionController(
         headers.contentType = MediaType.APPLICATION_NDJSON
         compression?.let { headers.add(HttpHeaders.CONTENT_ENCODING, it.compressionName) }
 
-        val totalRecords = submissionDatabaseService.countReleasedSubmissions(organism)
-        headers.add(X_TOTAL_RECORDS, totalRecords.toString())
-        // TODO(https://github.com/loculus-project/loculus/issues/2778)
-        // There's a possibility that the totalRecords change between the count and the actual query
-        // this is not too bad, if the client ends up with a few more records than expected
-        // We just need to make sure the etag used is from before the count
-        // Alternatively, we could read once to file while counting and then stream the file
-
-        val streamBody = streamTransactioned(compression, endpoint = "get-released-data", organism = organism) {
+        // Count while writing so the response header stays exact.
+        val spooled = streamSpoolService.spool(compression, endpoint = "get-released-data") {
             releasedDataModel.getReleasedData(organism)
         }
-        return ResponseEntity.ok().headers(headers).body(streamBody)
+        try {
+            request.setAttribute(RELEASED_DATA_SPOOL_ATTRIBUTE, spooled)
+            headers.add(X_TOTAL_RECORDS, spooled.recordCount.toString())
+            headers.contentLength = spooled.file.length()
+
+            return ResponseEntity.ok().headers(headers)
+                .body(spooledStreamBody(spooled, endpoint = "get-released-data", organism))
+        } catch (error: Throwable) {
+            spooled.close()
+            throw error
+        }
     }
 
     @Operation(description = GET_DATA_TO_EDIT_SEQUENCE_VERSION_DESCRIPTION)
@@ -464,11 +483,6 @@ open class SubmissionController(
             parsedAccessionVersions,
         )
         headers.add(X_TOTAL_RECORDS, totalRecords.toString())
-        // TODO(https://github.com/loculus-project/loculus/issues/2778)
-        // There's a possibility that the totalRecords change between the count and the actual query
-        // this is not too bad, if the client ends up with a few more records than expected
-        // We just need to make sure the etag used is from before the count
-        // Alternatively, we could read once to file while counting and then stream the file
 
         val streamBody = streamTransactioned(compression, endpoint = "get-submitted-metadata", organism = organism) {
             submissionDatabaseService.streamSubmittedMetadata(
@@ -673,6 +687,27 @@ open class SubmissionController(
         MDC.remove(ORGANISM_MDC_KEY)
     }
 
+    private fun spooledStreamBody(spooled: SpooledStream, endpoint: String, organism: Organism) =
+        StreamingResponseBody { responseBodyStream ->
+            val startTime = System.currentTimeMillis()
+            try {
+                check(spooled.beginTransfer()) { "The spooled response was closed before transfer started" }
+                MDC.put(REQUEST_ID_MDC_KEY, requestIdContext.requestId)
+                MDC.put(ORGANISM_MDC_KEY, organism.name)
+                spooled.file.inputStream().use { fileStream ->
+                    fileStream.copyTo(responseBodyStream)
+                }
+                log.info { "[$endpoint] Response completed in ${System.currentTimeMillis() - startTime}ms" }
+            } catch (error: Exception) {
+                log.error(error) { "[$endpoint] Response failed after ${System.currentTimeMillis() - startTime}ms" }
+                throw error
+            } finally {
+                spooled.close()
+                MDC.remove(REQUEST_ID_MDC_KEY)
+                MDC.remove(ORGANISM_MDC_KEY)
+            }
+        }
+
     fun parseFileMapping(fileMapping: String?, organism: Organism): SubmissionIdFilesMap? {
         val fileMappingParsed = fileMapping?.let {
             if (!backendConfig.getInstanceConfig(organism).schema.submissionDataTypes.files.enabled) {