Replace SDIFFSTORE cohort diff with streamed client-side diff

[adrianliu0815]

Problem

RedisCohortStorage computes the added/removed members between two cohort versions with two server-side SDIFFSTORE calls over the full old/new member sets. A set difference is O(N) and Redis executes it as a single atomic command on a single-threaded shard, so for multi-million member cohorts the diff blocks that shard — and its replicas, which re-execute the replicated command — for the entire duration. Every concurrent command routed to that shard queues behind it, including the SMEMBERS reads that serve /sdk/v2/memberships requests.

Observed in a production deployment running ~10M-member cohorts (Redis Cluster, 6 shards):

• SDIFFSTORE p50 ~88ms, but p75 ~2s / p95 ~8s / p99 ~9.9s during periods when the large cohorts republish (hourly warehouse-import recomputes → ~2 diffs per version pickup).
• Request-path SMEMBERS (normally p99 ~1.4ms) showed multi-second tail stalls that correlate 1:1 with slow SDIFFSTORE executions — over 3 days at 10-minute resolution, every SMEMBERS stall bin had a concurrent slow SDIFFSTORE and every SDIFFSTORE >5s bin produced SMEMBERS stalls.
• At ~10M members the diff reliably crosses the default 10s Lettuce command timeout (Default.REDIS_COMMAND_TIMEOUT_MILLIS), so the refresh throws, the version never publishes, and the next sync cycle re-downloads and re-diffs — a retry loop that re-runs the ~10s shard block while the cohort stays stale. A client-side timeout cannot cancel the server-side execution, so tightening timeouts cannot mitigate this.

Change

Compute the diff client-side from streamed data instead, so the largest single Redis command issued during a cohort update is one SSCAN page regardless of cohort size:

• SSCAN the existing member set in REDIS_SCAN_CHUNK_SIZE (1000) chunks into memory, then SSCAN the new set: members absent from the existing set are additions (flushed incrementally via the existing pipelined SADD path); members present in both are drained, so whatever remains afterward is exactly the set of removals (flushed via pipelined SREM).
• To bound proxy memory, members are hash-partitioned into ceil(size / diffPartitionMaxMembers) partitions (default 2M ⇒ at most ~2M member strings held at once, ~5 partitions for a 10M cohort) and diffed one partition at a time, at the cost of one extra SSCAN pass over both sets per additional partition.
• Both keys are immutable at diff time (the new set is fully written before complete(), the existing set is the previously published version); they are scanned from the primary connection to avoid misclassification from replica lag on the just-written new set.
• The CohortTemporary added/removed staging keys are no longer written.

Trade-offs: refresh wall-clock for very large cohorts increases (multiple sequential SSCAN passes instead of one blocking command) and the member data now transits to the proxy during diffing — but the refresh is a background loop, and in exchange serving latency is fully decoupled from cohort size and the diff can no longer exceed the command timeout.

Testing

• New test exercising a version update with adds and removals across multiple hash partitions (diffPartitionMaxMembers = 3 forces the partitioned path).
• Existing CohortStorageTest update/delete/membership tests pass unchanged on the new path (:core:test green, ktlint clean).

---

Happy to adjust the approach (e.g., gate behind a configuration flag, different partition sizing, or exposing diffPartitionMaxMembers via RedisConfiguration) if you'd prefer.

🤖 Generated with Claude Code

---

Note

Medium Risk
Changes core cohort publish semantics and Redis load patterns for every version update; incorrect diff logic would corrupt user cohort memberships, though behavior is covered by existing and new tests.

Overview
Cohort version updates no longer run two full-set SDIFFSTORE calls (and temporary added/removed keys). RedisCohortStorage now diffs old vs new member sets by streaming both cohort keys with SSCAN, computing adds/removals in the proxy, and applying them through the existing pipelined SADD / SREM membership paths.

Large cohorts are handled in hash partitions (default cap ~2M members per partition via diffPartitionMaxMembers) so only one partition of the prior version is held in memory at a time, at the cost of extra SSCAN passes. Diff scans use the primary Redis connection so replica lag does not misclassify members on the just-written new set.

A new test forces a small partition size to verify adds and removals across the multi-partition code path.

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

[kyeh-amp]

Thanks for the excellent writeup and fix — the diagnosis is consistent with the code, and the approach (never issue an O(cohort-size) Redis command on the refresh path) is the right shape. Approving with a few non-blocking suggestions.

