feat(backend): S3 garbage collection

backend/docs/db/schema.sql

@@ -935,6 +935,13 @@ ALTER TABLE ONLY public.user_groups_table
 CREATE INDEX data_use_terms_table_accession_idx ON public.data_use_terms_table USING btree (accession);
 
 
+--
+-- Name: files_upload_requested_at_idx; Type: INDEX; Schema: public; Owner: postgres
+--
+
+CREATE INDEX files_upload_requested_at_idx ON public.files USING btree (upload_requested_at);
+
+
 --
 -- Name: flyway_schema_history_s_idx; Type: INDEX; Schema: public; Owner: postgres
 --

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

@@ -43,6 +43,9 @@ object BackendSpringProperty {
     const val ENABLE_SEQSETS = "loculus.enable-seqsets"
 
     const val S3_ENABLED = "loculus.s3.enabled"
+    const val S3_GC_ENABLED = "loculus.s3.gc-enabled"
+    const val S3_GC_GRACE_PERIOD_MINUTES = "loculus.s3.gc-grace-period-minutes"
+    const val S3_GC_INITIAL_DELAY_MINUTES = "loculus.s3.gc-initial-delay-minutes"
     const val S3_BUCKET_ENDPOINT = "loculus.s3.bucket.endpoint"
     const val S3_BUCKET_INTERNAL_ENDPOINT = "loculus.s3.bucket.internal-endpoint"
     const val S3_BUCKET_BUCKET = "loculus.s3.bucket.bucket"

backend/src/main/kotlin/org/loculus/backend/service/files/FilesDatabaseService.kt

@@ -1,10 +1,15 @@
 package org.loculus.backend.service.files
 
+import kotlinx.datetime.LocalDateTime
+import org.jetbrains.exposed.sql.SqlExpressionBuilder.eq
 import org.jetbrains.exposed.sql.and
+import org.jetbrains.exposed.sql.deleteWhere
 import org.jetbrains.exposed.sql.insert
+import org.jetbrains.exposed.sql.kotlin.datetime.KotlinLocalDateTimeColumnType
 import org.jetbrains.exposed.sql.not
+import org.jetbrains.exposed.sql.statements.StatementType
+import org.jetbrains.exposed.sql.transactions.transaction
 import org.jetbrains.exposed.sql.update
-import org.loculus.backend.utils.DatabaseConstants
 import org.loculus.backend.utils.DateProvider
 import org.loculus.backend.utils.chunkedForDatabase
 import org.springframework.stereotype.Service
@@ -26,6 +31,10 @@ class FilesDatabaseService(private val dateProvider: DateProvider) {
         }
     }
 
