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

[vaibhav-jain-exp]

Summary

Builds on @adrianliu0815's #38 and #39 (both authored commits preserved on this branch) and adds the gating and hardening we need before the streamed diff can be on a release path. Supersedes #38; if #39 merges first, the identical hunks here resolve cleanly.

From #38 (unchanged): the client-side streamed cohort diff — SSCAN both member-set versions in hash partitions and compute adds/removals in the proxy, so no O(cohort-size) blocking command ever runs on a Redis shard.

From #39 (cherry-picked): pending-version 24h TTL at ingestion, PERSIST at promotion, and retiring the previous version only after successful promotion — so failed refreshes can neither strand full member sets in Redis nor expire the still-published live version.

Added here:

1. Opt-in flag — AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED (default false). The SDIFFSTORE path remains the default with a command stream identical to main, plus Prevent orphaned cohort version sets when a refresh fails mid-flight #39's temp-key backstop TTLs armed immediately after each SDIFFSTORE (previously addedKey sat unprotected through the second, multi-second blocking command — the likeliest crash point).
2. Tuning knobs, validated > 0 at startup — AMPLITUDE_COHORT_DIFF_PARTITION_MAX_MEMBERS (default 2M) and AMPLITUDE_COHORT_DIFF_SCAN_CHUNK_SIZE (default 1000). At 40M members the defaults cost ~1.6M sequential SSCAN round trips per refresh (~10-45 min); partition 8-10M + chunk 5000 brings that to ~1-4 min. Note booleanEnv only accepts true — 1/yes silently mean off.
3. Lock renewal — new atomic Redis.renewLock() (compare-and-expire, mirroring releaseLock); the diff renews the 900s cohort-loading lock every 512 scan pages and aborts fail-closed if renewal fails, so a diff outlasting the lock TTL can't run concurrently with a second instance's load of the same cohort.
4. Scan-truncation guards — new Redis.scard(); every SSCAN pass cross-checks raw returned count against cardinality and aborts on shortfall. A cluster failover or slot migration can silently invalidate an SSCAN cursor mid-scan; on this code path a truncated scan becomes false removals of unchanged users that no later diff repairs. A cross-partition distinct-count check additionally catches duplicate-masked skips on the existing set.
5. Missing-previous-key fallback — if the published version's members key is gone at diff time, the update degrades to SDIFFSTORE semantics (all additions, no removals) under the same scan guard, with an error log — instead of failing every retry forever.
6. Bugbot's finding on Prevent orphaned cohort version sets when a refresh fails mid-flight #39, fixed both ways — complete() now publishes the description before PERSIST, so a refresh dying pre-publish always self-cleans via the pending TTL; and the inverse crash window (live set left with a 24h fuse that nothing would ever clear, which a plain order-swap would create) is repaired within one sync cycle by an O(1) ensureCurrentVersionPersisted on the not-modified loader path.
7. Fixed in passing — lock values in both Redis connections were built with ${'$'} escapes, producing the same literal string on every instance and making releaseLock's compare-and-delete vacuous across pods.

Testing

• CohortStorageTest 14/14 (union of Replace SDIFFSTORE cohort diff with streamed client-side diff #38/Prevent orphaned cohort version sets when a refresh fails mid-flight #39 tests plus: flag-off SDIFFSTORE update path, >1000 flush-threshold streamed update, missing-key fallback, knob validation, fault-injected crash-at-publish TTL survival, stray-TTL self-heal); full ./gradlew test ktlintCheck green.
• Refresh lifecycle exercised against live Redis in both flag states: success, mid-diff failures, fault-injected lock loss, missing-key fallback, first load, empty cohort, and unmodified re-complete — which now leaves no TTL on the live version's key (on main, the finally-block expire armed a cohortSyncIntervalMillis fuse on it).

Rollout guidance for large-cohort deployments

Flag on with partition 8-10M and chunk 5000 for 10-40M member cohorts; keep the cohort sync interval at or above expected diff duration; watch per-cohort diff duration and added/removed counts (a full-size "added" spike with zero removals is the signature of a missing-previous-key fallback).

---

Note

High Risk
Changes cohort refresh semantics, per-user membership writes, and distributed locking in Redis; incorrect diff or lock handling could corrupt memberships or allow concurrent loads on multi-million-member cohorts.

Overview
Adds an opt-in streamed client-side cohort diff (AMPLITUDE_COHORT_STREAMED_DIFF_ENABLED, default off) so very large refreshes avoid blocking SDIFFSTORE on a Redis shard; when disabled, behavior stays on the existing server-side diff path with safer temp-key TTL handling.

Streamed diff (when enabled): partitions member sets by hash, diffs via chunked SSCAN + SCARD completeness checks, pipelines membership updates, renews the cohort loading lock during long scans, and degrades to “all additions” if the previous members key is missing.

Promotion / recovery: complete() now publishes the description before PERSIST on the new member set; ensureCurrentVersionPersisted on unchanged sync cycles clears a stray TTL on the live set. Redis: scard, renewLock, per-instance lock values (fixes broken ${'$'} lock strings), and env-tunable partition/chunk sizes validated at startup.

Tests cover both diff paths, crash-at-publish TTL, and TTL self-heal.

^{Reviewed by Cursor Bugbot for commit 1556659. Bugbot is set up for automated code reviews on this repo. Configure here.}

[cursor]

Cursor Bugbot has reviewed your changes and found 1 potential issue.

Bugbot Autofix prepared a fix for the issue found in the latest run.

• ✅ Fixed: Partition count ignores Redis cardinality
  • Changed partition count calculation to use existingCard and newCard from SCARD instead of description sizes existingSize and newSize, ensuring memory bounds are enforced based on actual Redis data.

Or push these changes by commenting:

@cursor push 9c51505bdf

Preview (9c51505bdf)

diff --git a/core/src/main/kotlin/cohort/CohortStorage.kt b/core/src/main/kotlin/cohort/CohortStorage.kt
--- a/core/src/main/kotlin/cohort/CohortStorage.kt
+++ b/core/src/main/kotlin/cohort/CohortStorage.kt
@@ -362,7 +362,7 @@ internal class RedisCohortStorage(
             applyAllAdditions(newCohortKey, description, existingSize, newCard, cohortIdSet)
             return
         }
-        val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1
+        val partitions = ((maxOf(existingCard, newCard, 1L) - 1) / diffPartitionMaxMembers).toInt() + 1
         var addedCount = 0L
         var removedCount = 0L
         var existingDistinct = 0L

@@ -362,7 +362,7 @@ internal class RedisCohortStorage(
             applyAllAdditions(newCohortKey, description, existingSize, newCard, cohortIdSet)
             return
         }
-        val partitions = ((maxOf(existingSize, newSize, 1) - 1) / diffPartitionMaxMembers) + 1
+        val partitions = ((maxOf(existingCard, newCard, 1L) - 1) / diffPartitionMaxMembers).toInt() + 1
         var addedCount = 0L
         var removedCount = 0L
         var existingDistinct = 0L

_{You can send follow-ups to the cloud agent here.}

^{Reviewed by Cursor Bugbot for commit f09c9f9. Configure here.}