Streamed cohort diff behind a flag, hardened for 10-40M member cohorts (supersedes #38, includes #39)

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

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

core/src/main/kotlin/util/redis/Redis.kt

@@ -36,6 +36,8 @@ internal interface Redis {
 
     suspend fun smembers(key: RedisKey): Set<String>
 
+    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
 }

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<Long>(script, io.lettuce.core.ScriptOutputType.INTEGER, arrayOf(key.value), lockValue, ttlSeconds.toString())
+            }
+
+        return result == 1L
+    }
+
     suspend fun pipeline(block: suspend RedisAdvancedClusterAsyncCommands<String, String>.() -> List<RedisFuture<*>>) {
         withContext(pipelineContext) {
             val conn = pipelineConnection.await()

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<Long>(luaScript, io.lettuce.core.ScriptOutputType.INTEGER, arrayOf(key.value), lockValue, ttlSeconds.toString())
+            }
+
+        return result == 1L
+    }
+
     private suspend fun pipeline(block: RedisAsyncCommands<String, String>.() -> List<RedisFuture<*>>) {
         withContext(pipelineContext) {
             val connection = pipelineConnection.await()