Async sticky bucket service support for GrowthBookClient

[madhuchavva]

Closes #127

Summary

Adds async sticky bucket service support to GrowthBookClient, mirroring the JS SDK's design: evaluation stays synchronous and CPU-only, while all sticky bucket I/O moves off the event loop — reads are prefetched (awaited natively for async services, executor-offloaded for sync ones) and writes are fire-and-forget. Use case: network-backed sticky bucketing (Redis, DynamoDB) in high-concurrency asyncio services without blocking the loop.

Why not an async evaluation core

Feature evaluation (core.py) is pure CPU — hashing, bucketing, condition matching — and never performs I/O. The TS SDK, the most async-native GrowthBook SDK, deliberately keeps evalFeature/run synchronous for the same reason (sticky-bucket-service.ts is Promise-based; GrowthBook.ts evalFeature is not). An async core would force either an event loop into the sync GrowthBook class or a duplicated core_async.py, and would double the cases.json conformance surface. This PR reproduces the JS shape instead: async at the I/O edges only.

What's included

New public API (common_types.py, exported from package root):

from growthbook import AbstractAsyncStickyBucketService

class RedisStickyBucketService(AbstractAsyncStickyBucketService):
    async def get_assignments(self, attributeName, attributeValue):
        return await self.redis.get(self.get_key(attributeName, attributeValue))

    async def save_assignments(self, doc):
        await self.redis.set(self.get_key(doc["attributeName"], doc["attributeValue"]), doc)

    # optional: override get_all_assignments() with a single MGET

Plus GrowthBookClient.flush_sticky_bucket_saves() for serverless environments (close() drains automatically). Existing sync AbstractStickyBucketService implementations keep working with both clients, unchanged. The sync GrowthBook class rejects async services at construction with a ValueError (it has no loop to await them on).

• Read path (growthbook_client.py::_refresh_sticky_buckets) — before, a sync service call ran directly on the event loop, guarded by a bool flag that was not a lock:

  # before
  while not self._sticky_bucket_cache_lock:
      ...
      assignments = self.options.sticky_bucket_service.get_all_assignments(attributes)  # blocks the loop
  
  # after
  async with self._sticky_bucket_lock:                      # real asyncio.Lock, coalesces concurrent refreshes
      if isinstance(service, AbstractAsyncStickyBucketService):
          assignments = await service.get_all_assignments(attributes)
      else:
          assignments = await loop.run_in_executor(None, service.get_all_assignments, attributes)

  The bool flag was dead weight while the body had no awaits, but becomes a real race once suspension points exist; the asyncio.Lock also coalesces concurrent refreshes for identical attributes (verified by test: 10 concurrent evals → 1 fetch).

• Write path (core.py step 13.5) — before, core called the service directly from synchronous run_experiment, blocking the loop mid-await client.eval_feature(...):

  # before
  evalContext.global_ctx.options.sticky_bucket_service.save_assignments(doc)  # blocks the loop
  
  # after
  if evalContext.save_sticky_bucket_doc:          # async client wires this; schedules fire-and-forget
      evalContext.save_sticky_bucket_doc(doc)
  else:                                           # sync client: byte-identical to before
      evalContext.global_ctx.options.sticky_bucket_service.save_assignments(doc)

  The callback is a defaulted field on EvaluationContext (not a threaded parameter, which eval_prereqs re-entry would drop). The in-memory assignment doc is still updated synchronously during eval, so read-your-writes holds while persistence completes in the background — the same mechanism as JS's stickyBucketAssignmentDocs. Scheduled saves are held as strong refs (bare create_task results can be GC'd mid-write), failures are logged and never raised into eval.

• Bug fix (own commit) — core replaced the assignment-docs dict when it was empty instead of mutating it in place (if not docs: docs = {}), severing the shared-cache reference. With an initially-empty service, every re-eval saw no existing assignments, flagged changed=True, and re-saved. Now mutates in place; regression test included.

• Async user callbacks — on_experiment_viewed, on_feature_usage, and subscription callbacks may now be coroutine functions; they are scheduled on the loop (previously a returned coroutine was silently dropped) and drained in close(). Sync callbacks are invoked exactly as before.

• Cleanups — stop_refresh runs blocking stopAutoRefresh(timeout=10) in the executor (was blocking the loop up to 10s on shutdown, inconsistent with the SSE teardown path that already used the executor); is_on/is_off/get_feature_value delegate to eval_feature instead of duplicating its body.

Measured impact

Benchmark harness is checked in at tests/scripts/benchmark_async_client.py (simulated asyncio service: 100 concurrent request handlers, 1000 requests, sticky bucket service with 1 ms simulated network latency; feature loading mocked). "Before" rows were produced by running the same harness against main (v2.3.1).

