From b733a440db18cccb24042c1b544d54cb9a913214 Mon Sep 17 00:00:00 2001 From: Adrian Liu Date: Wed, 15 Jul 2026 23:44:34 -0700 Subject: [PATCH 1/6] Replace SDIFFSTORE cohort diff with streamed client-side diff A set difference over the full old/new member sets executes as a single blocking command on one single-threaded shard - multi-second for multi-million member cohorts - stalling every concurrent command on that shard for its whole duration (replicas included, since they re-execute the replicated command). At 10M members the diff reliably exceeds the default 10s command timeout, failing the refresh outright. Stream both sets via SSCAN in bounded chunks instead and diff client-side, hash-partitioned to cap proxy memory, so the largest single Redis command issued during a cohort update is one SSCAN page regardless of cohort size. Co-Authored-By: Claude Fable 5 --- core/src/main/kotlin/cohort/CohortStorage.kt | 151 ++++++++++++++---- .../test/kotlin/cohort/CohortStorageTest.kt | 41 +++++ 2 files changed, 157 insertions(+), 35 deletions(-) diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt index cc30f0b..2db7e71 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/core/src/main/kotlin/cohort/CohortStorage.kt @@ -24,6 +24,13 @@ import kotlin.time.Duration // Constants private const val REDIS_SCAN_CHUNK_SIZE: Int = 1000 +/** + * Maximum number of existing-set members held in proxy memory at once while diffing two cohort + * versions. Cohorts larger than this are diffed in ceil(size / limit) hash partitions, trading one + * extra SSCAN pass over both member sets per additional partition for bounded memory. + */ +private const val DIFF_PARTITION_MAX_MEMBERS: Int = 2_000_000 + @Serializable internal data class CohortDescription( @SerialName("cohortId") val id: String, @@ -228,6 +235,7 @@ internal class RedisCohortStorage( private val scanLimit: Long, private val pipelineBatchSize: Int, private val cohortBlobCache: CohortBlobCache, + private val diffPartitionMaxMembers: Int = DIFF_PARTITION_MAX_MEMBERS, ) : CohortStorage { companion object { val log by logger() @@ -265,6 +273,112 @@ internal class RedisCohortStorage( } } + /** + * Compute and apply the added/removed membership updates between the existing and new cohort + * member sets without issuing any O(cohort size) Redis command. + * + * SDIFFSTORE was previously used here, but a set difference over the full old/new member sets + * executes as a single blocking command on one (single-threaded) shard — multi-second for + * multi-million member cohorts — stalling every concurrent command on that shard for its whole + * duration, replicas included since they re-execute the replicated command. Instead, both sets + * are streamed via SSCAN in [REDIS_SCAN_CHUNK_SIZE] chunks and diffed client-side, so the + * largest single Redis command issued is one SSCAN page regardless of cohort size. + * + * To bound proxy memory, members are hash-partitioned into ceil(size / [diffPartitionMaxMembers]) + * partitions and diffed one partition at a time: at most one partition of the existing set is + * held in memory at once, at the cost of one extra SSCAN pass over both sets per additional + * partition. + * + * Both keys are immutable at this point (the new set is fully written before [CohortIngestionWriter.complete], + * the existing set is the previously published version), but they are scanned from the primary + * connection to avoid misclassifying members due to replica lag on the just-written new set. + */ + private suspend fun applyMembershipDiff( + existingCohortKey: RedisKey, + newCohortKey: RedisKey, + description: CohortDescription, + existingSize: Int, + newSize: Int, + ): Pair { + val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1 + val cohortIdSet = setOf(description.id) + var addedCount = 0L + var removedCount = 0L + for (partition in 0 until partitions) { + val existingPartition = HashSet() + redis.sscanChunked(existingCohortKey, REDIS_SCAN_CHUNK_SIZE) { chunk -> + for (member in chunk) { + if (partitions == 1 || partitionOf(member, partitions) == partition) { + existingPartition.add(member) + } + } + } + val added = mutableListOf() + redis.sscanChunked(newCohortKey, REDIS_SCAN_CHUNK_SIZE) { chunk -> + for (member in chunk) { + if (partitions > 1 && partitionOf(member, partitions) != partition) { + continue + } + // Members present in both versions are drained from the existing partition + // here; whatever remains after the scan is exactly this partition's removals. + if (!existingPartition.remove(member)) { + added.add(member) + if (added.size >= REDIS_SCAN_CHUNK_SIZE) { + addedCount += + flushMembershipUpdates(added, description.groupType, cohortIdSet) { updates, batchSize -> + redis.saddPipeline(updates, batchSize) + } + } + } + } + } + addedCount += + flushMembershipUpdates(added, description.groupType, cohortIdSet) { updates, batchSize -> + redis.saddPipeline(updates, batchSize) + } + val removed = mutableListOf() + for (member in existingPartition) { + removed.add(member) + if (removed.size >= REDIS_SCAN_CHUNK_SIZE) { + removedCount += + flushMembershipUpdates(removed, description.groupType, cohortIdSet) { updates, batchSize -> + redis.sremPipeline(updates, batchSize) + } + } + } + removedCount += + flushMembershipUpdates(removed, description.groupType, cohortIdSet) { updates, batchSize -> + redis.sremPipeline(updates, batchSize) + } + } + return addedCount to removedCount + } + + private suspend fun flushMembershipUpdates( + members: MutableList, + groupType: String, + cohortIdSet: Set, + pipeline: suspend (updates: List>>, batchSize: Int) -> Unit, + ): Long { + if (members.isEmpty()) return 0L + val updates = + members.map { userId -> + RedisKey.UserCohortMemberships(prefix, projectId, groupType, userId) to cohortIdSet + } + pipeline(updates, pipelineBatchSize) + val flushed = members.size.toLong() + members.clear() + return flushed + } + + private fun partitionOf( + member: String, + partitions: Int, + ): Int { + val hash = member.hashCode() % partitions + return if (hash < 0) hash + partitions else hash + } + override suspend fun getCohort(cohortId: String): Cohort? { val description = getCohortDescription(cohortId) ?: return null val members = getCohortMembers(cohortId, description.groupType, description.lastModified) @@ -367,51 +481,18 @@ internal class RedisCohortStorage( prev.lastModified, ) - // Create temporary keys for differences - server-side operations - val addedKey = - RedisKey.CohortTemporary( - prefix, - projectId, - description.id, - "added_${System.currentTimeMillis()}", - ) - val removedKey = - RedisKey.CohortTemporary( - prefix, - projectId, - 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) + val (addedCount, removedCount) = + applyMembershipDiff(existingCohortKey, newCohortKey, description, prev.size, finalSize) log.info( "cohort={} diff: addedCount={}, removedCount={}", description.id, addedCount, removedCount, ) - - // Process added users in streamed chunks - if (addedCount > 0) { - processMembershipUpdates(addedKey, description.groupType, description.id) { updates, batchSize -> - redis.saddPipeline(updates, batchSize) - } - } - - // Process removed users in streamed chunks - if (removedCount > 0) { - processMembershipUpdates(removedKey, description.groupType, description.id) { updates, batchSize -> - redis.sremPipeline(updates, batchSize) - } - } } finally { - // Clean up temporary and existing keys - redis.del(addedKey) - redis.del(removedKey) redis.expire(existingCohortKey, ttl) redis.expire(existingBlobKey, ttl) } diff --git a/core/src/test/kotlin/cohort/CohortStorageTest.kt b/core/src/test/kotlin/cohort/CohortStorageTest.kt index a729540..a65e24f 100644 --- a/core/src/test/kotlin/cohort/CohortStorageTest.kt +++ b/core/src/test/kotlin/cohort/CohortStorageTest.kt @@ -204,6 +204,47 @@ class CohortStorageTest { } } + @Test + fun `test redis, cohort update diffs client side across partitions`(): Unit = + runBlocking { + // diffPartitionMaxMembers=3 forces multiple hash partitions for a 10-member cohort so + // the partitioned code path is exercised, not just the single-partition fast path. + val cohortStorage = + RedisCohortStorage( + "12345", + Duration.INFINITE, + "amplitude ", + redis, + redis, + 1000, + 1000, + CohortBlobCache(), + diffPartitionMaxMembers = 3, + ) + val v1 = cohort("p", lastModified = 1, members = (1..10).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v1.toCohortDescription()) + acc.addMembers(v1.members.toList()) + acc.complete(v1.members.size) + } + for (member in 1..10) { + assertEquals(setOf("p"), cohortStorage.getCohortMemberships("User", "u$member")) + } + // v2 removes u1-u5 and adds u11-u15 + val v2 = cohort("p", lastModified = 2, members = (6..15).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v2.toCohortDescription()) + acc.addMembers(v2.members.toList()) + acc.complete(v2.members.size) + } + for (removed in 1..5) { + assertEquals(emptySet(), cohortStorage.getCohortMemberships("User", "u$removed")) + } + for (retained in 6..15) { + assertEquals(setOf("p"), cohortStorage.getCohortMemberships("User", "u$retained")) + } + } + @Test fun `test redis, put large cohort, no OutOfMemoryError`(): Unit = runBlocking { From cccfd9f92aa4c213ab2bfb8393b1dfffbc7f8f75 Mon Sep 17 00:00:00 2001 From: Adrian Liu Date: Thu, 16 Jul 2026 20:48:20 -0700 Subject: [PATCH 2/6] 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 | 61 ++++++++++++++----- 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, 106 insertions(+), 16 deletions(-) diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt index 2db7e71..25ee51f 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/core/src/main/kotlin/cohort/CohortStorage.kt @@ -20,6 +20,7 @@ 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 @@ -31,6 +32,15 @@ private const val REDIS_SCAN_CHUNK_SIZE: Int = 1000 */ private const val DIFF_PARTITION_MAX_MEMBERS: Int = 2_000_000 +/** + * 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, @@ -450,6 +460,7 @@ internal class RedisCohortStorage( description.lastModified, ) private var existingDescription: CohortDescription? = null + private var pendingTtlApplied = false override suspend fun addMembers(members: List) { if (existingDescription == null) { @@ -462,6 +473,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 + } } } @@ -481,21 +497,14 @@ internal class RedisCohortStorage( prev.lastModified, ) - val existingBlobKey = RedisKey.CohortBlob(prefix, projectId, description.id, prev.lastModified) - - try { - val (addedCount, removedCount) = - applyMembershipDiff(existingCohortKey, newCohortKey, description, prev.size, finalSize) - log.info( - "cohort={} diff: addedCount={}, removedCount={}", - description.id, - addedCount, - removedCount, - ) - } finally { - redis.expire(existingCohortKey, ttl) - redis.expire(existingBlobKey, ttl) - } + val (addedCount, removedCount) = + applyMembershipDiff(existingCohortKey, newCohortKey, description, prev.size, finalSize) + log.info( + "cohort={} diff: addedCount={}, removedCount={}", + description.id, + addedCount, + removedCount, + ) } else { // No previous cohort: all members are additions processMembershipUpdates(newCohortKey, description.groupType, description.id) { updates, batchSize -> @@ -512,13 +521,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 a65e24f..a727f73 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 @@ -280,4 +281,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, From 147b26021895c43e5022d3ab3c7daa41d576e0fe Mon Sep 17 00:00:00 2001 From: Vaibhav Jain Date: Fri, 17 Jul 2026 11:58:05 -0700 Subject: [PATCH 3/6] Gate streamed cohort diff behind a flag; harden it for very large cohorts The streamed client-side diff is now opt-in via AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED (default false); the server-side SDIFFSTORE diff remains the default path with a command stream identical to before, plus the pending-version backstop TTLs from the previous commit armed immediately after each SDIFFSTORE (not after both, so addedKey is not left unprotected through the second, multi-second blocking command). Two knobs are exposed (and validated > 0 at startup) for tuning at 10M+ member scale, where the 2M default partition size costs one full SSCAN pass over both member sets per partition: AMPLITUDE_COHORT_DIFF_PARTITION_MAX_MEMBERS and AMPLITUDE_COHORT_DIFF_SCAN_CHUNK_SIZE. Hardening for the streamed path, whose diffs can run for minutes on multi-million member cohorts: - Renew the cohort-loading lock every 512 scan pages/flushes so a diff outlasting the lock TTL (900s) doesn't let a second instance start a concurrent load of the same cohort; abort the refresh (retried next sync cycle) if renewal fails rather than continue unlocked. New Redis.renewLock() uses an atomic compare-and-expire mirroring releaseLock(). - Cross-check each scan's raw returned count against SCARD and abort the refresh on a shortfall: SSCAN silently skips members when its cursor is invalidated mid-scan (cluster failover/slot migration), and a truncated scan here becomes false removals of unchanged users that no later diff repairs. A distinct-count check across partitions additionally catches duplicate-masked skips on the existing set. - If the existing members key is missing at diff time, apply all new members as additions with no removals -- SDIFFSTORE semantics for that state -- under the same scan guard, and log an error instead of silently publishing. Also fixed in passing: lock values in RedisConnection and RedisClusterConnection were built with ${'$'} escapes, producing the same literal string on every instance and making releaseLock's compare-and-delete vacuous across pods. Co-Authored-By: Claude Fable 5 --- core/src/main/kotlin/Config.kt | 20 ++ core/src/main/kotlin/cohort/CohortStorage.kt | 247 ++++++++++++++++-- core/src/main/kotlin/util/redis/Redis.kt | 11 + .../util/redis/RedisClusterConnection.kt | 33 ++- .../main/kotlin/util/redis/RedisConnection.kt | 37 ++- .../test/kotlin/cohort/CohortStorageTest.kt | 141 ++++++++++ core/src/test/kotlin/test/InMemoryRedis.kt | 17 ++ 7 files changed, 485 insertions(+), 21 deletions(-) diff --git a/core/src/main/kotlin/Config.kt b/core/src/main/kotlin/Config.kt index 7823f2e..2ffb27e 100644 --- a/core/src/main/kotlin/Config.kt +++ b/core/src/main/kotlin/Config.kt @@ -226,6 +226,9 @@ data class RedisConfiguration( val connectionTimeoutMillis: Long = Default.REDIS_CONNECTION_TIMEOUT_MILLIS, val commandTimeoutMillis: Long = Default.REDIS_COMMAND_TIMEOUT_MILLIS, val pipelineBatchSize: Int = Default.REDIS_PIPELINE_BATCH_SIZE, + val streamedCohortDiffEnabled: Boolean = Default.COHORT_STREAMED_DIFF_ENABLED, + val cohortDiffPartitionMaxMembers: Int = Default.COHORT_DIFF_PARTITION_MAX_MEMBERS, + val cohortDiffScanChunkSize: Int = Default.COHORT_DIFF_SCAN_CHUNK_SIZE, ) { companion object { fun fromEnv(): RedisConfiguration? { @@ -240,6 +243,12 @@ data class RedisConfiguration( val connectionTimeoutMillis = longEnv(EnvKey.REDIS_CONNECTION_TIMEOUT_MILLIS, Default.REDIS_CONNECTION_TIMEOUT_MILLIS)!! val commandTimeoutMillis = longEnv(EnvKey.REDIS_COMMAND_TIMEOUT_MILLIS, Default.REDIS_COMMAND_TIMEOUT_MILLIS)!! val pipelineBatchSize = intEnv(EnvKey.REDIS_PIPELINE_BATCH_SIZE, Default.REDIS_PIPELINE_BATCH_SIZE)!! + val streamedCohortDiffEnabled = + booleanEnv(EnvKey.COHORT_STREAMED_DIFF_ENABLED, Default.COHORT_STREAMED_DIFF_ENABLED) + val cohortDiffPartitionMaxMembers = + intEnv(EnvKey.COHORT_DIFF_PARTITION_MAX_MEMBERS, Default.COHORT_DIFF_PARTITION_MAX_MEMBERS)!! + val cohortDiffScanChunkSize = + intEnv(EnvKey.COHORT_DIFF_SCAN_CHUNK_SIZE, Default.COHORT_DIFF_SCAN_CHUNK_SIZE)!! RedisConfiguration( uri = redisUri, @@ -251,6 +260,9 @@ data class RedisConfiguration( connectionTimeoutMillis = connectionTimeoutMillis, commandTimeoutMillis = commandTimeoutMillis, pipelineBatchSize = pipelineBatchSize, + streamedCohortDiffEnabled = streamedCohortDiffEnabled, + cohortDiffPartitionMaxMembers = cohortDiffPartitionMaxMembers, + cohortDiffScanChunkSize = cohortDiffScanChunkSize, ) } else { null @@ -316,6 +328,10 @@ object EnvKey { const val REDIS_COMMAND_TIMEOUT_MILLIS = "AMPLITUDE_REDIS_COMMAND_TIMEOUT_MILLIS" const val REDIS_PIPELINE_BATCH_SIZE = "AMPLITUDE_REDIS_PIPELINE_BATCH_SIZE" + const val COHORT_STREAMED_DIFF_ENABLED = "AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED" + const val COHORT_DIFF_PARTITION_MAX_MEMBERS = "AMPLITUDE_COHORT_DIFF_PARTITION_MAX_MEMBERS" + const val COHORT_DIFF_SCAN_CHUNK_SIZE = "AMPLITUDE_COHORT_DIFF_SCAN_CHUNK_SIZE" + const val METRICS_PORT = "AMPLITUDE_METRICS_PORT" const val METRICS_PATH = "AMPLITUDE_METRICS_PATH" const val METRICS_LOG_FAILURES = "AMPLITUDE_METRICS_LOG_FAILURES" @@ -357,6 +373,10 @@ object Default { const val REDIS_COMMAND_TIMEOUT_MILLIS = 10000L const val REDIS_PIPELINE_BATCH_SIZE = 100 + const val COHORT_STREAMED_DIFF_ENABLED = false + const val COHORT_DIFF_PARTITION_MAX_MEMBERS = 2_000_000 + const val COHORT_DIFF_SCAN_CHUNK_SIZE = 1000 + const val METRICS_PORT = 9090 const val METRICS_PATH = "metrics" const val METRICS_LOG_FAILURES = false diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt index 25ee51f..e744a3a 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/core/src/main/kotlin/cohort/CohortStorage.kt @@ -41,6 +41,13 @@ private const val DIFF_PARTITION_MAX_MEMBERS: Int = 2_000_000 */ private val PENDING_VERSION_TTL = 24.hours +/** + * How often (in SSCAN pages / pipeline flushes) the streamed diff renews the cohort-loading lock. + * A single partition pass over a multi-million member cohort can run for minutes — longer than the + * lock TTL — so renewing only between passes is not enough to keep a second instance out. + */ +private const val LOCK_RENEWAL_SCAN_PAGES: Int = 512 + @Serializable internal data class CohortDescription( @SerialName("cohortId") val id: String, @@ -134,6 +141,9 @@ internal fun getCohortStorage( redisConfiguration.scanLimit, redisConfiguration.pipelineBatchSize, CohortBlobCache(), + redisConfiguration.streamedCohortDiffEnabled, + redisConfiguration.cohortDiffPartitionMaxMembers, + redisConfiguration.cohortDiffScanChunkSize, ) } else { InMemoryCohortStorage() @@ -245,15 +255,28 @@ internal class RedisCohortStorage( private val scanLimit: Long, private val pipelineBatchSize: Int, private val cohortBlobCache: CohortBlobCache, + private val streamedDiffEnabled: Boolean = false, private val diffPartitionMaxMembers: Int = DIFF_PARTITION_MAX_MEMBERS, + private val diffScanChunkSize: Int = REDIS_SCAN_CHUNK_SIZE, ) : CohortStorage { companion object { val log by logger() } + init { + // A non-positive partition max makes the partition-count arithmetic skip the diff loop + // entirely (publishing a version with zero membership updates); a non-positive chunk + // size is rejected by SSCAN. Fail at startup instead. + require(diffPartitionMaxMembers > 0) { "diffPartitionMaxMembers must be > 0, got $diffPartitionMaxMembers" } + require(diffScanChunkSize > 0) { "diffScanChunkSize must be > 0, got $diffScanChunkSize" } + } + // Track inflight blob loads to avoid duplicate reads private val inflightBlobLoads = ConcurrentHashMap>() + // TTLs the loading locks were acquired with, so long-running diffs can renew them + private val loadingLockTtls = ConcurrentHashMap() + /** * Stream a Redis Set via SSCAN and pipeline membership updates in sub-batches. * The provided [pipeline] function is invoked once per SSCAN chunk with the @@ -291,7 +314,7 @@ internal class RedisCohortStorage( * executes as a single blocking command on one (single-threaded) shard — multi-second for * multi-million member cohorts — stalling every concurrent command on that shard for its whole * duration, replicas included since they re-execute the replicated command. Instead, both sets - * are streamed via SSCAN in [REDIS_SCAN_CHUNK_SIZE] chunks and diffed client-side, so the + * are streamed via SSCAN in [diffScanChunkSize] chunks and diffed client-side, so the * largest single Redis command issued is one SSCAN page regardless of cohort size. * * To bound proxy memory, members are hash-partitioned into ceil(size / [diffPartitionMaxMembers]) @@ -302,6 +325,15 @@ internal class RedisCohortStorage( * Both keys are immutable at this point (the new set is fully written before [CohortIngestionWriter.complete], * the existing set is the previously published version), but they are scanned from the primary * connection to avoid misclassifying members due to replica lag on the just-written new set. + * + * A missed member is not a transient error here: one skipped by the existing-set scan becomes a + * false addition, and one skipped by the new-set scan becomes a false removal of an unchanged + * user — neither is ever repaired by a later diff. SSCAN can silently skip members when the + * cursor is invalidated mid-scan (cluster failover or slot migration), so each scan cross-checks + * the raw returned count against SCARD and aborts the refresh (to be retried) on a shortfall. + * The loading lock is renewed every [LOCK_RENEWAL_SCAN_PAGES] pages so that a diff outlasting + * the lock TTL does not let a second instance start a concurrent load; if renewal fails the + * diff aborts rather than run unlocked. */ private suspend fun applyMembershipDiff( existingCohortKey: RedisKey, @@ -309,22 +341,40 @@ internal class RedisCohortStorage( description: CohortDescription, existingSize: Int, newSize: Int, - ): Pair { - val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1 + ) { + val existingCard = redis.scard(existingCohortKey) + val newCard = redis.scard(newCohortKey) val cohortIdSet = setOf(description.id) + if (existingCard == 0L) { + applyAllAdditions(newCohortKey, description, existingSize, newCard, cohortIdSet) + return + } + val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1 var addedCount = 0L var removedCount = 0L + var existingDistinct = 0L + var scanPages = 0 for (partition in 0 until partitions) { - val existingPartition = HashSet() - redis.sscanChunked(existingCohortKey, REDIS_SCAN_CHUNK_SIZE) { chunk -> + renewLoadingLock(description.id) + val expectedPartitionMembers = (existingCard / partitions).toInt() + 1 + val existingPartition = HashSet((expectedPartitionMembers / 0.75f).toInt() + 1) + var existingScanned = 0L + redis.sscanChunked(existingCohortKey, diffScanChunkSize) { chunk -> + existingScanned += chunk.size + if (++scanPages % LOCK_RENEWAL_SCAN_PAGES == 0) renewLoadingLock(description.id) for (member in chunk) { if (partitions == 1 || partitionOf(member, partitions) == partition) { existingPartition.add(member) } } } + checkScanComplete(description.id, existingCohortKey, existingScanned, existingCard) + existingDistinct += existingPartition.size val added = mutableListOf() - redis.sscanChunked(newCohortKey, REDIS_SCAN_CHUNK_SIZE) { chunk -> + var newScanned = 0L + redis.sscanChunked(newCohortKey, diffScanChunkSize) { chunk -> + newScanned += chunk.size + if (++scanPages % LOCK_RENEWAL_SCAN_PAGES == 0) renewLoadingLock(description.id) for (member in chunk) { if (partitions > 1 && partitionOf(member, partitions) != partition) { continue @@ -333,7 +383,7 @@ internal class RedisCohortStorage( // here; whatever remains after the scan is exactly this partition's removals. if (!existingPartition.remove(member)) { added.add(member) - if (added.size >= REDIS_SCAN_CHUNK_SIZE) { + if (added.size >= diffScanChunkSize) { addedCount += flushMembershipUpdates(added, description.groupType, cohortIdSet) { updates, batchSize -> redis.saddPipeline(updates, batchSize) @@ -342,6 +392,7 @@ internal class RedisCohortStorage( } } } + checkScanComplete(description.id, newCohortKey, newScanned, newCard) addedCount += flushMembershipUpdates(added, description.groupType, cohortIdSet) { updates, batchSize -> redis.saddPipeline(updates, batchSize) @@ -349,11 +400,12 @@ internal class RedisCohortStorage( val removed = mutableListOf() for (member in existingPartition) { removed.add(member) - if (removed.size >= REDIS_SCAN_CHUNK_SIZE) { + if (removed.size >= diffScanChunkSize) { removedCount += flushMembershipUpdates(removed, description.groupType, cohortIdSet) { updates, batchSize -> redis.sremPipeline(updates, batchSize) } + if (++scanPages % LOCK_RENEWAL_SCAN_PAGES == 0) renewLoadingLock(description.id) } } removedCount += @@ -361,7 +413,101 @@ internal class RedisCohortStorage( redis.sremPipeline(updates, batchSize) } } - return addedCount to removedCount + // Duplicates in an SSCAN pass can mask skipped members from the raw-count check above. + // Each existing member hashes to exactly one partition, so the distinct members collected + // across all passes must equal the cardinality; a shortfall means a member was skipped and + // would have been misclassified. + checkScanComplete(description.id, existingCohortKey, existingDistinct, existingCard) + log.info( + "cohort={} diff: addedCount={}, removedCount={}", + description.id, + addedCount, + removedCount, + ) + } + + /** + * Degraded-mode update for when the previous version's members key is gone even though its + * description is still published. SDIFFSTORE semantics for this state were "every new member + * is an addition, nothing is removed" — match that rather than failing the refresh forever, + * but loudly: members dropped between the two versions keep a stale membership until a later + * update removes them. + */ + private suspend fun applyAllAdditions( + newCohortKey: RedisKey, + description: CohortDescription, + existingSize: Int, + newCard: Long, + cohortIdSet: Set, + ) { + if (existingSize > 0) { + log.error( + "cohort={} existing members key missing at diff time (expected size {}); " + + "applying all new members as additions with no removals", + description.id, + existingSize, + ) + } + renewLoadingLock(description.id) + val added = mutableListOf() + var addedCount = 0L + var newScanned = 0L + var scanPages = 0 + redis.sscanChunked(newCohortKey, diffScanChunkSize) { chunk -> + newScanned += chunk.size + if (++scanPages % LOCK_RENEWAL_SCAN_PAGES == 0) renewLoadingLock(description.id) + for (member in chunk) { + added.add(member) + if (added.size >= diffScanChunkSize) { + addedCount += + flushMembershipUpdates(added, description.groupType, cohortIdSet) { updates, batchSize -> + redis.saddPipeline(updates, batchSize) + } + } + } + } + checkScanComplete(description.id, newCohortKey, newScanned, newCard) + addedCount += + flushMembershipUpdates(added, description.groupType, cohortIdSet) { updates, batchSize -> + redis.saddPipeline(updates, batchSize) + } + log.info( + "cohort={} diff: addedCount={}, removedCount={}", + description.id, + addedCount, + 0L, + ) + } + + /** + * SSCAN returns every member of an unmodified set at least once, so the raw returned count + * (duplicates included) must be >= the set's cardinality. A shortfall means the scan was + * silently truncated (e.g. the cursor was invalidated by a cluster failover or slot + * migration mid-scan); diffing from a truncated scan would corrupt per-user memberships. + */ + private fun checkScanComplete( + cohortId: String, + key: RedisKey, + scanned: Long, + cardinality: Long, + ) { + check(scanned >= cardinality) { + "cohort=$cohortId SSCAN of ${key.value} returned $scanned members but SCARD reported " + + "$cardinality; aborting diff to retry (scan was silently truncated)" + } + } + + private suspend fun renewLoadingLock(cohortId: String) { + // Only renew when this instance acquired the lock (tests call complete() directly) + val ttlSeconds = loadingLockTtls[cohortId] ?: return + val lockKey = RedisKey.CohortLoadingLock(prefix, projectId, cohortId) + // Fail closed: continuing without the lock lets a second instance run a concurrent + // diff of the same cohort, interleaving membership writes. Aborting leaves the + // description unpublished and the refresh retries next sync cycle. + check(redis.renewLock(lockKey, ttlSeconds)) { + "cohort=$cohortId loading-lock renewal failed (expired or taken over mid-diff); " + + "aborting refresh to avoid interleaving with a concurrent load" + } } private suspend fun flushMembershipUpdates( @@ -389,6 +535,71 @@ internal class RedisCohortStorage( return if (hash < 0) hash + partitions else hash } + /** + * Compute and apply the added/removed membership updates with two server-side SDIFFSTORE + * commands into temporary keys. Atomic and delta-sized on the apply side, but a set + * difference over multi-million member sets blocks its (single-threaded) shard for the + * whole computation and can exceed the Redis command timeout; enable the streamed + * client-side diff for cohorts at that scale. + */ + private suspend fun applySdiffstoreDiff( + existingCohortKey: RedisKey, + newCohortKey: RedisKey, + description: CohortDescription, + ) { + // Create temporary keys for differences - server-side operations + val addedKey = + RedisKey.CohortTemporary( + prefix, + projectId, + description.id, + "added_${System.currentTimeMillis()}", + ) + val removedKey = + RedisKey.CohortTemporary( + prefix, + projectId, + description.id, + "removed_${System.currentTimeMillis()}", + ) + try { + // Server-side set operations - no memory transfer to client! + // Backstop TTLs are armed immediately after each SDIFFSTORE (not after both): + // the second SDIFFSTORE is itself a multi-second blocking command on large + // cohorts, so addedKey would otherwise sit unprotected through the likeliest + // crash window. If this process dies before the DELs below run, the temp keys + // self-clean instead of persisting forever. + val addedCount = redis.sdiffstore(addedKey, newCohortKey, existingCohortKey) + redis.expire(addedKey, PENDING_VERSION_TTL) + val removedCount = redis.sdiffstore(removedKey, existingCohortKey, newCohortKey) + redis.expire(removedKey, PENDING_VERSION_TTL) + log.info( + "cohort={} diff: addedCount={}, removedCount={}", + description.id, + addedCount, + removedCount, + ) + + // Process added users in streamed chunks + if (addedCount > 0) { + processMembershipUpdates(addedKey, description.groupType, description.id) { updates, batchSize -> + redis.saddPipeline(updates, batchSize) + } + } + + // Process removed users in streamed chunks + if (removedCount > 0) { + processMembershipUpdates(removedKey, description.groupType, description.id) { updates, batchSize -> + redis.sremPipeline(updates, batchSize) + } + } + } finally { + // Clean up temporary keys + redis.del(addedKey) + redis.del(removedKey) + } + } + override suspend fun getCohort(cohortId: String): Cohort? { val description = getCohortDescription(cohortId) ?: return null val members = getCohortMembers(cohortId, description.groupType, description.lastModified) @@ -497,14 +708,11 @@ internal class RedisCohortStorage( prev.lastModified, ) - val (addedCount, removedCount) = + if (streamedDiffEnabled) { applyMembershipDiff(existingCohortKey, newCohortKey, description, prev.size, finalSize) - log.info( - "cohort={} diff: addedCount={}, removedCount={}", - description.id, - addedCount, - removedCount, - ) + } else { + applySdiffstoreDiff(existingCohortKey, newCohortKey, description) + } } else { // No previous cohort: all members are additions processMembershipUpdates(newCohortKey, description.groupType, description.id) { updates, batchSize -> @@ -597,10 +805,15 @@ internal class RedisCohortStorage( ): Boolean { val lockKey = RedisKey.CohortLoadingLock(prefix, projectId, cohortId) log.debug("Acquiring lock for cohort $cohortId") - return redis.acquireLock(lockKey, lockTimeoutSeconds.toLong()) + val acquired = redis.acquireLock(lockKey, lockTimeoutSeconds.toLong()) + if (acquired) { + loadingLockTtls[cohortId] = lockTimeoutSeconds.toLong() + } + return acquired } override suspend fun releaseCohortLoadingLock(cohortId: String) { + loadingLockTtls.remove(cohortId) val lockKey = RedisKey.CohortLoadingLock(prefix, projectId, cohortId) val released = redis.releaseLock(lockKey) if (!released) { diff --git a/core/src/main/kotlin/util/redis/Redis.kt b/core/src/main/kotlin/util/redis/Redis.kt index 8ef64d2..205d512 100644 --- a/core/src/main/kotlin/util/redis/Redis.kt +++ b/core/src/main/kotlin/util/redis/Redis.kt @@ -36,6 +36,8 @@ internal interface Redis { suspend fun smembers(key: RedisKey): Set + suspend fun scard(key: RedisKey): Long + suspend fun sismember( key: RedisKey, value: String, @@ -107,4 +109,13 @@ internal interface Redis { * Only releases locks that were acquired by this Redis instance. */ suspend fun releaseLock(key: RedisKey): Boolean + + /** + * Extend the TTL of a lock previously acquired via [acquireLock] on this instance. + * Returns false if the lock is no longer held by this instance (expired or taken over). + */ + suspend fun renewLock( + key: RedisKey, + ttlSeconds: Long, + ): Boolean } diff --git a/core/src/main/kotlin/util/redis/RedisClusterConnection.kt b/core/src/main/kotlin/util/redis/RedisClusterConnection.kt index c07bd17..8a7ff7a 100644 --- a/core/src/main/kotlin/util/redis/RedisClusterConnection.kt +++ b/core/src/main/kotlin/util/redis/RedisClusterConnection.kt @@ -152,6 +152,10 @@ internal class RedisClusterConnection( return connection.run { smembers(key.value) } } + override suspend fun scard(key: RedisKey): Long { + return connection.run { scard(key.value) } + } + override suspend fun sismember( key: RedisKey, value: String, @@ -282,7 +286,7 @@ internal class RedisClusterConnection( ttlSeconds: Long, ): Boolean { // Generate unique lock value internally - val lockValue = "${'$'}{System.currentTimeMillis()}-${'$'}{Thread.currentThread().id}-${'$'}{java.util.UUID.randomUUID()}" + val lockValue = "${System.currentTimeMillis()}-${Thread.currentThread().id}-${java.util.UUID.randomUUID()}" val result = connection.run { @@ -330,6 +334,33 @@ internal class RedisClusterConnection( return released } + override suspend fun renewLock( + key: RedisKey, + ttlSeconds: Long, + ): Boolean { + val lockValue = + synchronized(activeLocks) { + activeLocks[key.value] + } ?: return false + + // Use Lua script for atomic compare-and-expire + val script = + """ + if redis.call("get", KEYS[1]) == ARGV[1] then + return redis.call("expire", KEYS[1], ARGV[2]) + else + return 0 + end + """.trimIndent() + + val result = + connection.run { + eval(script, io.lettuce.core.ScriptOutputType.INTEGER, arrayOf(key.value), lockValue, ttlSeconds.toString()) + } + + return result == 1L + } + suspend fun pipeline(block: suspend RedisAdvancedClusterAsyncCommands.() -> List>) { withContext(pipelineContext) { val conn = pipelineConnection.await() diff --git a/core/src/main/kotlin/util/redis/RedisConnection.kt b/core/src/main/kotlin/util/redis/RedisConnection.kt index 3aec9bb..0b528c9 100644 --- a/core/src/main/kotlin/util/redis/RedisConnection.kt +++ b/core/src/main/kotlin/util/redis/RedisConnection.kt @@ -152,6 +152,10 @@ internal class RedisConnection( return connection.run { smembers(key.value) } } + override suspend fun scard(key: RedisKey): Long { + return connection.run { scard(key.value) } + } + override suspend fun sismember( key: RedisKey, value: String, @@ -282,7 +286,7 @@ internal class RedisConnection( ttlSeconds: Long, ): Boolean { // Generate unique lock value internally - val lockValue = "${'$'}{System.currentTimeMillis()}-${'$'}{Thread.currentThread().id}-${'$'}{java.util.UUID.randomUUID()}" + val lockValue = "${System.currentTimeMillis()}-${Thread.currentThread().id}-${java.util.UUID.randomUUID()}" val result = connection.run { @@ -323,15 +327,42 @@ internal class RedisConnection( val released = result == 1L if (!released) { - log.warn("Failed to release lock for key ${'$'}{key.value} - lock may have expired or been taken by another process") + log.warn("Failed to release lock for key ${key.value} - lock may have expired or been taken by another process") } released } else { - log.warn("Attempted to release lock for key ${'$'}{key.value} but no active lock found") + log.warn("Attempted to release lock for key ${key.value} but no active lock found") false } } + override suspend fun renewLock( + key: RedisKey, + ttlSeconds: Long, + ): Boolean { + val lockValue = + synchronized(activeLocks) { + activeLocks[key.value] + } ?: return false + + // Use Lua script for atomic compare-and-expire + val luaScript = + """ + if redis.call("GET", KEYS[1]) == ARGV[1] then + return redis.call("EXPIRE", KEYS[1], ARGV[2]) + else + return 0 + end + """.trimIndent() + + val result = + connection.run { + eval(luaScript, io.lettuce.core.ScriptOutputType.INTEGER, arrayOf(key.value), lockValue, ttlSeconds.toString()) + } + + return result == 1L + } + private suspend fun pipeline(block: RedisAsyncCommands.() -> List>) { withContext(pipelineContext) { val connection = pipelineConnection.await() diff --git a/core/src/test/kotlin/cohort/CohortStorageTest.kt b/core/src/test/kotlin/cohort/CohortStorageTest.kt index a727f73..638bc4c 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.assertFailsWith import kotlin.test.assertFalse import kotlin.test.assertNotNull import kotlin.test.assertNull @@ -220,6 +221,7 @@ class CohortStorageTest { 1000, 1000, CohortBlobCache(), + streamedDiffEnabled = true, diffPartitionMaxMembers = 3, ) val v1 = cohort("p", lastModified = 1, members = (1..10).map { "u$it" }.toSet()) @@ -246,6 +248,145 @@ class CohortStorageTest { } } + @Test + fun `test redis, cohort update uses sdiffstore path by default`(): Unit = + runBlocking { + // Same update scenario as the streamed-diff test, but with the flag left at its + // default (off) so the server-side SDIFFSTORE path is exercised. + val cohortStorage = + RedisCohortStorage( + "12345", + Duration.INFINITE, + "amplitude ", + redis, + redis, + 1000, + 1000, + CohortBlobCache(), + ) + val v1 = cohort("p", lastModified = 1, members = (1..10).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v1.toCohortDescription()) + acc.addMembers(v1.members.toList()) + acc.complete(v1.members.size) + } + val v2 = cohort("p", lastModified = 2, members = (6..15).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v2.toCohortDescription()) + acc.addMembers(v2.members.toList()) + acc.complete(v2.members.size) + } + for (removed in 1..5) { + assertEquals(emptySet(), cohortStorage.getCohortMemberships("User", "u$removed")) + } + for (retained in 6..15) { + assertEquals(setOf("p"), cohortStorage.getCohortMemberships("User", "u$retained")) + } + } + + @Test + fun `test redis, streamed diff flushes adds and removals above chunk threshold`(): Unit = + runBlocking { + // >1000 additions and >1000 removals in a single partition so the incremental + // flush branches (added/removed buffers reaching diffScanChunkSize) are exercised. + val cohortStorage = + RedisCohortStorage( + "12345", + Duration.INFINITE, + "amplitude ", + redis, + redis, + 1000, + 1000, + CohortBlobCache(), + streamedDiffEnabled = true, + ) + val v1 = cohort("p", lastModified = 1, members = (1..2500).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v1.toCohortDescription()) + acc.addMembers(v1.members.toList()) + acc.complete(v1.members.size) + } + // v2 removes u1-u1300 and adds u2501-u3800 + val v2 = cohort("p", lastModified = 2, members = (1301..3800).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v2.toCohortDescription()) + acc.addMembers(v2.members.toList()) + acc.complete(v2.members.size) + } + for (removed in 1..1300) { + assertEquals(emptySet(), cohortStorage.getCohortMemberships("User", "u$removed")) + } + for (retained in 1301..3800) { + assertEquals(setOf("p"), cohortStorage.getCohortMemberships("User", "u$retained")) + } + } + + @Test + fun `test redis, rejects non-positive diff tuning values`() { + assertFailsWith { + RedisCohortStorage( + "12345", Duration.INFINITE, "amplitude ", redis, redis, 1000, 1000, CohortBlobCache(), + diffPartitionMaxMembers = 0, + ) + } + assertFailsWith { + RedisCohortStorage( + "12345", Duration.INFINITE, "amplitude ", redis, redis, 1000, 1000, CohortBlobCache(), + diffPartitionMaxMembers = -2_000_000, + ) + } + assertFailsWith { + RedisCohortStorage( + "12345", Duration.INFINITE, "amplitude ", redis, redis, 1000, 1000, CohortBlobCache(), + diffScanChunkSize = 0, + ) + } + } + + @Test + fun `test redis, streamed diff falls back to all additions when existing members key is missing`(): Unit = + runBlocking { + val cohortStorage = + RedisCohortStorage( + "12345", + Duration.INFINITE, + "amplitude ", + redis, + redis, + 1000, + 1000, + CohortBlobCache(), + streamedDiffEnabled = true, + ) + val v1 = cohort("p", lastModified = 1, members = (1..5).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v1.toCohortDescription()) + acc.addMembers(v1.members.toList()) + acc.complete(v1.members.size) + } + // Simulate the published version's members key expiring/being deleted out from + // under the diff. + redis.del(RedisKey.CohortMembers("amplitude ", "12345", "p", "User", 1)) + val v2 = cohort("p", lastModified = 2, members = (3..7).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v2.toCohortDescription()) + acc.addMembers(v2.members.toList()) + acc.complete(v2.members.size) + } + // The refresh must still publish the new version... + assertEquals(2L, cohortStorage.getCohortDescription("p")?.lastModified) + // ...with every new member present. + for (added in 3..7) { + assertEquals(setOf("p"), cohortStorage.getCohortMemberships("User", "u$added")) + } + // Members dropped between versions keep a stale membership in this degraded state + // (matching SDIFFSTORE semantics for a missing existing key). + for (stale in 1..2) { + assertEquals(setOf("p"), cohortStorage.getCohortMemberships("User", "u$stale")) + } + } + @Test fun `test redis, put large cohort, no OutOfMemoryError`(): Unit = runBlocking { diff --git a/core/src/test/kotlin/test/InMemoryRedis.kt b/core/src/test/kotlin/test/InMemoryRedis.kt index 680398a..d50a0bb 100644 --- a/core/src/test/kotlin/test/InMemoryRedis.kt +++ b/core/src/test/kotlin/test/InMemoryRedis.kt @@ -36,6 +36,9 @@ internal class InMemoryRedis : Redis { override suspend fun del(key: RedisKey) { kv.remove(key.value) + sets.remove(key.value) + hashes.remove(key.value) + expirations.remove(key.value) } override suspend fun sadd( @@ -63,6 +66,10 @@ internal class InMemoryRedis : Redis { return sets[key.value] ?: emptySet() } + override suspend fun scard(key: RedisKey): Long { + return (sets[key.value]?.size ?: 0).toLong() + } + override suspend fun sismember( key: RedisKey, value: String, @@ -189,4 +196,14 @@ internal class InMemoryRedis : Redis { false // No active lock to release } } + + override suspend fun renewLock( + key: RedisKey, + ttlSeconds: Long, + ): Boolean { + val keyStr = key.value + val expectedValue = activeLocks[keyStr] ?: return false + // Note: InMemory doesn't implement TTL expiration for simplicity + return kv[keyStr] == expectedValue + } } From 5032aeab2f7cae5f5d7785aae42a8a54c9db9f61 Mon Sep 17 00:00:00 2001 From: Vaibhav Jain Date: Tue, 21 Jul 2026 12:07:42 -0700 Subject: [PATCH 4/6] docs: applyMembershipDiff is the opt-in alternative, not a replacement The kdoc predated the flag gating and still described the streamed diff as having replaced SDIFFSTORE; with the flag off (the default) the SDIFFSTORE path is what runs. Co-Authored-By: Claude Fable 5 --- core/src/main/kotlin/cohort/CohortStorage.kt | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt index e744a3a..c3af6ee 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/core/src/main/kotlin/cohort/CohortStorage.kt @@ -310,12 +310,13 @@ internal class RedisCohortStorage( * Compute and apply the added/removed membership updates between the existing and new cohort * member sets without issuing any O(cohort size) Redis command. * - * SDIFFSTORE was previously used here, but a set difference over the full old/new member sets - * executes as a single blocking command on one (single-threaded) shard — multi-second for - * multi-million member cohorts — stalling every concurrent command on that shard for its whole - * duration, replicas included since they re-execute the replicated command. Instead, both sets - * are streamed via SSCAN in [diffScanChunkSize] chunks and diffed client-side, so the - * largest single Redis command issued is one SSCAN page regardless of cohort size. + * Opt-in alternative to [applySdiffstoreDiff] (the default): a SDIFFSTORE over the full + * old/new member sets executes as a single blocking command on one (single-threaded) shard — + * multi-second for multi-million member cohorts — stalling every concurrent command on that + * shard for its whole duration, replicas included since they re-execute the replicated + * command. Here both sets are instead streamed via SSCAN in [diffScanChunkSize] chunks and + * diffed client-side, so the largest single Redis command issued is one SSCAN page regardless + * of cohort size. * * To bound proxy memory, members are hash-partitioned into ceil(size / [diffPartitionMaxMembers]) * partitions and diffed one partition at a time: at most one partition of the existing set is From f09c9f97b065f921392fe4562bb420f23bcfe530 Mon Sep 17 00:00:00 2001 From: Vaibhav Jain Date: Tue, 21 Jul 2026 12:12:31 -0700 Subject: [PATCH 5/6] Publish before persist, and self-heal a stray TTL on the live version Cursor bugbot on PR #39 flagged that persist(newCohortKey) ran before the description hset: a refresh dying between the two left a persisted (no-TTL) member set with the old description still published -- the pending TTL's self-clean guarantee was defeated exactly one command before it mattered. (In practice the next retry re-arms the TTL on the same key, so the leak only stuck when the upstream version was superseded first -- but the window was real.) Simply swapping the order would trade that leak for something worse: a crash after hset but before persist would leave the LIVE version with a 24h fuse and no code path to clear it (the next sync short-circuits on lastModified), breaking reads a day later. So this does both halves: - complete() now publishes first and persists immediately after, so a refresh that dies anywhere before publish always leaves a set that self-cleans via its pending TTL. - CohortLoader clears any TTL found on the published version's member set on every not-modified sync cycle (new O(1) CohortStorage.ensureCurrentVersionPersisted), repairing the one-command publish->persist crash window within one sync interval, instead of never. This also defuses the stray-fuse race where a lock-breaching concurrent writer re-arms the TTL on an already-promoted set. New tests: a fault-injected hset failure proves the pending TTL survives a crash at publish; ensureCurrentVersionPersisted clears a stray TTL on the published version. Co-Authored-By: Claude Fable 5 --- core/src/main/kotlin/cohort/CohortLoader.kt | 6 ++ core/src/main/kotlin/cohort/CohortStorage.kt | 35 +++++++++-- .../test/kotlin/cohort/CohortStorageTest.kt | 58 +++++++++++++++++++ 3 files changed, 95 insertions(+), 4 deletions(-) diff --git a/core/src/main/kotlin/cohort/CohortLoader.kt b/core/src/main/kotlin/cohort/CohortLoader.kt index 94f0c01..cc089e0 100644 --- a/core/src/main/kotlin/cohort/CohortLoader.kt +++ b/core/src/main/kotlin/cohort/CohortLoader.kt @@ -47,6 +47,12 @@ internal class CohortLoader( } } else { log.debug("loadCohort: cohort not modified - cohortId={}", cohortId) + // Repair path: a previous refresh that died between publishing the + // description and clearing the pending ingestion TTL left the live + // member set with an expiry armed; clear it here (O(1) no-op otherwise). + if (storageCohort != null) { + cohortStorage.ensureCurrentVersionPersisted(storageCohort) + } } } catch (t: Throwable) { // Don't throw if we fail to download the cohort. We diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt index c3af6ee..88a6864 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/core/src/main/kotlin/cohort/CohortStorage.kt @@ -116,6 +116,14 @@ internal interface CohortStorage { * Release the distributed lock for cohort loading. */ suspend fun releaseCohortLoadingLock(cohortId: String) + + /** + * Ensure the published (current) version's member set has no expiry armed. A refresh that + * dies between publishing the description and clearing the pending ingestion TTL leaves the + * live set with a fuse; calling this on every not-modified sync cycle repairs that within + * one cycle. O(1), safe to call unconditionally. + */ + suspend fun ensureCurrentVersionPersisted(description: CohortDescription) } // Streaming ingestion writer contract @@ -244,6 +252,10 @@ internal class InMemoryCohortStorage : CohortStorage { override suspend fun releaseCohortLoadingLock(cohortId: String) { // No-op for in-memory storage } + + override suspend fun ensureCurrentVersionPersisted(description: CohortDescription) { + // No-op for in-memory storage: nothing expires + } } internal class RedisCohortStorage( @@ -730,16 +742,19 @@ internal class RedisCohortStorage( val b64 = Base64.getEncoder().encodeToString(gzBytes) redis.set(blobKey, b64) - // 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) + // Promote this version: publish the cohort description (only after successful + // blob store), then clear the pending TTL so the now-current member set does + // not expire. Publish-then-persist means a refresh that dies anywhere before + // the publish always leaves a set that self-cleans via its pending TTL; the + // one-command window where the published set still carries the TTL is repaired + // by [ensureCurrentVersionPersisted] on the next sync cycle. val updatedDescription = description.copy(size = finalSize) val jsonEncodedDescription = json.encodeToString(updatedDescription) redis.hset( RedisKey.CohortDescriptions(prefix, projectId), mapOf(description.id to jsonEncodedDescription), ) + redis.persist(newCohortKey) // Retire the previous version only after successful promotion. A failed // refresh must leave the previous (still published) version untouched — @@ -813,6 +828,18 @@ internal class RedisCohortStorage( return acquired } + override suspend fun ensureCurrentVersionPersisted(description: CohortDescription) { + redis.persist( + RedisKey.CohortMembers( + prefix, + projectId, + description.id, + description.groupType, + description.lastModified, + ), + ) + } + override suspend fun releaseCohortLoadingLock(cohortId: String) { loadingLockTtls.remove(cohortId) val lockKey = RedisKey.CohortLoadingLock(prefix, projectId, cohortId) diff --git a/core/src/test/kotlin/cohort/CohortStorageTest.kt b/core/src/test/kotlin/cohort/CohortStorageTest.kt index 638bc4c..d291734 100644 --- a/core/src/test/kotlin/cohort/CohortStorageTest.kt +++ b/core/src/test/kotlin/cohort/CohortStorageTest.kt @@ -9,6 +9,7 @@ import com.amplitude.cohort.InMemoryCohortStorage import com.amplitude.cohort.RedisCohortStorage import com.amplitude.cohort.toCohortDescription import com.amplitude.util.json +import com.amplitude.util.redis.Redis import com.amplitude.util.redis.RedisKey import kotlinx.coroutines.runBlocking import test.InMemoryRedis @@ -464,4 +465,61 @@ class CohortStorageTest { // published cohort still reads correctly assertEquals(cohortV1, cohortStorage.getCohort(cohortV1.id)) } + + @Test + fun `refresh that fails at publish keeps the pending ttl so the set self-cleans`(): Unit = + runBlocking { + // Fails the description hset on demand to simulate a refresh dying at the + // publish step — the pending version must keep its self-clean TTL. + var failHset = false + val faultyRedis = + object : Redis by redis { + override suspend fun hset( + key: RedisKey, + values: Map, + ) { + if (failHset) throw RuntimeException("injected hset failure") + redis.hset(key, values) + } + } + val cohortStorage = + RedisCohortStorage("12345", Duration.INFINITE, "amplitude ", faultyRedis, redis, 1000, 1000, CohortBlobCache()) + 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) + } + val cohortV2 = cohort("a", lastModified = 2, members = setOf("1", "3")) + failHset = true + val acc = cohortStorage.createWriter(cohortV2.toCohortDescription()) + acc.addMembers(cohortV2.members.toList()) + assertFailsWith { acc.complete(cohortV2.members.size) } + failHset = false + + val v2Key = RedisKey.CohortMembers("amplitude ", "12345", cohortV2.id, "User", 2) + // publish never happened, so the pending TTL must still be armed (self-clean) + assertTrue(redis.expirations.containsKey(v2Key.value)) + // and the published description still points at v1 + assertEquals(1L, cohortStorage.getCohortDescription("a")?.lastModified) + } + + @Test + fun `ensureCurrentVersionPersisted clears a stray ttl on the published version`(): Unit = + runBlocking { + val cohortStorage = RedisCohortStorage("12345", Duration.INFINITE, "amplitude ", redis, redis, 1000, 1000, CohortBlobCache()) + 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) + } + val v1Key = RedisKey.CohortMembers("amplitude ", "12345", cohortV1.id, "User", 1) + // simulate a refresh that died between publish and persist: live set with a fuse + redis.expire(v1Key, Duration.INFINITE) + assertTrue(redis.expirations.containsKey(v1Key.value)) + + cohortStorage.ensureCurrentVersionPersisted(cohortV1.toCohortDescription()) + assertFalse(redis.expirations.containsKey(v1Key.value)) + } } From 7632c5597f6ded383a99adfefee8a5579a8417f2 Mon Sep 17 00:00:00 2001 From: Vaibhav Jain Date: Tue, 21 Jul 2026 12:52:23 -0700 Subject: [PATCH 6/6] Derive diff partition count from SCARD, not description sizes Cursor bugbot on #40: partition count used the description-reported sizes (prev.size / finalSize) while the scan guards two lines above already fetch the authoritative SCARD cardinalities. A crashed ingest can leave a version key holding more members than its published size; partitioning by the smaller number lets a single in-memory partition exceed diffPartitionMaxMembers -- the bound partitioning exists to enforce. Partition count now comes from max(existingCard, newCard). The newSize parameter became unused and is removed. New test publishes a version whose description size understates the key's cardinality and verifies diff behavior (removals applied, retained/added memberships correct, residue members untouched -- matching SDIFFSTORE semantics). Co-Authored-By: Claude Fable 5 --- core/src/main/kotlin/cohort/CohortStorage.kt | 9 ++-- .../test/kotlin/cohort/CohortStorageTest.kt | 50 +++++++++++++++++++ 2 files changed, 56 insertions(+), 3 deletions(-) diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt index 88a6864..5bf3f73 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/core/src/main/kotlin/cohort/CohortStorage.kt @@ -353,7 +353,6 @@ internal class RedisCohortStorage( newCohortKey: RedisKey, description: CohortDescription, existingSize: Int, - newSize: Int, ) { val existingCard = redis.scard(existingCohortKey) val newCard = redis.scard(newCohortKey) @@ -362,7 +361,11 @@ internal class RedisCohortStorage( applyAllAdditions(newCohortKey, description, existingSize, newCard, cohortIdSet) return } - val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1 + // Partition count is derived from the SCARD cardinalities, not the description-reported + // sizes: a crashed ingest can leave a version key holding more members than its published + // size, and an understated partition count would hold more than [diffPartitionMaxMembers] + // in memory at once — the bound this partitioning exists to enforce. + val partitions = ((maxOf(existingCard, newCard, 1L) - 1L) / diffPartitionMaxMembers).toInt() + 1 var addedCount = 0L var removedCount = 0L var existingDistinct = 0L @@ -722,7 +725,7 @@ internal class RedisCohortStorage( ) if (streamedDiffEnabled) { - applyMembershipDiff(existingCohortKey, newCohortKey, description, prev.size, finalSize) + applyMembershipDiff(existingCohortKey, newCohortKey, description, prev.size) } else { applySdiffstoreDiff(existingCohortKey, newCohortKey, description) } diff --git a/core/src/test/kotlin/cohort/CohortStorageTest.kt b/core/src/test/kotlin/cohort/CohortStorageTest.kt index d291734..5f40044 100644 --- a/core/src/test/kotlin/cohort/CohortStorageTest.kt +++ b/core/src/test/kotlin/cohort/CohortStorageTest.kt @@ -345,6 +345,56 @@ class CohortStorageTest { } } + @Test + fun `test redis, streamed diff partitions by actual cardinality when description size understates the key`(): Unit = + runBlocking { + val cohortStorage = + RedisCohortStorage( + "12345", + Duration.INFINITE, + "amplitude ", + redis, + redis, + 1000, + 1000, + CohortBlobCache(), + streamedDiffEnabled = true, + diffPartitionMaxMembers = 3, + ) + // v1 publishes with size 4... + val v1 = cohort("p", lastModified = 1, members = (1..4).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v1.toCohortDescription()) + acc.addMembers(v1.members.toList()) + acc.complete(v1.members.size) + } + // ...but the version key actually holds 10 members (residue of a crashed ingest), + // so the description understates the cardinality the diff must partition over. + redis.sadd( + RedisKey.CohortMembers("amplitude ", "12345", "p", "User", 1), + (5..10).map { "u$it" }.toSet(), + ) + // v2 keeps u3-u8 and adds u11-u12 + val v2 = cohort("p", lastModified = 2, members = ((3..8) + (11..12)).map { "u$it" }.toSet()) + run { + val acc = cohortStorage.createWriter(v2.toCohortDescription()) + acc.addMembers(v2.members.toList()) + acc.complete(v2.members.size) + } + for (removed in (1..2) + (9..10)) { + assertEquals(emptySet(), cohortStorage.getCohortMemberships("User", "u$removed")) + } + // Retained members with real memberships keep them; newly added members gain them. + for (member in (3..4) + (11..12)) { + assertEquals(setOf("p"), cohortStorage.getCohortMemberships("User", "u$member")) + } + // Residue members present in both version keys are (correctly) untouched by the + // diff — same as SDIFFSTORE semantics — so they never gain a membership. + for (residue in 5..8) { + assertEquals(emptySet(), cohortStorage.getCohortMemberships("User", "u$residue")) + } + } + @Test fun `test redis, streamed diff falls back to all additions when existing members key is missing`(): Unit = runBlocking {