🎯 Fix Cache Key Collision and Partial Results during Tenant Shuffle-Sharding Rebalance in Query Frontend

[rasoolharlym8]

📝 Description

During tenant shuffle-sharding rebalancing events under concurrent query load, the Query Frontend intermittently merges cached query fragments generated by different (and potentially outdated) shard assignments. This results in a race condition where a PromQL query returns a partial result set to the client while still responding with an incorrect HTTP 200 OK status.

This issue is transient and isolated to the window of time when a tenant's shard assignment is transitioning. The root cause lies in the cache key generation logic, which does not currently account for active shard routing topologies or epoch versions, combined with a lack of validation during shard ownership transitions in the Query Frontend.

🎯 Acceptance Criteria

• Topology-Aware Cache Keys: Modify the cache key generation algorithm in the Query Frontend to incorporate the active shard assignment identifier, routing table hash, or epoch version.
• Stale Cache Invalidation: Ensure that any query fragment cache writes initiated under an older shard configuration are discarded if a rebalance occurs mid-flight.
• Strict Error Handling: If a query execution spans a shard transition and detects mismatched or incomplete shard responses, the Query Frontend must fail the query (e.g., returning an HTTP 500/503 or triggering an internal retry) instead of returning a partial result set with an HTTP 200 OK.
• Thread-Safe Routing Updates: Ensure tenant routing table updates in the Query Frontend are atomic and thread-safe under high concurrent query loads.

🛠️ Technical Specifications & Context

Within the rasoolharlym8/CockroachDB repository (specifically looking at the query distribution and caching layers, typically located in paths like pkg/queryfrontend/, pkg/frontend/queryrange/, or equivalent query-routing directories):

1. Cache Key Generation:

   • Locate the cache key generator (e.g., Extractor or KeyGenerator implementations in the query range split-and-merge middleware).
   • Update the key generation struct/function to append the tenant's current shard configuration hash or epoch version:

     // Example conceptual change
     func (g *KeyGenerator) GenerateKey(req Request, tenantID string, shardEpoch int64) string {
         return fmt.Sprintf("%s:%d:%s", tenantID, shardEpoch, req.Hash())
     }

2. Shard Transition Validation:

   • In the query execution path (where sub-queries/fragments are dispatched to downstream shard targets), capture the routing table epoch at the start of the request.
   • Before writing merged results back to the query cache, verify that the routing table epoch has not changed. If a change is detected, discard the cache write and discard/retry the query.

3. Concurrency Control:

   • Review the tenant routing state storage in the Query Frontend. Ensure that updates to the routing table use atomic pointers (sync/atomic) or read-write locks (sync.RWMutex) to prevent dirty reads during rebalancing.

🧪 Verification & Testing

• Unit Tests:
  • Add unit tests in the query frontend package to verify that cache keys differ when the same query is executed under two different shard configurations.
  • Write a test simulating a mid-flight shard configuration change to verify that the cache write-back is safely aborted.
• Integration & Chaos Testing:
  • Implement a mock test case simulating a tenant shuffle-sharding rebalance.
  • Generate concurrent PromQL query load during the rebalance and assert that:
    1. No partial results are returned to the client.
    2. All successful responses (HTTP 200 OK) contain complete datasets.
    3. Any interrupted queries are retried internally or fail gracefully with an appropriate non-200 error code.
• Race Detection:
  • Run the test suite with the Go race detector enabled (go test -race ./...) to ensure no data races exist in the updated routing table lookup and cache key generation paths.

---

---

This repo is using Opire - what does it mean? 👇

💵 Everyone can add rewards for this issue commenting /reward 100 (replace 100 with the amount).
🕵️‍♂️ If someone starts working on this issue to earn the rewards, they can comment /try to let everyone know!
🙌 And when they open the PR, they can comment /claim #3 either in the PR description or in a PR's comment.

🪙 Also, everyone can tip any user commenting /tip 20 @rasoolharlym8 (replace 20 with the amount, and @rasoolharlym8 with the user to tip).

📖 If you want to learn more, check out our documentation.

[rasoolharlym8]

/reward 10

[opirebot]

😅 Ups... you cannot create a reward of $10.00. It needs to be at least $20.00

[testdasislyt]

Hi team,

This issue caught my eye, and I'm interested in helping resolve it. The description points to cache key collisions stemming from different shard assignment states during rebalance.

My initial thought would be to ensure the cache key for query fragments explicitly incorporates the current shard assignment version or epoch. This would prevent a newly rebalanced Query Frontend instance from serving stale fragments generated under an older assignment, ensuring uniqueness per tenant and per active shard configuration. I'd start by investigating how shard assignment versions are currently exposed (or could be exposed) to the Query Frontend and integrated into the cache key generation logic.

Is this issue still available for contributors?

Thanks!

[jamboriu]

/try

Check if you're the only one trying 👇

You're the first one to try to solve this issue, go for it! 😉