From 26f1232ae46a9cbac45cc7fec002b8d8a28e298c Mon Sep 17 00:00:00 2001 From: Adrian Liu Date: Thu, 16 Jul 2026 20:48:20 -0700 Subject: [PATCH] Prevent orphaned cohort version sets on failed refreshes A cohort version's member set is written to Redis with no TTL, and only receives one when a *successor* refresh completes successfully. Any refresh that dies mid-flight (Redis command timeout, pod restart, CohortTooLargeException) therefore strands the full member set in Redis forever. For multi-million-member cohorts this leaks ~0.5 GB per failed refresh, all pinned to the same cluster slot by the cohort hash tag. In one production cluster we found 58 orphaned version sets (~25 GB) for a single 9.8M-member cohort, accumulated over two months of intermittent refresh failures. Fix, in three parts: - Pending version sets self-clean: apply a 24h TTL to the new version's member set as soon as ingestion starts writing it, and PERSIST it when the version is promoted (description published). An aborted refresh now leaves nothing behind once the TTL fires. - Temp diff keys get the same backstop TTL right after SDIFFSTORE, in case the process dies before the finally-block DELs run. - Retire the previous version only after successful promotion. Previously the EXPIRE of the old (still published) version set ran in a finally block, so a refresh that failed mid-update expired the *live* version, breaking cohort reads until the next successful refresh. Co-Authored-By: Claude Fable 5 --- core/src/main/kotlin/cohort/CohortStorage.kt | 47 +++++++++++++++++-- core/src/main/kotlin/util/redis/Redis.kt | 2 + .../util/redis/RedisClusterConnection.kt | 6 +++ .../main/kotlin/util/redis/RedisConnection.kt | 6 +++ .../test/kotlin/cohort/CohortStorageTest.kt | 43 +++++++++++++++++ core/src/test/kotlin/test/InMemoryRedis.kt | 4 ++ 6 files changed, 103 insertions(+), 5 deletions(-) diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt index cc30f0b..20a8804 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/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) { 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) + } } } } diff --git a/core/src/main/kotlin/util/redis/Redis.kt b/core/src/main/kotlin/util/redis/Redis.kt index c70a12c..8ef64d2 100644 --- a/core/src/main/kotlin/util/redis/Redis.kt +++ b/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>>, batchSize: Int, diff --git a/core/src/main/kotlin/util/redis/RedisClusterConnection.kt b/core/src/main/kotlin/util/redis/RedisClusterConnection.kt index c6f521f..c07bd17 100644 --- a/core/src/main/kotlin/util/redis/RedisClusterConnection.kt +++ b/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>>, batchSize: Int, diff --git a/core/src/main/kotlin/util/redis/RedisConnection.kt b/core/src/main/kotlin/util/redis/RedisConnection.kt index 94d11a4..3aec9bb 100644 --- a/core/src/main/kotlin/util/redis/RedisConnection.kt +++ b/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>>, batchSize: Int, diff --git a/core/src/test/kotlin/cohort/CohortStorageTest.kt b/core/src/test/kotlin/cohort/CohortStorageTest.kt index a729540..31ed1f2 100644 --- a/core/src/test/kotlin/cohort/CohortStorageTest.kt +++ b/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)) + } } diff --git a/core/src/test/kotlin/test/InMemoryRedis.kt b/core/src/test/kotlin/test/InMemoryRedis.kt index fba6685..680398a 100644 --- a/core/src/test/kotlin/test/InMemoryRedis.kt +++ b/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>>, batchSize: Int,