Scenario	Throughput	Event-loop lag (max)	Sticky docs persisted
Before (v2.3.1): sync sticky service, runs on the event loop	348 rps	2,869 ms — loop frozen for the whole run	1000
After (this PR): same sync sticky service, offloaded to executor	303 rps	8.7 ms	1000
After (this PR): async sticky service (new ABC), awaited natively	342 rps	10.5 ms	1000
After (this PR): async service, single hot user — cache-hit path	18,059 rps	0.4 ms	1
Control (both trees): no sticky service	~55,000 rps	—	—

Wins and observations

• The customer-reported problem is the loop-lag column, and it's fixed. On v2.3.1, a coroutine sharing the event loop with GrowthBook — a health check, another request handler, a websocket ping — was not scheduled once during the entire 2.9 s run: every sticky fetch and save ran on the loop. After this PR, worst-case scheduling delay is ~9 ms (a ~330× improvement), with no change required to existing sync service implementations.
• v2.3.1's per-eval latency numbers were an illusion. They looked fine (p50 2.6 ms) only because requests executed serially — each eval measured its own blocking work while every other request in the process waited. The "After" latencies are honest queueing numbers under real concurrency.
• Core evaluation is unchanged: ~10 µs/eval in the no-service control on both trees. The callback seam added to core.py has no measurable cost, and this number is also the argument for not making evaluation itself async — there is no I/O in it to overlap.
• The ~13% throughput gap between the sync-offloaded and native-async rows is the thread-hop tax of run_in_executor. That delta is the concrete reason for customers with network-backed services to implement the new AbstractAsyncStickyBucketService rather than keep a sync service, beyond avoiding executor-pool pressure.
• No writes were lost in any scenario despite fire-and-forget persistence — close() drains in-flight saves (1000/1000 persisted).

Trade-offs

• Two ABCs instead of one. A service author supporting both clients maintains a sync and an async implementation. The JS SDK chose the inverse (async-primary interface + sync adapter); doing that here would break every existing AbstractStickyBucketService subclass. A thin sync → async adapter can be added later without breaking anything.
• isinstance dispatch requires subclassing the ABC. A duck-typed async service that doesn't inherit AbstractAsyncStickyBucketService is treated as sync and its coroutine dropped (with an error log). A Protocol would be structurally looser but loses the shared get_key/get_all_assignments defaults.
• Fire-and-forget writes have a durability window. A crash between eval and write loses the assignment, and failed saves are logged, not retried. This matches JS exactly; callers needing stronger guarantees can await client.flush_sticky_bucket_saves().
• Concurrent saves for one key are unordered (also matches JS). Docs merge from shared in-memory state under the eval lock, so each scheduled doc is a superset of the previous — last-write-wins converges.
• Sync services share the loop's default ThreadPoolExecutor (max ≈ min(32, cpu+4) workers). A pathologically slow sticky backend can crowd out other run_in_executor users in the process.
• Async tracking callbacks are deduped at schedule time, not completion: if the scheduled coroutine later fails, that impression is not retried (a failing sync callback is retried on the next eval). Follow-up below.

What's NOT changed

• GrowthBook (sync) behavior is byte-identical: it sets no callback, so core takes the exact pre-PR code path. The full sync conformance suite passes untouched.
• Remote-eval mode still rejects sticky bucketing (matches JS).
• Write timing for the async client is now eventual (fire-and-forget) rather than synchronous — the only intentional behavior change, and it matches JS. Per-key ordering of concurrent saves is unordered, also matching JS; docs merge from shared in-memory state, so last-write-wins converges.

Follow-ups (Pipeline)

• The sticky bucket prefetch cache is single-slot (keyed on the last attributes dict), so multi-user workloads with distinct attributes refetch per eval and serialize behind the refresh lock (~340 rps ceiling at 1 ms service latency in the benchmark above; same ceiling as main, which additionally froze the loop). Fix: keyed LRU cache with per-key inflight coalescing.
• _eval_lock still serializes evaluations in CDN mode, including across the sticky prefetch await. Fix: immutable feature-snapshot swap.
• The async-tracking dedup asymmetry noted above. Fix: un-mark the dedup key when a scheduled tracking callback fails.
• Sync services could get a dedicated bounded executor pool and a save retry policy if demand appears.
• Async feature-cache/plugin migration; plugins remain thread-based.

Test plan

• Full suite: pytest tests/ -q --ignore=tests/scripts — 807 passed (was 785)
• Every cases.json stickyBucket conformance case now runs against BOTH service flavors (sync parametrization + async)
• Loop-not-blocked proofs for the read and write paths — deterministic event gates, no wall-clock sleeps: the sync service blocks its thread on a gate only a loop-side coroutine can release, so a blocked loop fails the test instead of flaking it
• Concurrency: 10 concurrent evals with identical attributes → exactly 1 service fetch; attribute change invalidates the cache
• Fire-and-forget semantics: doc visible in-memory immediately after eval; persisted after flush_sticky_bucket_saves(); failing save logged, never raised; unchanged assignment not re-saved (regression test for the dict-replacement bug)
• GrowthBook(sticky_bucket_service=AsyncService()) raises ValueError
• Async callback test: coroutine on_experiment_viewed/on_feature_usage/subscription all invoked and drained by close()