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,