+    fun deleteFileEntry(fileId: UUID) {
+        FilesTable.deleteWhere { FilesTable.idColumn eq fileId }
+    }
+
     fun getGroupIds(fileIds: Set<FileId>): Map<FileId, Int> = fileIds.chunkedForDatabase({ chunk ->
         FilesTable.select(FilesTable.idColumn, FilesTable.groupIdColumn)
             .where { FilesTable.idColumn inList chunk }
@@ -57,6 +66,48 @@ class FilesDatabaseService(private val dateProvider: DateProvider) {
         chunk.filterNot { it in existingIds }
     }, 1).toSet()
 
+    fun getOrphanedFileIds(threshold: LocalDateTime): Set<FileId> {
+        val sql = """           
+            -- check for files for which an upload was requested > threshold days ago
+            -- but are not referenced by a submission. For this, check the submitted_data
+            -- and processed_data jsonb objects (but not archive_of_submitted_data)
+            WITH referenced AS (
+              -- fetch ids for files uploaded by users and referenced in submissions
+              SELECT (fil->>'fileId')::uuid AS file_id              
+              FROM sequence_entries,
+                 LATERAL jsonb_each(COALESCE(NULLIF(submitted_data->'files', 'null'::jsonb),'{}'::jsonb)) AS cat(k,v),
+                 LATERAL jsonb_array_elements(cat.v) AS fil
+              UNION
+              -- also need to check processed_data since preprocessing
+              -- can create files that are never referenced in submissions
+              SELECT (fil->>'fileId')::uuid AS file_id
+              FROM sequence_entries_preprocessed_data sepd
+              JOIN sequence_entries se
+                  ON se.accession = sepd.accession
+                 AND se.version   = sepd.version,
+                 LATERAL jsonb_each(COALESCE(NULLIF(sepd.processed_data->'files', 'null'::jsonb),'{}'::jsonb)) AS cat(k,v),
+                 LATERAL jsonb_array_elements(cat.v) AS fil
+            )
+            SELECT f.id FROM files f
+              LEFT JOIN referenced r ON r.file_id = f.id
+              WHERE r.file_id IS NULL
+                  AND f.upload_requested_at < ?;
+        """.trimIndent()
+        return transaction {
+            exec(
+                sql,
+                listOf(KotlinLocalDateTimeColumnType() to threshold),
+                explicitStatementType = StatementType.SELECT,
+            ) { rs ->
+                buildSet<FileId> {
+                    while (rs.next()) {
+                        add(rs.getObject("id", UUID::class.java))
+                    }
+                }
+            } ?: emptySet()
+        }
+    }
+
     /**
      * Return the subset of file IDs for which the file size hasn't been checked yet or
      * no file has been uploaded yet (and therefore there's no file size).

backend/src/main/kotlin/org/loculus/backend/service/files/S3Service.kt

@@ -14,6 +14,7 @@ import software.amazon.awssdk.services.s3.model.CompleteMultipartUploadRequest
 import software.amazon.awssdk.services.s3.model.CompletedMultipartUpload
 import software.amazon.awssdk.services.s3.model.CompletedPart
 import software.amazon.awssdk.services.s3.model.CreateMultipartUploadRequest
+import software.amazon.awssdk.services.s3.model.DeleteObjectRequest
 import software.amazon.awssdk.services.s3.model.GetObjectRequest
 import software.amazon.awssdk.services.s3.model.HeadObjectRequest
 import software.amazon.awssdk.services.s3.model.PutObjectRequest
@@ -176,6 +177,18 @@ class S3Service(private val s3Config: S3Config) {
         Unit
     }
 
+    fun deleteFile(fileId: FileId) = s3ErrorMapping {
+        val config = getS3BucketConfig()
+        s3Client.deleteObject(
+            DeleteObjectRequest
+                .builder()
+                .bucket(config.bucket)
+                .key(getFileIdPath(fileId))
+                .build(),
+        )
+        Unit
+    }
+
     private fun assertIsEnabled() {
         if (!s3Config.enabled) {
             throw IllegalStateException("S3 is not enabled")
@@ -258,6 +271,11 @@ fun <T> s3ErrorMapping(block: () -> T): T {
                     "by part number.",
             )
 
+            "NoSuchKey" -> UnprocessableEntityException(
+                "NoSuchKey: The referenced file does not exist in storage. Uploaded files that are not " +
+                    "referenced by a submission for too long are cleaned up automatically.",
+            )
+
             else -> RuntimeException("Unexpected S3 error: ${e.awsErrorDetails().errorCode()}")
         }
     }

backend/src/main/kotlin/org/loculus/backend/service/submission/S3GarbageCollectionTask.kt

@@ -0,0 +1,92 @@
+package org.loculus.backend.service.submission
+
+import kotlinx.datetime.DateTimeUnit
+import kotlinx.datetime.minus
+import kotlinx.datetime.toLocalDateTime
+import org.loculus.backend.config.BackendSpringProperty
+import org.loculus.backend.log.AuditLogger
+import org.loculus.backend.service.files.FilesDatabaseService
+import org.loculus.backend.service.files.S3Service
+import org.loculus.backend.utils.DateProvider
+import org.springframework.beans.factory.annotation.Value
+import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty
+import org.springframework.scheduling.annotation.Scheduled
+import org.springframework.stereotype.Component
+import java.util.UUID
+import java.util.concurrent.TimeUnit
+import kotlin.math.max
+
+private val log = mu.KotlinLogging.logger {}
+
+@Component
+@ConditionalOnProperty("loculus.s3.enabled", havingValue = "true")
+class S3GarbageCollectionTask(
+    private val filesDatabaseService: FilesDatabaseService,
+    private val s3Service: S3Service,
+    private val dateProvider: DateProvider,
+    private val auditLogger: AuditLogger,
+    @Value("\${${BackendSpringProperty.S3_GC_GRACE_PERIOD_MINUTES}}") private val gracePeriod: Int,
+    @Value("\${${BackendSpringProperty.S3_GC_ENABLED}}") private val deleteOrphans: Boolean,
+) {
+
+    /**
+     * Deletes S3 objects older than `loculus.s3.gc-grace-period-minutes` that are
+     * not referenced in submitted_data or processed_data
+     */
+    @Scheduled(
+        initialDelayString = "\${${BackendSpringProperty.S3_GC_INITIAL_DELAY_MINUTES}}",
+        fixedDelay = 60 * 24,
+        timeUnit = TimeUnit.MINUTES,
+    )
+    fun task() {
+        // `gracePeriod` must be at least 1 or files produced by preprocessing may be
+        // garbage collected before they're attached to sequence entries
+        val gracePeriod = max(gracePeriod, 1)
+        log.info {
+            "Running S3 garbage collection task to clean up orphan files at least $gracePeriod " +
+                "minutes old (garbageCollectionEnabled = $deleteOrphans)"
+        }
+
+        val threshold = dateProvider.getCurrentInstant()
+            .minus(gracePeriod, DateTimeUnit.MINUTE, DateProvider.timeZone)
+            .toLocalDateTime(DateProvider.timeZone)
+        val orphans = filesDatabaseService.getOrphanedFileIds(threshold)
+
+        if (!deleteOrphans) {
+            log.info { "S3 garbage collection task would have deleted ${orphans.size} files: $orphans" }
+            return
+        }
+
+        var deleteFailures = 0
+        orphans.forEach { fileId: UUID ->
+            try {
+                s3Service.deleteFile(fileId)
+                filesDatabaseService.deleteFileEntry(fileId)
+            } catch (e: Exception) {
+                log.warn("Failed to delete $fileId", e)
+                deleteFailures++
+            }
+        }
+
+        if (orphans.isNotEmpty()) {
+            log.info {
+                "S3 garbage collection task deleted ${orphans.size - deleteFailures} orphan(s) not referenced by a " +
+                    "submission after $gracePeriod minutes"
+            }
+            auditLogger
+                .log(
+                    "CLEANUP",
+                    "S3 garbage collection task deleted ${orphans.size - deleteFailures} orphan(s) " +
+                        "not referenced by a submission after $gracePeriod minutes",
+                )
+
+            if (deleteFailures > 0) {
+                log.warn {
+                    "S3 garbage collection task unsuccessfully attempted to delete $deleteFailures orphan file(s)"
+                }
+            }
+        } else {
+            log.info { "S3 garbage collection task identified no orphan files on S3" }
+        }
+    }
+}

