Fix LRUQueryCache hashCode/equals contract and account query RAM once per distinct query

[sgup432]

Description

LRUQueryCache was recently refactored (releasing in Lucene 11.0) to partition the cache, and its key structure was changed into a composite one (QueryCacheKey).

I found a bug in QueryCacheKey: it compares queries by value in equals() but hashes them with System.identityHashCode(query), breaking the hashCode/equals contract. For example, we cache new TermQuery("color","red") and then look it up with another equal new TermQuery("color","red"). The two are equal, but their identity hashes differ, so the lookup resolves to the wrong partition (or the wrong bucket within that partition's HashMap) and misses. In practice this neutered the cache.

The PR first fixes that by hashing on query.hashCode() so it matches equals(), and memoizing the composite hash to avoid recomputing it.

As part of the PR review, we also found issues where query bytes were over-counted as we were not tracking the canonical representation of the query. So now we route each query to a single partition (by the query alone rather than (segment, query)) and canonicalizes it there to one shared Query instance, so a query cached across many segments is no longer duplicated. This in turn lets the cache account a query's RAM once per distinct query instead of once per (segment, query) entry, fixing an over-count that inflated ramBytesUsed and caused premature eviction as a result.

[sgup432]

@msfroh Can you take a look at this?

[msfroh]

@sgup432 -- I was comparing this to the 10.x branch logic. I see that there, we have the top-level uniqueQueries, which keeps track of the distinct queries in the cache. When a query is (potentially) added, we replace it with the canonical representation (i.e. the one already in uniqueQueries). That's why we're able to use IdentityHashMap at the leaf level, since we're guaranteed that if the query is in the cache, then we're using the canonical instance.

Now that you've partitioned the cache at the top layer, I'm wondering if it makes sense to move the uniqueQueries map (and associated mostRecentlyUsedQueries view of the keySet()) into LRUQueryCachePartition. Then you can canonicalize the query before you create the queryCacheKey here.

Even with your change as written, if two equal queries happen to land on two different segments, the associated QueryCacheKey objects will reference different instances of the same Query object, which wastes some memory. So replacing with the canonical instance of a Query has some benefit.

[sgup432]

Even with your change as written, if two equal queries happen to land on two different segments, the associated QueryCacheKey objects will reference different instances of the same Query object, which wastes some memory. So replacing with the canonical instance of a Query has some benefit.

Even if we reintroduce(it was removed earlier) the uniqueQueries into partition level, I guess we will still have the same problem i.e. same query being held in memory for two different partitions? As the partition number is decided based on the composite key - (segment, query) for better throughput.
Until unless we route to partition by queries only, this would avoid us storing the redundant query objects in memory. But the tradeoff being a hot-query segments concentrate in one partition.

[sgup432]

If we have to replace the query with its canonical representation, I think it has to done at a global level so that each partition uses the shared canonical query. There might be cases(and I have seen many) where users send very large queries and they end up getting cached.

To have a map at a global level might introduce some regression but we can use a ConcurrentHashMap<Query, Query>.

The harder part is to remove the query from this map when this particular query is removed from all the segments/partitions. This can be done by either introducing a reference counting based mechanism like ConcurrentHashMap<Query, Tuple<Query, int>>, so when it goes to 0, we remove it. Or simply rely on GC by doing ConcurrentHashMap<Query, WeakReferences<Query>>. I believe from performance perspective both might be same though need to verify. But each has its own tradeoff, reference count based mechanism is determinstic but involves more complexity and WeakReference mechanism is more easy to implement but non-deterministic.

[sgup432]

@msfroh Maybe we can take the canonical change as part of a different PR as might be more involved?

[msfroh]

Hmm... since eviction is managed on a per-partition basis, I don't think we can reasonably pull uniqueQueries up to the root, since we won't know when things should be evicted.

Actually, looking at the implementation, I'm not convinced that we are properly accounting for query size in ramBytesUsed, since the cache entries are on the segment/query pairs. But we increment the query size contribution for each entry, even though it's probably the same query instance.

I agree that this change fixes the immediate bug, but I'm worried that there are other issues with the new cache implementation.

[sgup432]

Actually, looking at the implementation, I'm not convinced that we are properly accounting for query size in ramBytesUsed, since the cache entries are on the segment/query pairs. But we increment the query size contribution for each entry, even though it's probably the same query instance.

But considering we are not doing any canonicalization at this point, so those N entries(across segments) hold N same Query instances. So counting the query size N times is actually accurate right now?

Actually you are right. Within a single search, CachingWrapperWeight passes the same in.getQuery() instance for every leaf, so all N QueryCacheKeys reference one shared Query object. So counting the query size N times does over-counts in this case.

Hmm... since eviction is managed on a per-partition basis, I don't think we can reasonably pull uniqueQueries up to the root, since we won't know when things should be evicted.

Each partition does invoke onCacheEntryEvicted which can probably be used or something similar at a global level to handle this. Though there are other points like clear()(clearing all entries per partition) which might be tricky to handle. We can still use WeakReferences logic and rely on GC to clear those entries, it may not be a very bad idea.
OR
Route by query only instead of (segment, query). Then all of a query's segment-entries land in one partition. And have a per-partition uniqueQueries map.

[sgup432]

Route by query only instead of (segment, query). Then all of a query's segment-entries land in one partition. And have a per-partition uniqueQueries map.

I am going to try implement this approach and update this PR. I think this is the cleanest way to solve it IMO.

[sgup432]

@msfroh I have changed the logic. As mentioned above, instead of routing based on (query, segment), I am routing by query only. This simplifies a lot of things. Re-Introduced a uniqueQueries map which tracks the canonical representation of the query, also corrected the query bytes calculation.

It has also has a reference count based tracking per query. It is needed as N segments can have the same query now within same partition, and entry from the map is only remove once no segments is referencing the query.

While benchmarking the change and analyzing via JFR, I also saw some small performance optimizations opportunities like avoiding the redundant query hashCode calculation.

[sgup432]

Benchmarks:

Comparing the 16 partition vs 1 partition

┌────────────────────────────────────────┬────────┬──────────┬───────────────┬─────────┐
│                Workload                │ Metric │ V2 (16p) │ Baseline (1p) │ Speedup │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│ Uniform (distinct)                     │ mixed  │ 9.19M    │ 4.16M         │ 2.21×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│                                        │ get    │ 3.99M    │ 1.56M         │ 2.56×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│                                        │ put    │ 5.20M    │ 2.61M         │ 2.00×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│ Zipf (skewed)                          │ mixed  │ 12.53M   │ 6.86M         │ 1.83×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│                                        │ get    │ 7.07M    │ 2.53M         │ 2.80×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│                                        │ put    │ 5.45M    │ 4.33M         │ 1.26×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│ Zipf + invalidation (cleanup running)  │ get    │ 10.68M   │ 3.57M         │ 2.99×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│                                        │ put    │ 6.44M    │ 3.52M         │ 1.83×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│ Eviction-heavy                         │ mixed  │ 3.09M    │ 2.15M         │ 1.44×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│                                        │ get    │ 0.31M    │ 0.16M         │ 1.92×   │
├────────────────────────────────────────┼────────┼──────────┼───────────────┼─────────┤
│                                        │ put    │ 2.78M    │ 1.99M         │ 1.40×   │
└────────────────────────────────────────┴────────┴──────────┴───────────────┴─────────┘

Worst case for this cache:
Single hot query - as it is always routed to the same partition:

┌────────┬──────────┬───────────────┬───────┐
│ Metric │ V2 (16p) │ Baseline (1p) │ Ratio │
├────────┼──────────┼───────────────┼───────┤
│ mixed  │ 11.49M   │ 11.98M        │ 0.96× │
├────────┼──────────┼───────────────┼───────┤
│ get    │ 3.10M    │ 3.29M         │ 0.94× │
├────────┼──────────┼───────────────┼───────┤
│ put    │ 8.39M    │ 8.69M         │ 0.97× │
└────────┴──────────┴───────────────┴───────┘