Ensure that memory-only jobs do not get lost due to cache eviction.

app/src/main/java/org/thoughtcrime/securesms/jobs/FastJobStorage.kt

@@ -1,6 +1,7 @@
 package org.thoughtcrime.securesms.jobs
 
 import androidx.annotation.VisibleForTesting
+import kotlinx.collections.immutable.toImmutableSet
 import org.signal.core.util.Stopwatch
 import org.signal.core.util.logging.Log
 import org.thoughtcrime.securesms.database.JobDatabase
@@ -28,9 +29,9 @@ class FastJobStorage(private val jobDatabase: JobDatabase) : JobStorage {
 
   /**
    * We keep a set of job specs in memory to facilitate fast retrieval. This is important because the most common job storage pattern is
-   * [getNextEligibleJob], which needs to return full specs.
+   * [getNextEligibleJob], which needs to return full specs. Memory-only jobs are protected from eviction since they don't exist in the database.
    */
-  private val jobSpecCache: LRUCache<String, JobSpec> = LRUCache(JOB_CACHE_LIMIT)
+  private val jobSpecCache: JobSpecCache = JobSpecCache(JOB_CACHE_LIMIT)
 
   /**
    * We keep a set of constraints in memory, seeded by the same jobs in the [jobSpecCache]. It doesn't need to necessarily stay in sync with that cache, though.
@@ -181,7 +182,7 @@ class FastJobStorage(private val jobDatabase: JobDatabase) : JobStorage {
   override fun getJobsInQueue(queue: String): List<JobSpec> {
     return minimalJobs
       .filter { it.queueKey == queue }
-      .mapNotNull { it.toJobSpec() }
+      .map { it.toJobSpec() }
   }
 
   @Synchronized
@@ -348,7 +349,7 @@ class FastJobStorage(private val jobDatabase: JobDatabase) : JobStorage {
 
     val deleteIds: Set<String> = ids.toSet()
     minimalJobs.removeIf { deleteIds.contains(it.id) }
-    jobSpecCache.keys.removeAll(deleteIds)
+    deleteIds.forEach { jobSpecCache.remove(it) }
     eligibleJobs.removeIf { deleteIds.contains(it.id) }
     migrationJobs.removeIf { deleteIds.contains(it.id) }
 
@@ -596,10 +597,19 @@ class FastJobStorage(private val jobDatabase: JobDatabase) : JobStorage {
    * Converts a [MinimalJobSpec] to a [JobSpec]. We prefer using the cache, but if it's not found, we'll hit the database.
    * We consider this a "recent access" and will cache it for future use.
    */
-  private fun MinimalJobSpec.toJobSpec(): JobSpec? {
-    return jobSpecCache.getOrPut(this.id) {
-      jobDatabase.getJobSpec(this.id) ?: return null
+  private fun MinimalJobSpec.toJobSpec(): JobSpec {
+    val cachedJob = jobSpecCache[this.id]
+    if (cachedJob != null) {
+      return cachedJob
     }
+
+    val dbJob = jobDatabase.getJobSpec(this.id)
+    if (dbJob != null) {
+      jobSpecCache[dbJob.id] = dbJob
+      return dbJob
+    }
+
+    error("Failed to get a full JobSpec for ${this.id}")
   }
 
   private object EligibleMinJobComparator : Comparator<MinimalJobSpec> {
@@ -668,3 +678,35 @@ fun JobSpec.toMinimalJobSpec(): MinimalJobSpec {
     initialDelay = this.initialDelay
   )
 }
+
+/**
+ * A cache over top of LRUCache to protect memory-only jobs from eviction.
+ * Memory-only jobs are stored separately and never evicted since they don't exist in the database.
+ */
+private class JobSpecCache(capacity: Int) {
+
+  private val durableJobCache = LRUCache<String, JobSpec>(capacity)
+  private val memoryOnlyJobs = mutableMapOf<String, JobSpec>()
+
+  operator fun get(key: String): JobSpec? {
+    return memoryOnlyJobs[key] ?: durableJobCache[key]
+  }
+
+  operator fun set(key: String, value: JobSpec) {
+    if (value.isMemoryOnly) {
+      memoryOnlyJobs[key] = value
+    } else {
+      durableJobCache[key] = value
+    }
+  }
+
+  fun remove(key: String): JobSpec? {
+    return memoryOnlyJobs.remove(key) ?: durableJobCache.remove(key)
+  }
+
+  val keys: Set<String>
+    get() = (durableJobCache.keys + memoryOnlyJobs.keys).toImmutableSet()
+
+  val size: Int
+    get() = durableJobCache.size + memoryOnlyJobs.size
+}

app/src/test/java/org/thoughtcrime/securesms/jobs/FastJobStorageTest.kt

@@ -1309,6 +1309,47 @@ class FastJobStorageTest {
     assertThat(subject.getJobCountForFactory("f1")).isEqualTo(2)
   }
 
+  @Test
+  fun `getNextEligibleJob - memory job remains eligible after cache eviction`() {
+    val mostEligibleInMemoryJob = jobSpec(
+      id = "memory-job",
+      factoryKey = "memory-factory",
+      queueKey = "memory-queue",
+      isMemoryOnly = true,
+      globalPriority = Job.Parameters.PRIORITY_HIGH,
+      createTime = 0
+    )
+
+    val subject = FastJobStorage(mockDatabase(emptyList()))
+    subject.init()
+
+    subject.insertJobs(listOf(FullSpec(mostEligibleInMemoryJob, emptyList(), emptyList())))
+    assertThat(subject.getNextEligibleJob(100, NO_PREDICATE)).isEqualTo(mostEligibleInMemoryJob)
+
+    // Trigger cache eviction
+    val lessEligibleJobs = (1..2000).map { i ->
+      FullSpec(
+        jobSpec(
+          id = "job-$i",
+          factoryKey = "factory-$i",
+          queueKey = "queue-$i",
+          isMemoryOnly = false,
+          globalPriority = Job.Parameters.PRIORITY_DEFAULT,
+          createTime = i.toLong()
+        ),
+        emptyList(),
+        emptyList()
+      )
+    }
+    subject.insertJobs(lessEligibleJobs)
+
+    // The memory job should still be the most eligible despite cache eviction
+    val nextJob = subject.getNextEligibleJob(100, NO_PREDICATE)
+    assertThat(nextJob).isNotNull()
+    assertThat(nextJob!!.id).isEqualTo("memory-job")
+    assertThat(nextJob.isMemoryOnly).isEqualTo(true)
+  }
+
   private fun mockDatabase(fullSpecs: List<FullSpec> = emptyList()): JobDatabase {
     val jobs = fullSpecs.map { it.jobSpec }.toMutableList()
     val constraints = fullSpecs.map { it.constraintSpecs }.flatten().toMutableList()
@@ -1320,7 +1361,7 @@ class FastJobStorageTest {
     every { mock.getConstraintSpecs(any()) } returns constraints
     every { mock.getAllDependencySpecs() } returns dependencies
     every { mock.getConstraintSpecsForJobs(any()) } returns constraints
-    every { mock.getJobSpec(any()) } answers { jobs.first { it.id == firstArg() } }
+    every { mock.getJobSpec(any()) } answers { jobs.firstOrNull() { it.id == firstArg() } }
     every { mock.insertJobs(any()) } answers {
       val inserts: List<FullSpec> = firstArg()
       for (insert in inserts) {