Prevent orphaned cohort version sets when a refresh fails mid-flight

core/src/main/kotlin/cohort/CohortStorage.kt

@@ -20,10 +20,20 @@ import java.util.Base64
 import java.util.concurrent.ConcurrentHashMap
 import java.util.zip.GZIPOutputStream
 import kotlin.time.Duration
+import kotlin.time.Duration.Companion.hours
 
 // Constants
 private const val REDIS_SCAN_CHUNK_SIZE: Int = 1000
 
+/**
+ * TTL applied to a cohort version's member set (and diff temp keys) while a refresh is in
+ * flight, so that a refresh which dies mid-way (Redis timeout, pod restart, etc.) cannot
+ * strand the set in Redis forever. Cleared via PERSIST when the version is promoted to
+ * current. Must comfortably exceed the worst-case download + diff duration for the largest
+ * supported cohort.
+ */
+private val PENDING_VERSION_TTL = 24.hours
+
 @Serializable
 internal data class CohortDescription(
     @SerialName("cohortId") val id: String,
@@ -336,6 +346,7 @@ internal class RedisCohortStorage(
                     description.lastModified,
                 )
             private var existingDescription: CohortDescription? = null
+            private var pendingTtlApplied = false
 
             override suspend fun addMembers(members: List<String>) {
                 if (existingDescription == null) {
@@ -348,6 +359,11 @@ internal class RedisCohortStorage(
                 }
                 if (members.isNotEmpty()) {
                     redis.sadd(newCohortKey, members.toSet())
+                    if (!pendingTtlApplied) {
+                        // Self-clean if this refresh dies before the version is promoted in complete()
+                        redis.expire(newCohortKey, PENDING_VERSION_TTL)
+                        pendingTtlApplied = true
+                    }
                 }
             }
 
@@ -382,12 +398,15 @@ internal class RedisCohortStorage(
                                 description.id,
                                 "removed_${System.currentTimeMillis()}",
                             )
-                        val existingBlobKey = RedisKey.CohortBlob(prefix, projectId, description.id, prev.lastModified)
 
                         try {
                             // Server-side set operations - no memory transfer to client!
                             val addedCount = redis.sdiffstore(addedKey, newCohortKey, existingCohortKey)
                             val removedCount = redis.sdiffstore(removedKey, existingCohortKey, newCohortKey)
+                            // Backstop TTLs: if this process dies before the DELs below run,
+                            // the temp keys self-clean instead of persisting forever.
+                            redis.expire(addedKey, PENDING_VERSION_TTL)
+                            redis.expire(removedKey, PENDING_VERSION_TTL)
                             log.info(
                                 "cohort={} diff: addedCount={}, removedCount={}",
                                 description.id,
@@ -409,11 +428,9 @@ internal class RedisCohortStorage(
                                 }
                             }
                         } finally {
-                            // Clean up temporary and existing keys
+                            // Clean up temporary keys
                             redis.del(addedKey)
                             redis.del(removedKey)
-                            redis.expire(existingCohortKey, ttl)
-                            redis.expire(existingBlobKey, ttl)
                         }
                     } else {
                         // No previous cohort: all members are additions
@@ -431,13 +448,33 @@ internal class RedisCohortStorage(
                 val b64 = Base64.getEncoder().encodeToString(gzBytes)
                 redis.set(blobKey, b64)
 
-                // Publish the cohort description only after successful blob store
+                // Promote this version: clear the pending TTL so the now-current member set
+                // does not expire, then publish the cohort description (only after successful
+                // blob store).
+                redis.persist(newCohortKey)
                 val updatedDescription = description.copy(size = finalSize)
                 val jsonEncodedDescription = json.encodeToString(updatedDescription)
                 redis.hset(
                     RedisKey.CohortDescriptions(prefix, projectId),
                     mapOf(description.id to jsonEncodedDescription),
                 )
+
+                // Retire the previous version only after successful promotion. A failed
+                // refresh must leave the previous (still published) version untouched —
+                // expiring it in a finally block would break reads of the live cohort.
+                if (finalSize > 0 && prev != null && prev.lastModified != description.lastModified) {
+                    val previousCohortKey =
+                        RedisKey.CohortMembers(
+                            prefix,
+                            projectId,
+                            prev.id,
+                            prev.groupType,
+                            prev.lastModified,
+                        )
+                    val previousBlobKey = RedisKey.CohortBlob(prefix, projectId, description.id, prev.lastModified)
+                    redis.expire(previousCohortKey, ttl)
+                    redis.expire(previousBlobKey, ttl)
+                }
             }
         }
     }

