perf(lua): eliminate redundant ZSet member-prefix scan in zsetState#548
perf(lua): eliminate redundant ZSet member-prefix scan in zsetState#548
Conversation
|
Warning You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again! |
📝 WalkthroughWalkthroughRefactored ZSet type detection in the Redis adapter layer by introducing a Changes
Estimated code review effort🎯 3 (Moderate) | ⏱️ ~22 minutes Possibly related PRs
Poem
🚥 Pre-merge checks | ✅ 5✅ Passed checks (5 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
|
/gemini review |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request optimizes ZSet state detection by introducing zsetStorageHintAt to consolidate storage probes and reduce redundant scans. Feedback suggests further optimizing the TTL check for collection types and addressing redundant I/O in the rawZSetPhysTypeAt fallback path.
| } | ||
| h := zsetStorageHint{physType: physType, logType: physType, memberFound: memberFound} | ||
| if physType != redisTypeNone { | ||
| expired, err := r.hasExpiredTTLAt(ctx, key, readTS) |
There was a problem hiding this comment.
Since the physType is already known at this point, we can optimize the TTL check by skipping the probe for the !redis|str| key if the type is not a string. This avoids an unnecessary GetAt call on the hot path for collection types.
| expired, err := r.hasExpiredTTLAt(ctx, key, readTS) | |
| expired, err := r.hasExpired(ctx, key, readTS, isNonStringCollectionType(physType)) |
| return redisTypeZSet, false, nil | ||
| } | ||
| // Not a wide-column ZSet — full detection for other types. | ||
| physType, err := r.rawKeyTypeAt(ctx, key, readTS) |
There was a problem hiding this comment.
Calling rawKeyTypeAt here leads to redundant I/O because rawKeyTypeAt internally calls detectWideColumnType, which will re-scan the ZSet member prefix, meta key, and delta prefix that were just checked in this function. While this only affects non-ZSet or legacy ZSet keys, it could be optimized by refactoring rawKeyTypeAt to allow skipping the ZSet check.
…549) ## Summary Introduce a leader-local lease read path so leadership confirmation is amortised across reads. Implements the design in `docs/lease_read_design.md`. - `raftengine.LeaseProvider` (optional interface): `LeaseDuration()`, `AppliedIndex()`, `RegisterLeaderLossCallback()`. Implemented by the etcd engine; the hashicorp engine and test stubs leave it unimplemented and callers fall back to `LinearizableRead`. - etcd engine: `LeaseDuration = electionTimeout - 300ms` (700 ms with current tick config), `AppliedIndex` from the published Status. `refreshStatus` fires registered callbacks on leader -> non-leader transitions, and `shutdown()` fires them on close while still leader. - `kv.Coordinate` / `kv.ShardedCoordinator`: `LeaseRead` and `LeaseReadForKey`. Per-coordinator lease for `Coordinate`, per-shard lease for `ShardedCoordinator` (via `leaseRefreshingTxn` wrapper around `g.Txn`). Fast path returns the engine applied index when the lease is unexpired; slow path runs `LinearizableRead` and refreshes the lease on success. `Coordinate.Dispatch` and every `ShardedCoordinator` dispatch path refresh the lease on a successful commit using the pre-dispatch timestamp. - `leaseState`: lock-free `atomic.Pointer[time.Time]` plus a generation counter. `invalidate()` bumps `gen` before clearing `expiry`; `extend()` captures `gen` at entry and undoes its own CAS if `gen` moved, so a Dispatch that succeeded just before a leader-loss invalidate cannot resurrect the lease. - Callers switched: - `adapter/redis_lua_context.go` `newLuaScriptContext` — was full per-script ReadIndex, now lease-aware. - `adapter/redis.go` `get` — was no quorum check, now bounded by lease with `redisDispatchTimeout` context. - `adapter/dynamodb.go` `getItem` — was no quorum check, now bounded by lease via `LeaseReadForKey(dynamoTableMetaKey(tableName))` so sharded deployments consult the owning shard. Input parsing extracted into `parseGetItemInput` to stay under the cyclop limit. Remaining read paths (KEYS, EXISTS, ZSet/Hash/List/Set/Stream readers, DynamoDB query/scan/transactGet/batchGet) still rely on the lease being kept warm by Lua scripts and successful Dispatch calls; tracked as #557 follow-up. ## Motivation Profiling after PR #547/#548 showed average `redis.call()` time of 800 ms - 2.2 s, with `redis.call()` accounting for ~100% of Lua VM time. Investigation traced this to `newLuaScriptContext` calling `coordinator.LinearizableRead(ctx)` per script — a full etcd/raft `ReadOnlySafe` ReadIndex (heartbeat broadcast + quorum `MsgHeartbeatResp` wait) on every Lua script invocation. A lease-based check skips the broadcast under steady load. Stale-read window is bounded by `LeaseDuration < electionTimeout`, the same trade-off DynamoDB / Redis non-Lua already accept (and that change tightens for those two paths as a side benefit). ## Test plan - [x] `go build ./...` passes - [x] `go test ./adapter/... ./kv/... ./internal/...` passes - [x] `go test -race` on all lease tests passes - [x] `TestCoordinate_LeaseRead_AmortizesLinearizableRead` proves 100 LeaseRead calls within one lease window trigger exactly 1 underlying LinearizableRead - [ ] Sustained-load test: confirm Lua script throughput improves and `LinearizableRead` call rate drops below the script invocation rate - [ ] Partition-style test: confirm a stale leader stops serving reads at most `LeaseDuration` after losing quorum ## Follow-ups (tracked) - #553 lock-free `AppliedIndex` - #554 `atomic.Int64` lease expiry (remove heap alloc per extend) - #555 cached `LeaseProvider` type assertion - #556 background lease warm-up to flatten read-only sawtooth - #557 wrap remaining adapter read handlers in `LeaseRead` ## Notes - No metric for lease hit/miss yet (one of the open questions in the design doc). <!-- This is an auto-generated comment: release notes by coderabbit.ai --> ## Summary by CodeRabbit * **New Features** * Leader-local lease-based read APIs added (including per-shard lease reads) and optional engine lease capability to enable fast-path reads. * **Refactor** * DynamoDB/Redis read flows now consult lease checks with bounded timeouts and re-verify routing/ts to avoid stale reads; Redis request contexts and cancellation tightened for safer timeouts and shutdown. * Coordinators refresh/invalidate leases around commits and leader-loss events. * **Documentation** * Added lease-read design document. * **Tests** * Extensive unit tests for lease state, coordinator/sharded behavior, and leader-loss callbacks. * **Chores** * Coordinator cleanup ensured on server shutdown. <!-- end of auto-generated comment: release notes by coderabbit.ai -->
keyTypeAt → rawKeyTypeAt → detectWideColumnType already scans ZSetMemberScanPrefix to detect the type. zsetState then performed a second identical ScanAt to decide wide-column vs legacy-blob format. Add zsetStorageHintAt (+ rawZSetPhysTypeAt + zsetMetaOrDeltaExistsAt) that probes the member prefix once and carries the memberFound signal back to zsetState. The prior rawKeyTypeAt call in the redisTypeNone branch (for TTL-expired key detection) is also eliminated — physType / logType on the hint covers both cases in a single server round-trip. For the common ZSet-with-members path: 2 ScanAt → 1 ScanAt. For the TTL-expired ZSet path: 3 probe calls → 1–2 probe calls.
Rebase onto main + address gemini review findings on PR #548: - zsetState: consult c.cachedType(key) before calling the new zsetStorageHintAt. Without this guard, a pre-script storage probe would miss in-script mutations like SET / DEL / type-change: e.g. `SET k x; ZSCORE k m` returned the old score instead of WRONGTYPE. Mirrors the keyType() fallback order (cachedType first, storage second). - zsetStorageHintAt: use the skip-zset variant of rawKeyTypeAt in its fallback. Callers have already ruled out ZSet (member prefix + meta + delta probes came back empty), so re-issuing those three probes inside detectWideColumnType was redundant I/O on the non-ZSet fallback path. Adds rawKeyTypeAtSkipZSet and detectWideColumnTypeSkipZSet helpers. - zsetStorageHintAt: skip the `!redis|str|` TTL probe when physType is a non-string collection type, saving a GetAt on the hot path.
bootjp
force-pushed
the
perf/zset-io-optimization
branch
from
3dcc8ee to
5b0c56d
Compare
|
Rebased onto latest main and addressed gemini-code-assist's review (force-push, commit Correctness (regression caught during rebase)
Addressed review comments
Ready for re-review. |
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 1
🧹 Nitpick comments (1)
adapter/redis_lua_context.go (1)
721-742:typ == redisTypeZSetbranch appears unreachable and drops storage-hint state.
cachedTypeonly reportsredisTypeZSetwhenhasLoadedZSetValue(c.zsets[key])is true, which requiresc.zsets[key]to be non-nil withloaded && exists. ButzsetState(Lines 662–664) already short-circuits onc.zsets[k]being present, so by the timezsetStateFromCachedTyperuns, the ZSet branch at Lines 734–740 is effectively dead code.If this branch ever does become reachable (e.g., a future refactor populates
c.zsets[k]without going throughzsetState), it silently skips thezsetStorageHintAtprobe:legacyBlobBasestays false andphysicallyExistsAtStartstays false regardless of actual storage shape. Downstream,zsetCommitPlanwould then take the wide-column delta path on what may actually be a legacy-blob ZSet, and the TTL-expired-recreation cleanup wouldn't fire.Two reasonable paths forward:
- Drop the branch and let the
cached/non-ZSet fallthrough (returning wrongTypeError) be the only in-script cached-type outcome, since the live-ZSet case is already handled by the early-return inzsetState.- Keep it, but document the reachability contract and either compute
legacyBlobBasehere or delegate to the storage-hint path.♻️ Minimal change: drop the unreachable branch
func (c *luaScriptContext) zsetStateFromCachedType(key []byte, st *luaZSetState) (*luaZSetState, bool, error) { typ, cached := c.cachedType(key) if !cached { return nil, false, nil } if typ == redisTypeNone { st.loaded = true return st, true, nil } - if typ == redisTypeZSet { - // Live ZSet via in-script mutation — loadZSetAt under - // ensureZSetLoaded will surface the pre-script storage rows. - st.loaded = true - st.exists = true - return st, true, nil - } + // Note: cachedType only reports redisTypeZSet when c.zsets[key] + // is already populated, which zsetState's early return handles + // before reaching this helper. return nil, true, wrongTypeError() }🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed. In `@adapter/redis_lua_context.go` around lines 721 - 742, The typ == redisTypeZSet branch in zsetStateFromCachedType is effectively unreachable because cachedType returns redisTypeZSet only when hasLoadedZSetValue(c.zsets[key]) is true and zsetState already short-circuits when c.zsets[k] is present; remove this branch to avoid silently dropping storage-hint state and let the non-ZSet fallthrough return wrongTypeError, or alternatively if you want to keep it compute and populate legacyBlobBase/physicallyExistsAtStart here or call zsetStorageHintAt so zsetCommitPlan receives correct storage hints; update or delete the typ == redisTypeZSet block in zsetStateFromCachedType accordingly and ensure consistency with zsetState, cachedType, hasLoadedZSetValue, and zsetStorageHintAt.
🤖 Prompt for all review comments with AI agents
Verify each finding against the current code and only fix it if needed.
Inline comments:
In `@adapter/redis_compat_helpers.go`:
- Around line 160-184: rawZSetPhysTypeAt currently probes ZSet members first
(via prefixExistsAt/ ZSetMemberScanPrefix) which violates the documented "string
wins" invariant; fix by checking string/other non-ZSet key types before the
member-prefix scan: call rawKeyTypeAtSkipZSet (or the string-specific probe path
used by rawKeyTypeAt) prior to prefixExistsAt, keep the existing
zsetMetaOrDeltaExistsAt check, and return redisTypeString if the pre-scan finds
a string; update zsetStorageHintAt and any caching assumptions (e.g.,
redis_lua_context.go zsetState / keyType()) accordingly and add a unit test for
ZSet+string corruption to ensure the string-wins behavior.
---
Nitpick comments:
In `@adapter/redis_lua_context.go`:
- Around line 721-742: The typ == redisTypeZSet branch in
zsetStateFromCachedType is effectively unreachable because cachedType returns
redisTypeZSet only when hasLoadedZSetValue(c.zsets[key]) is true and zsetState
already short-circuits when c.zsets[k] is present; remove this branch to avoid
silently dropping storage-hint state and let the non-ZSet fallthrough return
wrongTypeError, or alternatively if you want to keep it compute and populate
legacyBlobBase/physicallyExistsAtStart here or call zsetStorageHintAt so
zsetCommitPlan receives correct storage hints; update or delete the typ ==
redisTypeZSet block in zsetStateFromCachedType accordingly and ensure
consistency with zsetState, cachedType, hasLoadedZSetValue, and
zsetStorageHintAt.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 171149cb-c69d-4300-abb7-583d1f676d35
📒 Files selected for processing (2)
adapter/redis_compat_helpers.goadapter/redis_lua_context.go
| // rawZSetPhysTypeAt detects whether a ZSet exists physically at readTS (ignoring | ||
| // TTL) and whether the detection was via the member-prefix scan (memberFound). | ||
| // For non-ZSet keys it falls back to rawKeyTypeAt. | ||
| func (r *RedisServer) rawZSetPhysTypeAt(ctx context.Context, key []byte, readTS uint64) (redisValueType, bool, error) { | ||
| // Single scan: probe member prefix (common path). | ||
| memberFound, err := r.prefixExistsAt(ctx, store.ZSetMemberScanPrefix(key), readTS) | ||
| if err != nil { | ||
| return redisTypeNone, false, err | ||
| } | ||
| if memberFound { | ||
| return redisTypeZSet, true, nil | ||
| } | ||
| // No member rows — check meta/delta for a memberless wide-column ZSet. | ||
| zsetOnly, err := r.zsetMetaOrDeltaExistsAt(ctx, key, readTS) | ||
| if err != nil { | ||
| return redisTypeNone, false, err | ||
| } | ||
| if zsetOnly { | ||
| return redisTypeZSet, false, nil | ||
| } | ||
| // Not a wide-column ZSet — probe other types without re-scanning | ||
| // the three ZSet-side prefixes we already ruled out above. | ||
| physType, err := r.rawKeyTypeAtSkipZSet(ctx, key, readTS) | ||
| return physType, false, err | ||
| } |
There was a problem hiding this comment.
🧩 Analysis chain
🏁 Script executed:
#!/bin/bash
# Look for existing tests that exercise the mixed string+zset corruption path
# to confirm whether the new zsetStorageHintAt ordering is already covered.
rg -nP --type=go -C3 '(string.*wins|mixed.*corrupt|replaceWithStringTxn|redisStrKey.*ZSet|ZSet.*redisStrKey)'
rg -nP --type=go -C2 '\brawKeyTypeAt\b' -g '*_test.go'Repository: bootjp/elastickv
Length of output: 3049
🏁 Script executed:
# Search for tests involving mixed string and zset scenarios
rg -nP --type=go 'test.*[Zz]set.*[Ss]tring|test.*[Ss]tring.*[Zz]set' -A 5 -B 2Repository: bootjp/elastickv
Length of output: 42
🏁 Script executed:
# Look for tests that may cover corruption recovery with multiple types
rg -nP --type=go -l 'corruption|mixed.*type' -g '*_test.go'Repository: bootjp/elastickv
Length of output: 205
🏁 Script executed:
# Find zsetState function in lua_context and check its implementation
rg -nP --type=go -A 15 'func.*zsetState' adapter/Repository: bootjp/elastickv
Length of output: 3070
🏁 Script executed:
# Compare ordering in rawZSetPhysTypeAt vs rawKeyTypeAt implementations
rg -nP --type=go -A 30 'func.*rawKeyTypeAt\(' adapter/redis_compat_helpers.goRepository: bootjp/elastickv
Length of output: 1368
🏁 Script executed:
# Check if hasHigherPriorityStringEncoding is used to enforce string priority elsewhere
rg -nP --type=go 'hasHigherPriorityStringEncoding' -B 2 -A 5Repository: bootjp/elastickv
Length of output: 7127
🏁 Script executed:
# Find zsetStorageHintAt implementation
rg -nP --type=go -A 20 'func.*zsetStorageHintAt'Repository: bootjp/elastickv
Length of output: 1693
🏁 Script executed:
# Search for ZSet tests that might cover string+zset corruption
rg -nP --type=go -B 3 -A 10 'TestRedis.*[Zz]set.*[Cc]orrup|TestRedis.*[Dd]ual.*[Zz]set' -g '*_test.go'Repository: bootjp/elastickv
Length of output: 42
🏁 Script executed:
# Check if there are any ZADD tests that exercise mixed encoding
rg -nP --type=go -B 2 -A 10 'ZADD.*string|string.*ZADD|corrupt.*ZADD' -g '*_test.go'Repository: bootjp/elastickv
Length of output: 42
🏁 Script executed:
# Find all references to rawZSetPhysTypeAt to understand its usage
rg -nP --type=go 'rawZSetPhysTypeAt' -B 3 -A 3Repository: bootjp/elastickv
Length of output: 3592
ZSet-first probe diverges from documented "string wins" tiebreaker.
rawKeyTypeAt documents the invariant (lines 219–225) that when a key carries both a string AND a collection entry (corruption recovery only), string probes run first so string wins. rawZSetPhysTypeAt inverts this: the ZSetMemberScanPrefix scan runs before rawKeyTypeAtSkipZSet, so when both a ZSet member row and a !redis|str|/bare/HLL entry exist for the same key, rawZSetPhysTypeAt returns redisTypeZSet where rawKeyTypeAt would return redisTypeString.
The comment in zsetStorageHintAt (lines 147–148) acknowledges this: "For any string types we reach via rawZSetPhysTypeAt's fallback (mixed corruption), we still need the string-side check too."
Downstream impact in redis_lua_context.go zsetState: if the hint returns redisTypeZSet for a corrupted key, the cached type bypasses the type-mismatch check that the full keyType() path would catch, allowing ZADD to proceed on a key whose actual keyType is string.
Unlike hash+string (guarded by TestRedis_HGET_FastPathGuardDualEncoding), there is no ZSet+string corruption test. Consider either adding one or reordering the probes in rawZSetPhysTypeAt to check strings before the member-prefix scan to preserve the documented invariant.
🤖 Prompt for AI Agents
Verify each finding against the current code and only fix it if needed.
In `@adapter/redis_compat_helpers.go` around lines 160 - 184, rawZSetPhysTypeAt
currently probes ZSet members first (via prefixExistsAt/ ZSetMemberScanPrefix)
which violates the documented "string wins" invariant; fix by checking
string/other non-ZSet key types before the member-prefix scan: call
rawKeyTypeAtSkipZSet (or the string-specific probe path used by rawKeyTypeAt)
prior to prefixExistsAt, keep the existing zsetMetaOrDeltaExistsAt check, and
return redisTypeString if the pre-scan finds a string; update zsetStorageHintAt
and any caching assumptions (e.g., redis_lua_context.go zsetState / keyType())
accordingly and add a unit test for ZSet+string corruption to ensure the
string-wins behavior.
1 participant
Summary
zsetStatepreviously scannedZSetMemberScanPrefixtwice — once viakeyTypeAt → rawKeyTypeAt → detectWideColumnType, then again to choose wide-column vs legacy-blob format.zsetStorageHintAt(withrawZSetPhysTypeAt+zsetMetaOrDeltaExistsAthelpers) probes the member prefix once and carries thememberFoundsignal back, also folding the TTL-expired-key detection into the same hint.ScanAt-> 1ScanAt.Motivation
Addresses Copilot review comment on PR #547 (line 689) about redundant I/O in
zsetState.Test plan
go test ./adapter/... -run ZSetSummary by CodeRabbit