backend/src/main/resources/application.properties

@@ -33,6 +33,9 @@ loculus.stream.batch-size=1000
 loculus.debug-mode=false
 
 loculus.s3.enabled=false
+loculus.s3.gc-enabled=false
+loculus.s3.gc-grace-period-minutes=1440
+loculus.s3.gc-initial-delay-minutes=15
 loculus.s3.bucket.endpoint=
 loculus.s3.bucket.bucket=
 loculus.s3.bucket.region=

backend/src/main/resources/db/migration/V1.31__add_files_upload_requested_at_index.sql

@@ -0,0 +1 @@
+CREATE INDEX IF NOT EXISTS files_upload_requested_at_idx ON files (upload_requested_at);

backend/src/test/kotlin/org/loculus/backend/service/TestHelpers.kt

@@ -0,0 +1,26 @@
+package org.loculus.backend.service
+
+import kotlinx.datetime.DateTimeUnit
+import kotlinx.datetime.LocalDateTime
+import kotlinx.datetime.minus
+import kotlinx.datetime.toLocalDateTime
+import org.jetbrains.exposed.sql.insert
+import org.jetbrains.exposed.sql.transactions.transaction
+import org.loculus.backend.service.files.FilesTable
+import org.loculus.backend.utils.DateProvider
+import java.util.UUID
+import kotlin.time.Clock
+
+fun insertFile(id: UUID, groupId: Int, requestedAt: LocalDateTime, uploader: String = "testuser") = transaction {
+    FilesTable.insert {
+        it[idColumn] = id
+        it[uploadRequestedAtColumn] = requestedAt
+        it[uploaderColumn] = uploader
+        it[groupIdColumn] = groupId
+        it[multipartCompleted] = true
+    }
+}
+
+fun daysAgo(days: Long): LocalDateTime = Clock.System.now()
+    .minus(days, DateTimeUnit.DAY, DateProvider.timeZone)
+    .toLocalDateTime(DateProvider.timeZone)