core/src/main/kotlin/util/redis/Redis.kt

@@ -80,6 +80,8 @@ internal interface Redis {
         ttl: Duration,
     )
 
+    suspend fun persist(key: RedisKey)
+
     suspend fun saddPipeline(
         commands: List<Pair<RedisKey, Set<String>>>,
         batchSize: Int,

core/src/main/kotlin/util/redis/RedisClusterConnection.kt

@@ -239,6 +239,12 @@ internal class RedisClusterConnection(
         }
     }
 
+    override suspend fun persist(key: RedisKey) {
+        connection.run {
+            persist(key.value)
+        }
+    }
+
     override suspend fun saddPipeline(
         commands: List<Pair<RedisKey, Set<String>>>,
         batchSize: Int,

core/src/main/kotlin/util/redis/RedisConnection.kt

@@ -239,6 +239,12 @@ internal class RedisConnection(
         }
     }
 
+    override suspend fun persist(key: RedisKey) {
+        connection.run {
+            persist(key.value)
+        }
+    }
+
     override suspend fun saddPipeline(
         commands: List<Pair<RedisKey, Set<String>>>,
         batchSize: Int,

core/src/test/kotlin/cohort/CohortStorageTest.kt

@@ -18,6 +18,7 @@ import java.util.Base64
 import java.util.zip.GZIPInputStream
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFalse
 import kotlin.test.assertNotNull
 import kotlin.test.assertNull
 import kotlin.test.assertTrue
@@ -239,4 +240,46 @@ class CohortStorageTest {
             assertEquals(3, stored?.size)
             assertEquals(setOf("u1", "u2", "u3"), stored?.members)
         }
+
+    @Test
+    fun `pending version set carries a ttl during ingestion, cleared on promotion`(): Unit =
+        runBlocking {
+            val cohortStorage = RedisCohortStorage("12345", Duration.INFINITE, "amplitude ", redis, redis, 1000, 1000, CohortBlobCache())
+            val cohortV1 = cohort("a", lastModified = 1, members = setOf("1", "2"))
+            val membersKey = RedisKey.CohortMembers("amplitude ", "12345", cohortV1.id, "User", 1)
+
+            val acc = cohortStorage.createWriter(cohortV1.toCohortDescription())
+            acc.addMembers(cohortV1.members.toList())
+            // pending (un-promoted) version set carries a TTL so an aborted refresh self-cleans
+            assertTrue(redis.expirations.containsKey(membersKey.value))
+
+            acc.complete(cohortV1.members.size)
+            // promotion clears the TTL so the current version never expires
+            assertFalse(redis.expirations.containsKey(membersKey.value))
+        }
+
+    @Test
+    fun `unpromoted refresh leaves previous version live and new version self-cleaning`(): Unit =
+        runBlocking {
+            val cohortStorage = RedisCohortStorage("12345", Duration.INFINITE, "amplitude ", redis, redis, 1000, 1000, CohortBlobCache())
+            // v1 promoted
+            val cohortV1 = cohort("a", lastModified = 1, members = setOf("1", "2"))
+            run {
+                val acc = cohortStorage.createWriter(cohortV1.toCohortDescription())
+                acc.addMembers(cohortV1.members.toList())
+                acc.complete(cohortV1.members.size)
+            }
+            // v2 refresh downloads members but never completes (simulates a mid-flight failure)
+            val cohortV2 = cohort("a", lastModified = 2, members = setOf("1", "3"))
+            cohortStorage.createWriter(cohortV2.toCohortDescription()).addMembers(cohortV2.members.toList())
+
+            val v1Key = RedisKey.CohortMembers("amplitude ", "12345", cohortV1.id, "User", 1)
+            val v2Key = RedisKey.CohortMembers("amplitude ", "12345", cohortV2.id, "User", 2)
+            // the published (current) version is untouched by the failed refresh
+            assertFalse(redis.expirations.containsKey(v1Key.value))
+            // the stranded pending version will self-clean via its TTL
+            assertTrue(redis.expirations.containsKey(v2Key.value))
+            // published cohort still reads correctly
+            assertEquals(cohortV1, cohortStorage.getCohort(cohortV1.id))
+        }
 }

core/src/test/kotlin/test/InMemoryRedis.kt

@@ -133,6 +133,10 @@ internal class InMemoryRedis : Redis {
         expirations[key.value] = ttl
     }
 
+    override suspend fun persist(key: RedisKey) {
+        expirations.remove(key.value)
+    }
+
     override suspend fun saddPipeline(
         commands: List<Pair<RedisKey, Set<String>>>,
         batchSize: Int,