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 {