From b733a440db18cccb24042c1b544d54cb9a913214 Mon Sep 17 00:00:00 2001 From: Adrian Liu Date: Wed, 15 Jul 2026 23:44:34 -0700 Subject: [PATCH] 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 {