Add optimize storage infrastructure for backupsv2.

2025-12-24 21:15:48 +00:00 · 2024-09-20 16:47:18 -04:00
parent 7935d12675
commit a10958ee13
16 changed files with 391 additions and 94 deletions
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/internal/backup/InternalBackupPlaygroundViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/internal/backup/InternalBackupPlaygroundViewModel.kt
@@ -383,7 +383,7 @@ class InternalBackupPlaygroundViewModel : ViewModel() {
          )
        } else {
          AppDependencies.jobManager.add(
-            RestoreAttachmentJob(
+            RestoreAttachmentJob.forInitialRestore(
              messageId = insertMessage.messageId,
              attachmentId = archivedAttachment.attachmentId
            )
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/storage/ManageStorageSettingsFragment.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/storage/ManageStorageSettingsFragment.kt
@@ -106,11 +106,11 @@ class ManageStorageSettingsFragment : ComposeFragment() {
            onSetChatLengthLimit = { navController.navigate("set-chat-length-limit") },
            onSyncTrimThreadDeletes = { viewModel.setSyncTrimDeletes(it) },
            onDeleteChatHistory = { navController.navigate("confirm-delete-chat-history") },
-            onToggleOnDeviceStorageOptimization = {
+            onToggleOnDeviceStorageOptimization = { enabled ->
              if (state.onDeviceStorageOptimizationState == ManageStorageSettingsViewModel.OnDeviceStorageOptimizationState.REQUIRES_PAID_TIER) {
                UpgradeToEnableOptimizedStorageSheet().show(parentFragmentManager, BottomSheetUtil.STANDARD_BOTTOM_SHEET_FRAGMENT_TAG)
              } else {
-                viewModel.setOptimizeStorage(it)
+                viewModel.setOptimizeStorage(enabled)
              }
            }
          )
@@ -535,6 +535,7 @@ private fun ManageStorageSettingsScreenPreview() {
      state = ManageStorageSettingsViewModel.ManageStorageState(
        keepMessagesDuration = KeepMessagesDuration.FOREVER,
        lengthLimit = ManageStorageSettingsViewModel.ManageStorageState.NO_LIMIT,
+        syncTrimDeletes = true,
        onDeviceStorageOptimizationState = ManageStorageSettingsViewModel.OnDeviceStorageOptimizationState.DISABLED
      )
    )
--- a/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/storage/ManageStorageSettingsViewModel.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/components/settings/app/storage/ManageStorageSettingsViewModel.kt
@@ -19,6 +19,8 @@ import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.database.SignalDatabase.Companion.media
 import org.thoughtcrime.securesms.database.ThreadTable
 import org.thoughtcrime.securesms.dependencies.AppDependencies
+import org.thoughtcrime.securesms.jobs.OptimizeMediaJob
+import org.thoughtcrime.securesms.jobs.RestoreOptimizedMediaJob
 import org.thoughtcrime.securesms.keyvalue.KeepMessagesDuration
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.util.RemoteConfig
@@ -95,7 +97,12 @@ class ManageStorageSettingsViewModel : ViewModel() {
    val storageState = getOnDeviceStorageOptimizationState()
    if (storageState >= OnDeviceStorageOptimizationState.DISABLED) {
      SignalStore.backup.optimizeStorage = enabled
-      store.update { it.copy(onDeviceStorageOptimizationState = if (enabled) OnDeviceStorageOptimizationState.ENABLED else OnDeviceStorageOptimizationState.DISABLED) }
+      store.update {
+        it.copy(
+          onDeviceStorageOptimizationState = if (enabled) OnDeviceStorageOptimizationState.ENABLED else OnDeviceStorageOptimizationState.DISABLED,
+          storageOptimizationStateChanged = true
+        )
+      }
    }
  }

@@ -112,6 +119,16 @@ class ManageStorageSettingsViewModel : ViewModel() {
    }
  }

+  override fun onCleared() {
+    if (state.value.storageOptimizationStateChanged) {
+      when (state.value.onDeviceStorageOptimizationState) {
+        OnDeviceStorageOptimizationState.DISABLED -> RestoreOptimizedMediaJob.enqueue()
+        OnDeviceStorageOptimizationState.ENABLED -> OptimizeMediaJob.enqueue()
+        else -> Unit
+      }
+    }
+  }
+
  enum class OnDeviceStorageOptimizationState {
    /**
     * The entire feature is not available and the option should not be displayed to the user.
@@ -136,11 +153,12 @@ class ManageStorageSettingsViewModel : ViewModel() {

  @Immutable
  data class ManageStorageState(
-    val keepMessagesDuration: KeepMessagesDuration = KeepMessagesDuration.FOREVER,
-    val lengthLimit: Int = NO_LIMIT,
-    val syncTrimDeletes: Boolean = true,
+    val keepMessagesDuration: KeepMessagesDuration,
+    val lengthLimit: Int,
+    val syncTrimDeletes: Boolean,
    val breakdown: MediaTable.StorageBreakdown? = null,
-    val onDeviceStorageOptimizationState: OnDeviceStorageOptimizationState
+    val onDeviceStorageOptimizationState: OnDeviceStorageOptimizationState,
+    val storageOptimizationStateChanged: Boolean = false
  ) {
    companion object {
      const val NO_LIMIT = 0
--- a/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/AttachmentTable.kt
@@ -79,7 +79,6 @@ import org.thoughtcrime.securesms.crypto.ModernDecryptingPartInputStream
 import org.thoughtcrime.securesms.crypto.ModernEncryptingPartOutputStream
 import org.thoughtcrime.securesms.database.MessageTable.SyncMessageId
 import org.thoughtcrime.securesms.database.SignalDatabase.Companion.messages
-import org.thoughtcrime.securesms.database.SignalDatabase.Companion.stickers
 import org.thoughtcrime.securesms.database.SignalDatabase.Companion.threads
 import org.thoughtcrime.securesms.database.model.databaseprotos.AudioWaveFormData
 import org.thoughtcrime.securesms.dependencies.AppDependencies
@@ -114,6 +113,9 @@ import java.security.NoSuchAlgorithmException
 import java.util.LinkedList
 import java.util.Optional
 import java.util.UUID
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.days
+import kotlin.time.Duration.Companion.hours

 class AttachmentTable(
  context: Context,
@@ -169,6 +171,7 @@ class AttachmentTable(
    const val ARCHIVE_TRANSFER_STATE = "archive_transfer_state"
    const val THUMBNAIL_RESTORE_STATE = "thumbnail_restore_state"
    const val ATTACHMENT_UUID = "attachment_uuid"
+    const val OFFLOAD_RESTORED_AT = "offload_restored_at"

    const val ATTACHMENT_JSON_ALIAS = "attachment_json"

@@ -276,7 +279,8 @@ class AttachmentTable(
        $THUMBNAIL_RANDOM BLOB DEFAULT NULL,
        $THUMBNAIL_RESTORE_STATE INTEGER DEFAULT ${ThumbnailRestoreState.NONE.value},
        $ATTACHMENT_UUID TEXT DEFAULT NULL,
-        $REMOTE_IV BLOB DEFAULT NULL
+        $REMOTE_IV BLOB DEFAULT NULL,
+        $OFFLOAD_RESTORED_AT INTEGER DEFAULT 0
      )
      """

@@ -473,15 +477,6 @@ class AttachmentTable(
      .flatten()
  }

-  fun getArchivableAttachments(): Cursor {
-    return readableDatabase
-      .select(*PROJECTION)
-      .from(TABLE_NAME)
-      .where("$ARCHIVE_MEDIA_ID IS NULL AND $REMOTE_DIGEST IS NOT NULL AND ($TRANSFER_STATE = ? OR $TRANSFER_STATE = ?)", TRANSFER_PROGRESS_DONE.toString(), TRANSFER_NEEDS_RESTORE.toString())
-      .orderBy("$ID DESC")
-      .run()
-  }
-
  fun getLocalArchivableAttachments(): List<LocalArchivableAttachment> {
    return readableDatabase
      .select(*PROJECTION)
@@ -520,6 +515,24 @@ class AttachmentTable(
      }
  }

+  fun getRestorableOptimizedAttachments(): List<RestorableAttachment> {
+    return readableDatabase
+      .select(ID, MESSAGE_ID, DATA_SIZE, REMOTE_DIGEST, REMOTE_KEY)
+      .from(TABLE_NAME)
+      .where("$TRANSFER_STATE = ?", TRANSFER_RESTORE_OFFLOADED)
+      .orderBy("$ID DESC")
+      .run()
+      .readToList {
+        RestorableAttachment(
+          attachmentId = AttachmentId(it.requireLong(ID)),
+          mmsId = it.requireLong(MESSAGE_ID),
+          size = it.requireLong(DATA_SIZE),
+          remoteDigest = it.requireBlob(REMOTE_DIGEST),
+          remoteKey = it.requireBlob(REMOTE_KEY)
+        )
+      }
+  }
+
  fun getRemainingRestorableAttachmentSize(): Long {
    return readableDatabase
      .select("SUM($DATA_SIZE)")
@@ -632,6 +645,52 @@ class AttachmentTable(
    }
  }

+  /**
+   * Marks eligible attachments as offloaded based on their received at timestamp, their last restore time,
+   * presence of thumbnail if media, and the full file being available in the archive.
+   *
+   * Marking offloaded only clears the strong references to the on disk file and clears other local file data like hashes.
+   * Another operation must run to actually delete the data from disk. See [deleteAbandonedAttachmentFiles].
+   */
+  fun markEligibleAttachmentsAsOptimized() {
+    val now = System.currentTimeMillis()
+
+    val subSelect = """
+      SELECT $TABLE_NAME.$ID 
+      FROM $TABLE_NAME 
+      INNER JOIN ${MessageTable.TABLE_NAME} ON ${MessageTable.TABLE_NAME}.${MessageTable.ID} = $TABLE_NAME.$MESSAGE_ID 
+      WHERE
+      (
+        $TABLE_NAME.$OFFLOAD_RESTORED_AT < ${now - 24.hours.inWholeMilliseconds} AND
+        $TABLE_NAME.$TRANSFER_STATE = $TRANSFER_PROGRESS_DONE AND
+        $TABLE_NAME.$ARCHIVE_TRANSFER_STATE = ${ArchiveTransferState.FINISHED.value} AND
+        (
+          $TABLE_NAME.$THUMBNAIL_FILE IS NOT NULL OR 
+          NOT ($TABLE_NAME.$CONTENT_TYPE like 'image/%' OR $TABLE_NAME.$CONTENT_TYPE like 'video/%')
+        ) AND
+        $TABLE_NAME.$DATA_FILE IS NOT NULL
+      )
+      AND
+      (
+        ${MessageTable.TABLE_NAME}.${MessageTable.DATE_RECEIVED} < ${now - 30.days.inWholeMilliseconds}
+      )
+    """
+
+    writableDatabase
+      .update(TABLE_NAME)
+      .values(
+        TRANSFER_STATE to TRANSFER_RESTORE_OFFLOADED,
+        DATA_FILE to null,
+        DATA_RANDOM to null,
+        TRANSFORM_PROPERTIES to null,
+        DATA_HASH_START to null,
+        DATA_HASH_END to null,
+        OFFLOAD_RESTORED_AT to 0
+      )
+      .where("$ID in ($subSelect)")
+      .run()
+  }
+
  /**
   * Returns the number of attachments that are in pending upload states to the archive cdn.
   */
@@ -834,13 +893,18 @@ class AttachmentTable(
      .map { file: File -> file.absolutePath }
      .toSet()

-    val filesInDb: Set<String> = readableDatabase
-      .select(DATA_FILE)
+    val filesInDb: MutableSet<String> = HashSet(filesOnDisk.size)
+
+    readableDatabase
+      .select(DATA_FILE, THUMBNAIL_FILE)
      .from(TABLE_NAME)
      .run()
-      .readToList { it.requireString(DATA_FILE) }
-      .filterNotNull()
-      .toSet() + stickers.allStickerFiles
+      .forEach { cursor ->
+        cursor.requireString(DATA_FILE)?.let { filesInDb += it }
+        cursor.requireString(THUMBNAIL_FILE)?.let { filesInDb += it }
+      }
+
+    filesInDb += SignalDatabase.stickers.allStickerFiles

    val onDiskButNotInDatabase: Set<String> = filesOnDisk - filesInDb

@@ -933,6 +997,10 @@ class AttachmentTable(
    notifyConversationListeners(messages.getThreadIdForMessage(mmsId))
  }

+  fun setThumbnailRestoreState(thumbnailAttachmentId: AttachmentId, thumbnailRestoreState: ThumbnailRestoreState) {
+    setThumbnailRestoreState(listOf(thumbnailAttachmentId), thumbnailRestoreState)
+  }
+
  fun setThumbnailRestoreState(thumbnailAttachmentIds: List<AttachmentId>, thumbnailRestoreState: ThumbnailRestoreState) {
    val prefix: String = when (thumbnailRestoreState) {
      ThumbnailRestoreState.IN_PROGRESS -> {
@@ -958,7 +1026,11 @@ class AttachmentTable(
    }
  }

-  fun setRestoreTransferState(restorableAttachments: Collection<RestorableAttachment>, state: Int) {
+  fun setRestoreTransferState(attachmentId: AttachmentId, state: Int) {
+    setRestoreTransferState(listOf(attachmentId), state)
+  }
+
+  fun setRestoreTransferState(restorableAttachments: Collection<AttachmentId>, state: Int) {
    val prefix = when (state) {
      TRANSFER_RESTORE_OFFLOADED -> "$TRANSFER_STATE != $TRANSFER_PROGRESS_PERMANENT_FAILURE AND"
      TRANSFER_RESTORE_IN_PROGRESS -> "($TRANSFER_STATE = $TRANSFER_NEEDS_RESTORE OR $TRANSFER_STATE = $TRANSFER_RESTORE_OFFLOADED) AND"
@@ -968,7 +1040,7 @@ class AttachmentTable(

    val setQueries = SqlUtil.buildCollectionQuery(
      column = ID,
-      values = restorableAttachments.map { it.attachmentId.id },
+      values = restorableAttachments,
      prefix = prefix
    )

@@ -991,7 +1063,7 @@ class AttachmentTable(
   * @return True if we had to change the digest as part of saving the file, otherwise false.
   */
  @Throws(MmsException::class)
-  fun finalizeAttachmentAfterDownload(mmsId: Long, attachmentId: AttachmentId, inputStream: LimitedInputStream, iv: ByteArray): Boolean {
+  fun finalizeAttachmentAfterDownload(mmsId: Long, attachmentId: AttachmentId, inputStream: LimitedInputStream, iv: ByteArray, offloadRestoredAt: Duration? = null): Boolean {
    Log.i(TAG, "[finalizeAttachmentAfterDownload] Finalizing downloaded data for $attachmentId. (MessageId: $mmsId, $attachmentId)")

    val existingPlaceholder: DatabaseAttachment = getAttachment(attachmentId) ?: throw MmsException("No attachment found for id: $attachmentId")
@@ -1064,6 +1136,10 @@ class AttachmentTable(
        values.put(UPLOAD_TIMESTAMP, 0)
      }

+      if (offloadRestoredAt != null) {
+        values.put(OFFLOAD_RESTORED_AT, offloadRestoredAt.inWholeMilliseconds)
+      }
+
      db.update(TABLE_NAME)
        .values(values)
        .where("$ID = ?", attachmentId.id)
--- a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/SignalDatabaseMigrations.kt
@@ -107,6 +107,7 @@ import org.thoughtcrime.securesms.database.helpers.migration.V245_DeletionTimest
 import org.thoughtcrime.securesms.database.helpers.migration.V246_DropThumbnailCdnFromAttachments
 import org.thoughtcrime.securesms.database.helpers.migration.V247_ClearUploadTimestamp
 import org.thoughtcrime.securesms.database.helpers.migration.V248_ArchiveTransferStateIndex
+import org.thoughtcrime.securesms.database.helpers.migration.V249_AttachmentOffloadRestoredAtColumn

 /**
 * Contains all of the database migrations for [SignalDatabase]. Broken into a separate file for cleanliness.
@@ -215,10 +216,11 @@ object SignalDatabaseMigrations {
    245 to V245_DeletionTimestampOnCallLinks,
    246 to V246_DropThumbnailCdnFromAttachments,
    247 to V247_ClearUploadTimestamp,
-    248 to V248_ArchiveTransferStateIndex
+    248 to V248_ArchiveTransferStateIndex,
+    249 to V249_AttachmentOffloadRestoredAtColumn
  )

-  const val DATABASE_VERSION = 248
+  const val DATABASE_VERSION = 249

  @JvmStatic
  fun migrate(context: Application, db: SQLiteDatabase, oldVersion: Int, newVersion: Int) {
--- a/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V249_AttachmentOffloadRestoredAtColumn.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/database/helpers/migration/V249_AttachmentOffloadRestoredAtColumn.kt
@@ -0,0 +1,19 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.database.helpers.migration
+
+import android.app.Application
+import net.zetetic.database.sqlcipher.SQLiteDatabase
+
+/**
+ * Adds the offload_restored_at column to attachments.
+ */
+@Suppress("ClassName")
+object V249_AttachmentOffloadRestoredAtColumn : SignalDatabaseMigration {
+  override fun migrate(context: Application, db: SQLiteDatabase, oldVersion: Int, newVersion: Int) {
+    db.execSQL("ALTER TABLE attachment ADD COLUMN offload_restored_at INTEGER DEFAULT 0;")
+  }
+}
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/AttachmentDownloadJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/AttachmentDownloadJob.kt
@@ -4,6 +4,7 @@
 */
 package org.thoughtcrime.securesms.jobs

+import androidx.annotation.MainThread
 import okio.Source
 import okio.buffer
 import org.greenrobot.eventbus.EventBus
@@ -82,6 +83,7 @@ class AttachmentDownloadJob private constructor(
    }

    @JvmStatic
+    @MainThread
    fun downloadAttachmentIfNeeded(databaseAttachment: DatabaseAttachment): String? {
      return when (val transferState = databaseAttachment.transferState) {
        AttachmentTable.TRANSFER_PROGRESS_DONE -> null
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/BackupMessagesJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/BackupMessagesJob.kt
@@ -37,14 +37,18 @@ class BackupMessagesJob private constructor(parameters: Parameters) : Job(parame
     */
    fun enqueue(pruneAbandonedRemoteMedia: Boolean = false) {
      val jobManager = AppDependencies.jobManager
+
+      val chain = jobManager.startChain(BackupMessagesJob())
+
      if (pruneAbandonedRemoteMedia) {
-        jobManager
-          .startChain(BackupMessagesJob())
-          .then(SyncArchivedMediaJob())
-          .enqueue()
-      } else {
-        jobManager.add(BackupMessagesJob())
+        chain.then(SyncArchivedMediaJob())
      }
+
+      if (SignalStore.backup.optimizeStorage && SignalStore.backup.backsUpMedia) {
+        chain.then(OptimizeMediaJob())
+      }
+
+      chain.enqueue()
    }
  }

@@ -53,7 +57,7 @@ class BackupMessagesJob private constructor(parameters: Parameters) : Job(parame
      .addConstraint(if (SignalStore.backup.backupWithCellular) NetworkConstraint.KEY else WifiConstraint.KEY)
      .setMaxAttempts(3)
      .setMaxInstancesForFactory(1)
-      .setQueue(BackfillDigestJob.QUEUE) // We want to ensure digests have been backfilled before this runs. Could eventually remove this constraint.b
+      .setQueue(BackfillDigestJob.QUEUE) // We want to ensure digests have been backfilled before this runs. Could eventually remove this constraint.
      .build()
  )

--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/BackupRestoreMediaJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/BackupRestoreMediaJob.kt
@@ -7,8 +7,8 @@ package org.thoughtcrime.securesms.jobs

 import org.signal.core.util.logging.Log
 import org.signal.core.util.withinTransaction
+import org.thoughtcrime.securesms.attachments.AttachmentId
 import org.thoughtcrime.securesms.database.AttachmentTable
-import org.thoughtcrime.securesms.database.AttachmentTable.RestorableAttachment
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.database.model.MmsMessageRecord
 import org.thoughtcrime.securesms.dependencies.AppDependencies
@@ -56,10 +56,10 @@ class BackupRestoreMediaJob private constructor(parameters: Parameters) : BaseJo

    do {
      val restoreThumbnailJobs: MutableList<RestoreAttachmentThumbnailJob> = mutableListOf()
-      val restoreFullAttachmentJobs: MutableMap<RestorableAttachment, RestoreAttachmentJob> = mutableMapOf()
+      val restoreFullAttachmentJobs: MutableList<RestoreAttachmentJob> = mutableListOf()

-      val restoreThumbnailOnlyAttachments: MutableList<RestorableAttachment> = mutableListOf()
-      val notRestorable: MutableList<RestorableAttachment> = mutableListOf()
+      val restoreThumbnailOnlyAttachmentsIds: MutableList<AttachmentId> = mutableListOf()
+      val notRestorable: MutableList<AttachmentId> = mutableListOf()

      val attachmentBatch = SignalDatabase.attachments.getRestorableAttachments(batchSize)
      val messageIds = attachmentBatch.map { it.mmsId }.toSet()
@@ -71,7 +71,7 @@ class BackupRestoreMediaJob private constructor(parameters: Parameters) : BaseJo
        val message = messageMap[attachment.mmsId]
        if (message == null && !isWallpaper) {
          Log.w(TAG, "Unable to find message for ${attachment.attachmentId}, mmsId: ${attachment.mmsId}")
-          notRestorable += attachment
+          notRestorable += attachment.attachmentId
          continue
        }

@@ -82,38 +82,36 @@ class BackupRestoreMediaJob private constructor(parameters: Parameters) : BaseJo
        )

        if (isWallpaper || shouldRestoreFullSize(message!!, restoreTime, SignalStore.backup.optimizeStorage)) {
-          restoreFullAttachmentJobs += attachment to RestoreAttachmentJob(
+          restoreFullAttachmentJobs += RestoreAttachmentJob.forInitialRestore(
            messageId = attachment.mmsId,
            attachmentId = attachment.attachmentId
          )
        } else {
-          restoreThumbnailOnlyAttachments += attachment
+          restoreThumbnailOnlyAttachmentsIds += attachment.attachmentId
        }
      }

      SignalDatabase.rawDatabase.withinTransaction {
        // Mark not restorable thumbnails and attachments as failed
-        SignalDatabase.attachments.setThumbnailRestoreState(notRestorable.map { it.attachmentId }, AttachmentTable.ThumbnailRestoreState.PERMANENT_FAILURE)
+        SignalDatabase.attachments.setThumbnailRestoreState(notRestorable, AttachmentTable.ThumbnailRestoreState.PERMANENT_FAILURE)
        SignalDatabase.attachments.setRestoreTransferState(notRestorable, AttachmentTable.TRANSFER_PROGRESS_FAILED)

-        // Mark restorable thumbnails and attachments as in progress
-        SignalDatabase.attachments.setThumbnailRestoreState(restoreThumbnailJobs.map { it.attachmentId }, AttachmentTable.ThumbnailRestoreState.IN_PROGRESS)
-        SignalDatabase.attachments.setRestoreTransferState(restoreFullAttachmentJobs.keys, AttachmentTable.TRANSFER_RESTORE_IN_PROGRESS)
-
        // Set thumbnail only attachments as offloaded
-        SignalDatabase.attachments.setRestoreTransferState(restoreThumbnailOnlyAttachments, AttachmentTable.TRANSFER_RESTORE_OFFLOADED)
-
-        jobManager.addAll(restoreThumbnailJobs + restoreFullAttachmentJobs.values)
+        SignalDatabase.attachments.setRestoreTransferState(restoreThumbnailOnlyAttachmentsIds, AttachmentTable.TRANSFER_RESTORE_OFFLOADED)
      }
+
+      // Intentionally enqueues one at a time for safer attachment transfer state management
+      restoreThumbnailJobs.forEach { jobManager.add(it) }
+      restoreFullAttachmentJobs.forEach { jobManager.add(it) }
    } while (restoreThumbnailJobs.isNotEmpty() && restoreFullAttachmentJobs.isNotEmpty() && notRestorable.isNotEmpty())

    SignalStore.backup.totalRestorableAttachmentSize = SignalDatabase.attachments.getRemainingRestorableAttachmentSize()

-    jobManager.add(CheckRestoreMediaLeftJob(RestoreAttachmentJob.constructQueueString()))
+    jobManager.add(CheckRestoreMediaLeftJob(RestoreAttachmentJob.constructQueueString(RestoreAttachmentJob.RestoreOperation.INITIAL_RESTORE)))
  }

  private fun shouldRestoreFullSize(message: MmsMessageRecord, restoreTime: Long, optimizeStorage: Boolean): Boolean {
-    return !optimizeStorage || ((restoreTime - message.dateSent) < 30.days.inWholeMilliseconds)
+    return !optimizeStorage || ((restoreTime - message.dateReceived) < 30.days.inWholeMilliseconds)
  }

  override fun onShouldRetry(e: Exception): Boolean = false
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/JobManagerFactories.java
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/JobManagerFactories.java
@@ -188,6 +188,7 @@ public final class JobManagerFactories {
      put(MultiDeviceViewOnceOpenJob.KEY,            new MultiDeviceViewOnceOpenJob.Factory());
      put(MultiDeviceViewedUpdateJob.KEY,            new MultiDeviceViewedUpdateJob.Factory());
      put(NullMessageSendJob.KEY,                    new NullMessageSendJob.Factory());
+      put(OptimizeMediaJob.KEY,                      new OptimizeMediaJob.Factory());
      put(OptimizeMessageSearchIndexJob.KEY,         new OptimizeMessageSearchIndexJob.Factory());
      put(PaymentLedgerUpdateJob.KEY,                new PaymentLedgerUpdateJob.Factory());
      put(PaymentNotificationSendJob.KEY,            new PaymentNotificationSendJob.Factory());
@@ -223,6 +224,7 @@ public final class JobManagerFactories {
      put(RestoreAttachmentJob.KEY,                  new RestoreAttachmentJob.Factory());
      put(RestoreAttachmentThumbnailJob.KEY,         new RestoreAttachmentThumbnailJob.Factory());
      put(RestoreLocalAttachmentJob.KEY,             new RestoreLocalAttachmentJob.Factory());
+      put(RestoreOptimizedMediaJob.KEY,              new RestoreOptimizedMediaJob.Factory());
      put(RetrieveProfileAvatarJob.KEY,              new RetrieveProfileAvatarJob.Factory());
      put(RetrieveProfileJob.KEY,                    new RetrieveProfileJob.Factory());
      put(RetrieveRemoteAnnouncementsJob.KEY,        new RetrieveRemoteAnnouncementsJob.Factory());
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/OptimizeMediaJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/OptimizeMediaJob.kt
@@ -0,0 +1,71 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.jobs
+
+import org.signal.core.util.logging.Log
+import org.thoughtcrime.securesms.database.SignalDatabase
+import org.thoughtcrime.securesms.dependencies.AppDependencies
+import org.thoughtcrime.securesms.jobmanager.Job
+import org.thoughtcrime.securesms.keyvalue.SignalStore
+import kotlin.time.Duration.Companion.days
+
+/**
+ * Optimizes media storage by relying on backups for full copies of files and only keeping thumbnails locally.
+ */
+class OptimizeMediaJob private constructor(parameters: Parameters) : Job(parameters) {
+
+  companion object {
+    private val TAG = Log.tag(OptimizeMediaJob::class)
+    const val KEY = "OptimizeMediaJob"
+
+    fun enqueue() {
+      if (!SignalStore.backup.optimizeStorage || !SignalStore.backup.backsUpMedia) {
+        Log.i(TAG, "Optimize media is not enabled, skipping. backsUpMedia: ${SignalStore.backup.backsUpMedia} optimizeStorage: ${SignalStore.backup.optimizeStorage}")
+        return
+      }
+
+      AppDependencies.jobManager.add(OptimizeMediaJob())
+    }
+  }
+
+  constructor() : this(
+    parameters = Parameters.Builder()
+      .setQueue("OptimizeMediaJob")
+      .setMaxInstancesForQueue(2)
+      .setLifespan(1.days.inWholeMilliseconds)
+      .setMaxAttempts(3)
+      .build()
+  )
+
+  override fun run(): Result {
+    if (!SignalStore.backup.optimizeStorage || !SignalStore.backup.backsUpMedia) {
+      Log.i(TAG, "Optimize media is not enabled, aborting. backsUpMedia: ${SignalStore.backup.backsUpMedia} optimizeStorage: ${SignalStore.backup.optimizeStorage}")
+      return Result.success()
+    }
+
+    Log.i(TAG, "Canceling any previous restore optimized media jobs and cleanup progress")
+    AppDependencies.jobManager.cancelAllInQueue(RestoreAttachmentJob.constructQueueString(RestoreAttachmentJob.RestoreOperation.RESTORE_OFFLOADED))
+    AppDependencies.jobManager.add(CheckRestoreMediaLeftJob(RestoreAttachmentJob.constructQueueString(RestoreAttachmentJob.RestoreOperation.RESTORE_OFFLOADED)))
+
+    Log.i(TAG, "Optimizing media in the db")
+    SignalDatabase.attachments.markEligibleAttachmentsAsOptimized()
+
+    Log.i(TAG, "Deleting abandoned attachment files")
+    SignalDatabase.attachments.deleteAbandonedAttachmentFiles()
+
+    return Result.success()
+  }
+
+  override fun serialize(): ByteArray? = null
+  override fun getFactoryKey(): String = KEY
+  override fun onFailure() = Unit
+
+  class Factory : Job.Factory<OptimizeMediaJob> {
+    override fun create(parameters: Parameters, serializedData: ByteArray?): OptimizeMediaJob {
+      return OptimizeMediaJob(parameters)
+    }
+  }
+}
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentJob.kt
@@ -35,6 +35,7 @@ import org.whispersystems.signalservice.api.push.exceptions.RangeException
 import java.io.File
 import java.io.IOException
 import java.util.concurrent.TimeUnit
+import kotlin.time.Duration.Companion.milliseconds

 /**
 * Download attachment from locations as specified in their record.
@@ -43,45 +44,80 @@ class RestoreAttachmentJob private constructor(
  parameters: Parameters,
  private val messageId: Long,
  private val attachmentId: AttachmentId,
-  private val offloaded: Boolean
+  private val manual: Boolean
 ) : BaseJob(parameters) {

  companion object {
    const val KEY = "RestoreAttachmentJob"
    private val TAG = Log.tag(RestoreAttachmentJob::class.java)

-    @JvmStatic
-    fun constructQueueString(): String {
-      // TODO: decide how many queues
-      return "RestoreAttachmentJob"
+    /**
+     * Create a restore job for the initial large batch of media on a fresh restore
+     */
+    fun forInitialRestore(attachmentId: AttachmentId, messageId: Long): RestoreAttachmentJob {
+      return RestoreAttachmentJob(
+        attachmentId = attachmentId,
+        messageId = messageId,
+        manual = false,
+        queue = constructQueueString(RestoreOperation.INITIAL_RESTORE)
+      )
    }

+    /**
+     * Create a restore job for the large batch of media on a full media restore after disabling optimize media.
+     *
+     * See [RestoreOptimizedMediaJob].
+     */
+    fun forOffloadedRestore(attachmentId: AttachmentId, messageId: Long): RestoreAttachmentJob {
+      return RestoreAttachmentJob(
+        attachmentId = attachmentId,
+        messageId = messageId,
+        manual = false,
+        queue = constructQueueString(RestoreOperation.RESTORE_OFFLOADED)
+      )
+    }
+
+    /**
+     * Restore an attachment when manually triggered by user interaction.
+     *
+     * @return job id of the restore
+     */
    @JvmStatic
    fun restoreAttachment(attachment: DatabaseAttachment): String {
      val restoreJob = RestoreAttachmentJob(
        messageId = attachment.mmsId,
        attachmentId = attachment.attachmentId,
-        offloaded = true
+        manual = true,
+        queue = constructQueueString(RestoreOperation.MANUAL)
      )
+
      AppDependencies.jobManager.add(restoreJob)
      return restoreJob.id
    }
+
+    /**
+     * There are three modes of restore and we use separate queues for each to facilitate canceling if necessary.
+     */
+    @JvmStatic
+    fun constructQueueString(restoreOperation: RestoreOperation): String {
+      return "RestoreAttachmentJob::${restoreOperation.name}"
+    }
  }

-  constructor(messageId: Long, attachmentId: AttachmentId, offloaded: Boolean = false) : this(
+  private constructor(messageId: Long, attachmentId: AttachmentId, manual: Boolean, queue: String) : this(
    Parameters.Builder()
-      .setQueue(constructQueueString())
+      .setQueue(queue)
      .addConstraint(NetworkConstraint.KEY)
-      .setLifespan(TimeUnit.DAYS.toMillis(1))
-      .setMaxAttempts(Parameters.UNLIMITED)
+      .setLifespan(TimeUnit.DAYS.toMillis(30))
+      .setMaxAttempts(3)
      .build(),
    messageId,
    attachmentId,
-    offloaded
+    manual
  )

-  override fun serialize(): ByteArray? {
-    return RestoreAttachmentJobData(messageId = messageId, attachmentId = attachmentId.id, offloaded = offloaded).encode()
+  override fun serialize(): ByteArray {
+    return RestoreAttachmentJobData(messageId = messageId, attachmentId = attachmentId.id, manual = manual).encode()
  }

  override fun getFactoryKey(): String {
@@ -89,15 +125,7 @@ class RestoreAttachmentJob private constructor(
  }

  override fun onAdded() {
-    if (offloaded) {
-      Log.i(TAG, "onAdded() messageId: $messageId  attachmentId: $attachmentId")
-
-      val attachment = SignalDatabase.attachments.getAttachment(attachmentId)
-      if (attachment?.transferState == AttachmentTable.TRANSFER_NEEDS_RESTORE || attachment?.transferState == AttachmentTable.TRANSFER_RESTORE_OFFLOADED) {
-        Log.i(TAG, "onAdded() Marking attachment restore progress as 'started'")
-        SignalDatabase.attachments.setTransferState(messageId, attachmentId, AttachmentTable.TRANSFER_RESTORE_IN_PROGRESS)
-      }
-    }
+    SignalDatabase.attachments.setRestoreTransferState(attachmentId, AttachmentTable.TRANSFER_RESTORE_IN_PROGRESS)
  }

  @Throws(Exception::class)
@@ -137,9 +165,13 @@ class RestoreAttachmentJob private constructor(
  }

  override fun onFailure() {
-    Log.w(TAG, format(this, "onFailure() messageId: $messageId  attachmentId: $attachmentId"))
+    if (isCanceled) {
+      SignalDatabase.attachments.setTransferState(messageId, attachmentId, AttachmentTable.TRANSFER_RESTORE_OFFLOADED)
+    } else {
+      Log.w(TAG, format(this, "onFailure() messageId: $messageId  attachmentId: $attachmentId"))

-    markFailed(messageId, attachmentId)
+      markFailed(messageId, attachmentId)
+    }
  }

  override fun onShouldRetry(exception: Exception): Boolean {
@@ -211,7 +243,7 @@ class RestoreAttachmentJob private constructor(
          )
      }

-      SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, downloadResult.dataStream, downloadResult.iv)
+      SignalDatabase.attachments.finalizeAttachmentAfterDownload(messageId, attachmentId, downloadResult.dataStream, downloadResult.iv, if (manual) System.currentTimeMillis().milliseconds else null)
    } catch (e: RangeException) {
      val transferFile = archiveFile ?: attachmentFile
      Log.w(TAG, "Range exception, file size " + transferFile.length(), e)
@@ -270,8 +302,12 @@ class RestoreAttachmentJob private constructor(
        parameters = parameters,
        messageId = data.messageId,
        attachmentId = AttachmentId(data.attachmentId),
-        offloaded = data.offloaded
+        manual = data.manual
      )
    }
  }
+
+  enum class RestoreOperation {
+    MANUAL, RESTORE_OFFLOADED, INITIAL_RESTORE
+  }
 }
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentThumbnailJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreAttachmentThumbnailJob.kt
@@ -73,6 +73,10 @@ class RestoreAttachmentThumbnailJob private constructor(
    return KEY
  }

+  override fun onAdded() {
+    SignalDatabase.attachments.setThumbnailRestoreState(attachmentId, AttachmentTable.ThumbnailRestoreState.IN_PROGRESS)
+  }
+
  @Throws(Exception::class, IOException::class, InvalidAttachmentException::class, InvalidMessageException::class, MissingConfigurationException::class)
  public override fun onRun() {
    Log.i(TAG, "onRun() messageId: $messageId  attachmentId: $attachmentId")
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreLocalAttachmentJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreLocalAttachmentJob.kt
@@ -9,7 +9,6 @@ import org.signal.core.util.Base64
 import org.signal.core.util.StreamUtil
 import org.signal.core.util.androidx.DocumentFileInfo
 import org.signal.core.util.logging.Log
-import org.signal.core.util.withinTransaction
 import org.signal.libsignal.protocol.InvalidMacException
 import org.signal.libsignal.protocol.InvalidMessageException
 import org.thoughtcrime.securesms.attachments.AttachmentId
@@ -44,14 +43,12 @@ class RestoreLocalAttachmentJob private constructor(
    private const val CONCURRENT_QUEUES = 2

    fun enqueueRestoreLocalAttachmentsJobs(mediaNameToFileInfo: Map<String, DocumentFileInfo>) {
-      var restoreAttachmentJobs: MutableList<Job>
+      val jobManager = AppDependencies.jobManager

      do {
        val possibleRestorableAttachments: List<RestorableAttachment> = SignalDatabase.attachments.getRestorableAttachments(500)
-        val restorableAttachments = ArrayList<RestorableAttachment>(possibleRestorableAttachments.size)
-        val notRestorableAttachments = ArrayList<RestorableAttachment>(possibleRestorableAttachments.size)
-
-        restoreAttachmentJobs = ArrayList(possibleRestorableAttachments.size)
+        val notRestorableAttachments = ArrayList<AttachmentId>(possibleRestorableAttachments.size)
+        val restoreAttachmentJobs: MutableList<Job> = ArrayList(possibleRestorableAttachments.size)

        possibleRestorableAttachments
          .forEachIndexed { index, attachment ->
@@ -63,22 +60,21 @@ class RestoreLocalAttachmentJob private constructor(
            }

            if (fileInfo != null) {
-              restorableAttachments += attachment
              restoreAttachmentJobs += RestoreLocalAttachmentJob(queueName(index), attachment, fileInfo)
            } else {
-              notRestorableAttachments += attachment
+              notRestorableAttachments += attachment.attachmentId
            }
          }

-        SignalDatabase.rawDatabase.withinTransaction {
-          SignalDatabase.attachments.setRestoreTransferState(restorableAttachments, AttachmentTable.TRANSFER_RESTORE_IN_PROGRESS)
-          SignalDatabase.attachments.setRestoreTransferState(notRestorableAttachments, AttachmentTable.TRANSFER_PROGRESS_FAILED)
+        // Mark not restorable attachments as failed
+        SignalDatabase.attachments.setRestoreTransferState(notRestorableAttachments, AttachmentTable.TRANSFER_PROGRESS_FAILED)

-          SignalStore.backup.totalRestorableAttachmentSize = SignalDatabase.attachments.getRemainingRestorableAttachmentSize()
-          AppDependencies.jobManager.addAll(restoreAttachmentJobs)
-        }
+        // Intentionally enqueues one at a time for safer attachment transfer state management
+        restoreAttachmentJobs.forEach { jobManager.add(it) }
      } while (restoreAttachmentJobs.isNotEmpty())

+      SignalStore.backup.totalRestorableAttachmentSize = SignalDatabase.attachments.getRemainingRestorableAttachmentSize()
+
      val checkDoneJobs = (0 until CONCURRENT_QUEUES)
        .map {
          CheckRestoreMediaLeftJob(queueName(it))
--- a/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreOptimizedMediaJob.kt
+++ b/app/src/main/java/org/thoughtcrime/securesms/jobs/RestoreOptimizedMediaJob.kt
@@ -0,0 +1,68 @@
+/*
+ * Copyright 2024 Signal Messenger, LLC
+ * SPDX-License-Identifier: AGPL-3.0-only
+ */
+
+package org.thoughtcrime.securesms.jobs
+
+import org.thoughtcrime.securesms.database.SignalDatabase
+import org.thoughtcrime.securesms.dependencies.AppDependencies
+import org.thoughtcrime.securesms.jobmanager.Job
+import org.thoughtcrime.securesms.keyvalue.SignalStore
+
+/**
+ * Restores any media that was previously optimized and off-loaded into the user's archive. Leverages
+ * the same archive restore progress/flow.
+ */
+class RestoreOptimizedMediaJob private constructor(parameters: Parameters) : Job(parameters) {
+
+  companion object {
+    const val KEY = "RestoreOptimizeMediaJob"
+
+    fun enqueue() {
+      val job = RestoreOptimizedMediaJob()
+      AppDependencies.jobManager.add(job)
+    }
+  }
+
+  private constructor() : this(
+    parameters = Parameters.Builder()
+      .setQueue("RestoreOptimizeMediaJob")
+      .setMaxInstancesForQueue(2)
+      .setMaxAttempts(3)
+      .build()
+  )
+
+  override fun run(): Result {
+    val jobManager = AppDependencies.jobManager
+
+    SignalDatabase
+      .attachments
+      .getRestorableOptimizedAttachments()
+      .forEach {
+        val job = RestoreAttachmentJob.forOffloadedRestore(
+          messageId = it.mmsId,
+          attachmentId = it.attachmentId
+        )
+
+        // Intentionally enqueues one at a time for safer attachment transfer state management
+        jobManager.add(job)
+      }
+
+    SignalStore.backup.totalRestorableAttachmentSize = SignalDatabase.attachments.getRemainingRestorableAttachmentSize()
+
+    AppDependencies.jobManager.add(CheckRestoreMediaLeftJob(RestoreAttachmentJob.constructQueueString(RestoreAttachmentJob.RestoreOperation.RESTORE_OFFLOADED)))
+
+    return Result.success()
+  }
+
+  override fun serialize(): ByteArray? = null
+  override fun getFactoryKey(): String = KEY
+  override fun onFailure() = Unit
+
+  class Factory : Job.Factory<RestoreOptimizedMediaJob> {
+    override fun create(parameters: Parameters, serializedData: ByteArray?): RestoreOptimizedMediaJob {
+      return RestoreOptimizedMediaJob(parameters)
+    }
+  }
+}
--- a/app/src/main/protowire/JobData.proto
+++ b/app/src/main/protowire/JobData.proto
@@ -122,7 +122,7 @@ message BackfillDigestJobData {
 message RestoreAttachmentJobData {
  uint64 messageId = 1;
  uint64 attachmentId = 2;
-  bool offloaded = 3;
+  bool manual = 3;
 }

 message CopyAttachmentToArchiveJobData {