Independently verified: :core:test passes on this branch including the new partition test, and the cluster analysis checks out — RedisKey.CohortMembers/CohortTemporary share the hash tag {projects:<pid>:cohort:<cid>} (RedisKey.kt), which is what SDIFFSTORE requires in cluster mode and exactly why the whole diff lands on one shard while UserCohortMemberships reads (no hash tag) spread across shards and stall when routed to it.

Correctness

Traced the algorithm, no bugs found:

• Partition cover is complete and disjoint. partitionOf (deterministic String.hashCode, correct negative handling) is applied identically to both scans, so per partition: old∩p loaded into a HashSet, new∩p drained from it — scan leftovers are exactly the adds, HashSet leftovers exactly the removals. The ceiling math is right, including the finalSize >= 1 guarantee at the call site.
• Immutability assumption holds. Both keys are versioned by lastModified, the new set is fully written before complete(), and concurrent loads are excluded by CohortLoadingLock. Scanning from the primary to avoid replica-lag misclassification on the just-written set is a good catch. Even a pathological SSCAN duplicate degrades to an idempotent SADD for a both-versions member — membership state stays correct, only the logged count inflates.
• Failure semantics are no worse than before. A mid-diff failure leaves partial SADD/SREMs, but the description isn't published, so the next sync recomputes the identical diff and converges idempotently — same non-transactional window the old path had during processMembershipUpdates. Strictly better in one respect: the old path could crash between sdiffstore and del and leak TTL-less CohortTemporary keys; this path writes no temp keys at all.
• Behavior parity preserved: the finally block still expires the old members/blob keys, the diff log line is unchanged, and the new constructor param defaults on an internal class — no API surface change.

Suggestions (non-blocking)

1. Yes to exposing diffPartitionMaxMembers via RedisConfiguration — taking you up on that offer. A 2M-member HashSet<String> is roughly 200–400MB of heap depending on ID length; not out of line with what the proxy already materializes elsewhere (getCohortMembers), but heap headroom varies by deployment.
2. Partition count trusts description sizes rather than actual cardinality. They should match SCARD in practice (published size is the real final count), but an O(1) SCARD on both keys would make the memory bound self-sufficient against a stale/incorrect description. Cheap hardening, optional.
3. Round-trip cost: a 10M cohort ⇒ 5 partitions × 2 sets × 10K SSCAN pages ≈ 100K sequential round trips per refresh (order of a minute at sub-ms RTT). Fine for a background loop and fully decoupled from serving, but a larger SSCAN COUNT for this path would cut round trips ~5x if refresh duration ever matters.
4. The memory bound is probabilistic — hash partitioning bounds the expected partition size; a skewed ID distribution could exceed it. Fine for real user IDs; at most worth a word in the doc comment.
5. Follow-up cleanup (we'll take it): Redis.sdiffstore (interface + both connection impls + test fake) and RedisKey.CohortTemporary are now dead in main code. All internal, safe to remove in a separate PR.

Tests

The new test genuinely exercises the partitioned path (diffPartitionMaxMembers = 3 ⇒ 4 partitions over 10 members) and the existing update test covers the single-partition path. One caveat: tests run against InMemoryRedis, whose sscanChunked is a plain chunked iteration, so real SSCAN cursor semantics aren't exercised — equally true of the old path, and your staging validation against a real cluster covers that gap in practice. We'll plan to canary this before cutting a release.

[vaibhav-jain-exp]

Thanks for this — the diagnosis is right, and we were able to reproduce the shard-stall correlation math you described. Before putting it on a release path we needed a few things layered on top: an opt-in flag (default off, SDIFFSTORE stays the default), env-exposed partition/chunk knobs (at 40M members the hardcoded 2M/1000 defaults mean ~1.6M sequential SSCAN round trips per refresh), loading-lock renewal for diffs that outlast the 900s lock TTL, and SCARD-based scan-truncation guards (a cluster failover mid-SSCAN can silently skip members, which on this code path becomes permanent false removals).

Rather than round-trip all of that through review comments, we've opened #40, which contains this PR's commits unchanged (you stay the author), your #39 cherry-picked on top, and the gating/hardening — see its description for the full list and rollout guidance for your 10-40M cohort scale.

If #40 looks good to you, we'd close this PR in its favor — but please poke holes in it first; you have the production deployment this was built for.