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/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 bf6ae8f..5bf3f73 100644 --- a/core/src/main/kotlin/cohort/CohortStorage.kt +++ b/core/src/main/kotlin/cohort/CohortStorage.kt @@ -25,6 +25,13 @@ import kotlin.time.Duration.Companion.hours // 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 + /** * 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 @@ -34,6 +41,13 @@ private const val REDIS_SCAN_CHUNK_SIZE: Int = 1000 */ 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, @@ -102,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 @@ -127,6 +149,9 @@ internal fun getCohortStorage( redisConfiguration.scanLimit, redisConfiguration.pipelineBatchSize, CohortBlobCache(), + redisConfiguration.streamedCohortDiffEnabled, + redisConfiguration.cohortDiffPartitionMaxMembers, + redisConfiguration.cohortDiffScanChunkSize, ) } else { InMemoryCohortStorage() @@ -227,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( @@ -238,14 +267,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 @@ -275,6 +318,304 @@ 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. + * + * 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 + * 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. + * + * 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, + newCohortKey: RedisKey, + description: CohortDescription, + existingSize: Int, + ) { + 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 + } + // 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 + var scanPages = 0 + for (partition in 0 until partitions) { + 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() + 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 + } + // 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 >= 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) + } + val removed = mutableListOf() + for (member in existingPartition) { + removed.add(member) + 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 += + flushMembershipUpdates(removed, description.groupType, cohortIdSet) { updates, batchSize -> + redis.sremPipeline(updates, batchSize) + } + } + // 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( + 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 + } + + /** + * 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) @@ -383,58 +724,10 @@ 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()}", - ) - - 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) + if (streamedDiffEnabled) { + applyMembershipDiff(existingCohortKey, newCohortKey, description, prev.size) + } else { + applySdiffstoreDiff(existingCohortKey, newCohortKey, description) } } else { // No previous cohort: all members are additions @@ -452,16 +745,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 — @@ -528,10 +824,27 @@ 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 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) 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 31ed1f2..5f40044 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 @@ -18,6 +19,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 @@ -205,6 +207,237 @@ 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(), + streamedDiffEnabled = true, + 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, 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 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 { + 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 { @@ -282,4 +515,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)) + } } 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 + } }