statistics: stabilize TestNonLiteInitStatsWithTableIDs case (#67971)

pkg/statistics/handle/bootstrap.go

@@ -334,6 +334,9 @@ func (h *Handle) initStatsHistogramsLite(ctx context.Context, cache statstypes.S
 			break
 		}
 		h.initStatsHistograms4ChunkLite(cache, iter)
+		// The same table may continue in the next chunk. Drain LFU async admission/rejection
+		// before the next chunk reads or mutates it again.
+		cache.WaitForAsyncUpdates()
 	}
 	return nil
 }
@@ -356,7 +359,14 @@ func (h *Handle) initStatsHistograms(is infoschema.InfoSchema, cache statstypes.
 		if req.NumRows() == 0 {
 			break
 		}
+<<<<<<< HEAD
 		h.initStatsHistograms4Chunk(is, cache, iter, false)
+=======
+		h.initStatsHistograms4Chunk(is, cache, iter, isFullCache(cache, totalMemory))
+		// The same table may continue in the next chunk. Drain LFU async admission/rejection
+		// before the next chunk reads or mutates it again.
+		cache.WaitForAsyncUpdates()
+>>>>>>> 933e59fbe58 (statistics: stabilize TestNonLiteInitStatsWithTableIDs case (#67971))
 	}
 	return nil
 }
@@ -485,7 +495,14 @@ func (h *Handle) initStatsTopN(cache statstypes.StatsCache, totalMemory uint64)
 		if req.NumRows() == 0 {
 			break
 		}
+<<<<<<< HEAD
 		h.initStatsTopN4Chunk(cache, iter, totalMemory)
+=======
+		h.initStatsTopN4Chunk(cache, iter, totalMemory, tablesWithBuckets)
+		// The same table may continue in the next chunk. Drain LFU async admission/rejection
+		// before the next chunk reads or mutates it again.
+		cache.WaitForAsyncUpdates()
+>>>>>>> 933e59fbe58 (statistics: stabilize TestNonLiteInitStatsWithTableIDs case (#67971))
 	}
 	return nil
 }
@@ -726,6 +743,9 @@ func (h *Handle) initStatsBucketsByPaging(cache statstypes.StatsCache, task init
 			break
 		}
 		h.initStatsBuckets4Chunk(cache, iter)
+		// The same table may continue in the next chunk. Drain LFU async admission/rejection
+		// before the next chunk reads or mutates it again.
+		cache.WaitForAsyncUpdates()
 	}
 	return nil
 }
@@ -780,6 +800,9 @@ func (h *Handle) InitStatsLite(ctx context.Context) (err error) {
 	if err != nil {
 		return errors.Trace(err)
 	}
+	// Required: initStatsMeta adds new tables without an internal wait; histogram loading
+	// reads them immediately.
+	cache.WaitForAsyncUpdates()
 	statslogutil.StatsLogger().Info("Complete loading the stats meta in the lite mode", zap.Duration("duration", time.Since(start)))
 	start = time.Now()
 	err = h.initStatsHistogramsLite(ctx, cache)
@@ -788,7 +811,25 @@ func (h *Handle) InitStatsLite(ctx context.Context) (err error) {
 		return errors.Trace(err)
 	}
 	statslogutil.StatsLogger().Info("Complete loading the histogram in the lite mode", zap.Duration("duration", time.Since(start)))
+<<<<<<< HEAD
 	h.Replace(cache)
+=======
+	// If tableIDs is empty, it means we load all the tables' stats meta and histograms.
+	// So we can replace the global cache with the new cache.
+	if len(tableIDs) == 0 {
+		h.Replace(cache)
+	} else {
+		tables := cache.Values()
+		for _, table := range tables {
+			intest.Assert(table != nil, "table should not be nil")
+			h.Put(table.PhysicalID, table)
+		}
+		// Targeted refresh publishes refreshed tables to global LFU; make async admissions visible before returning.
+		h.StatsCache.WaitForAsyncUpdates()
+		// Do not forget to close the new cache. Otherwise it would cause the goroutine leak issue.
+		cache.Close()
+	}
+>>>>>>> 933e59fbe58 (statistics: stabilize TestNonLiteInitStatsWithTableIDs case (#67971))
 	return nil
 }
 
@@ -820,6 +861,9 @@ func (h *Handle) InitStats(ctx context.Context, is infoschema.InfoSchema) (err e
 	if err != nil {
 		return errors.Trace(err)
 	}
+	// Required: initStatsMeta adds new tables without an internal wait; histogram loading
+	// reads them immediately.
+	cache.WaitForAsyncUpdates()
 	statslogutil.StatsLogger().Info("Complete loading the stats meta", zap.Duration("duration", time.Since(start)))
 	initstats.InitStatsPercentage.Store(initStatsPercentageInterval)
 	start = time.Now()
@@ -857,7 +901,29 @@ func (h *Handle) InitStats(ctx context.Context, is infoschema.InfoSchema) (err e
 	if err != nil {
 		return errors.Trace(err)
 	}
+<<<<<<< HEAD
 	h.Replace(cache)
+=======
+	// CalcPreScalar writes tables back; drain before replacing/publishing the cache.
+	cache.WaitForAsyncUpdates()
+	statslogutil.StatsLogger().Info("Complete loading the bucket", zap.Duration("duration", time.Since(start)))
+
+	// If tableIDs is empty, it means we load all the tables' stats.
+	// So we can replace the global cache with the new cache.
+	if len(tableIDs) == 0 {
+		h.Replace(cache)
+	} else {
+		tables := cache.Values()
+		for _, table := range tables {
+			intest.Assert(table != nil, "table should not be nil")
+			h.Put(table.PhysicalID, table)
+		}
+		// Targeted refresh publishes refreshed tables to global LFU; make async admissions visible before returning.
+		h.StatsCache.WaitForAsyncUpdates()
+		// Do not forget to close the new cache. Otherwise it would cause the goroutine leak issue.
+		cache.Close()
+	}
+>>>>>>> 933e59fbe58 (statistics: stabilize TestNonLiteInitStatsWithTableIDs case (#67971))
 	return nil
 }
 

pkg/statistics/handle/cache/internal/inner.go

@@ -43,4 +43,8 @@ type StatsCacheInner interface {
 	Close()
 	// TriggerEvict triggers the cache to evict some items
 	TriggerEvict()
+	// WaitForAsyncUpdates blocks until buffered asynchronous cache writes are visible to later Get calls.
+	// Use it after adding new items when following reads depend on them. LFU/Ristretto admits
+	// non-resident items asynchronously; init stats calls this between load phases/chunks.
+	WaitForAsyncUpdates()
 }

pkg/statistics/handle/cache/internal/lfu/lfu_cache.go

@@ -29,7 +29,12 @@ import (
 	"golang.org/x/exp/rand"
 )
 
-// LFU is a LFU based on the ristretto.Cache
+// LFU is a LFU based on the ristretto.Cache.
+//
+// NOTE: In Ristretto, updates to keys already resident in the primary store become visible
+// immediately, but keys not currently resident go through a buffered admission path and are
+// applied by a background worker. So a successful Set does not guarantee immediate Get
+// visibility for a non-resident key.
 type LFU struct {
 	cache *ristretto.Cache
 	// This is a secondary cache layer used to store all tables,
@@ -87,7 +92,27 @@ func adjustMemCost(totalMemCost int64) (result int64, err error) {
 	return totalMemCost, nil
 }
 
-// Get implements statsCacheInner
+// Get reads the primary Ristretto cache first and falls back to resultKeySet.
+//
+// NOTE: We keep this order to preserve Ristretto's admission/frequency behavior.
+// That means a resident primary-cache entry wins even if resultKeySet has already been updated to
+// a newer table snapshot; until the buffered Ristretto write catches up, Get can still observe
+// that older snapshot. We are aware of this stale-read window, but it is very rare in practice, so
+// we currently document it instead of changing the read order and taking on different concurrency
+// trade-offs.
+//
+// Extreme example (the old flaky pattern from TestNonLiteInitStatsWithTableIDs):
+//  1. InitStats(tbl1) publishes tbl1 to resultKeySet, but the new primary-cache entry is still in
+//     Ristretto's buffered admission path.
+//  2. A later targeted InitStats call that includes tbl1 again reads tbl1 from resultKeySet,
+//     fills in more index state, and republishes a newer tbl1 snapshot there.
+//  3. The older tbl1 can still reach Ristretto first, so topn/buckets sees that resident older
+//     snapshot, misses the index entry, and Put() can write that stale tbl1 back into the main
+//     cache.
+//
+// For phase-by-phase or chunk-by-chunk cache construction, callers that read after writes should
+// call WaitForAsyncUpdates before the next dependent read. InitStats does this between load phases
+// and chunks so a temporary LFU cache cannot carry an older table snapshot into later loading.
 func (s *LFU) Get(tid int64) (*statistics.Table, bool) {
 	result, ok := s.cache.Get(tid)
 	if !ok {
@@ -96,7 +121,10 @@ func (s *LFU) Get(tid int64) (*statistics.Table, bool) {
 	return result.(*statistics.Table), ok
 }
 
-// Put implements statsCacheInner
+// Put publishes to resultKeySet first, then hands the item to Ristretto.
+//
+// Resident-key updates become visible there immediately; non-resident keys still go through the
+// buffered admission path.
 func (s *LFU) Put(tblID int64, tbl *statistics.Table) bool {
 	cost := tbl.MemoryUsage().TotalTrackingMemUsage()
 	s.resultKeySet.AddKeyValue(tblID, tbl)
@@ -221,9 +249,9 @@ func (s *LFU) SetCapacity(maxCost int64) {
 	metrics.CostGauge.Set(float64(s.Cost()))
 }
 
-// wait blocks until all buffered writes have been applied. This ensures a call to Set()
-// will be visible to future calls to Get(). it is only used for test.
-func (s *LFU) wait() {
+// WaitForAsyncUpdates blocks until all buffered writes have been applied. This ensures a call to
+// Put is visible to future calls to Get.
+func (s *LFU) WaitForAsyncUpdates() {
 	s.cache.Wait()
 }
 

pkg/statistics/handle/cache/internal/lfu/lfu_cache_test.go

@@ -36,12 +36,12 @@ func TestLFUPutGetDel(t *testing.T) {
 	mockTable := testutil.NewMockStatisticsTable(1, 1, true, false, false)
 	mockTableID := int64(1)
 	lfu.Put(mockTableID, mockTable)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	lfu.Del(mockTableID)
 	v, ok := lfu.Get(mockTableID)
 	require.False(t, ok)
 	require.Nil(t, v)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, uint64(lfu.Cost()), lfu.metrics().CostAdded()-lfu.metrics().CostEvicted())
 	require.Equal(t, 0, len(lfu.Values()))
 }
@@ -58,15 +58,15 @@ func TestLFUFreshMemUsage(t *testing.T) {
 	lfu.Put(int64(1), t1)
 	lfu.Put(int64(2), t2)
 	lfu.Put(int64(3), t3)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, lfu.Cost(), 6*mockCMSMemoryUsage+6*mockCMSMemoryUsage)
 	t4 := testutil.NewMockStatisticsTable(2, 1, true, false, false)
 	lfu.Put(int64(1), t4)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, lfu.Cost(), 7*mockCMSMemoryUsage+6*mockCMSMemoryUsage)
 	t5 := testutil.NewMockStatisticsTable(2, 2, true, false, false)
 	lfu.Put(int64(1), t5)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, lfu.Cost(), 7*mockCMSMemoryUsage+7*mockCMSMemoryUsage)
 
 	t6 := testutil.NewMockStatisticsTable(1, 2, true, false, false)
@@ -76,7 +76,7 @@ func TestLFUFreshMemUsage(t *testing.T) {
 	t7 := testutil.NewMockStatisticsTable(1, 1, true, false, false)
 	lfu.Put(int64(1), t7)
 	require.Equal(t, lfu.Cost(), 6*mockCMSMemoryUsage+6*mockCMSMemoryUsage)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, uint64(lfu.Cost()), lfu.metrics().CostAdded()-lfu.metrics().CostEvicted())
 }
 
@@ -88,7 +88,7 @@ func TestLFUPutTooBig(t *testing.T) {
 	lfu.Put(int64(1), mockTable)
 	_, ok := lfu.Get(int64(1))
 	require.True(t, ok)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, uint64(lfu.Cost()), lfu.metrics().CostAdded()-lfu.metrics().CostEvicted())
 }
 
@@ -102,14 +102,14 @@ func TestCacheLen(t *testing.T) {
 	t2 := testutil.NewMockStatisticsTable(1, 1, true, false, false)
 	// put t2, t1 should be evicted 2 items and still exists in the list
 	lfu.Put(int64(2), t2)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, lfu.Len(), 2)
 	require.Equal(t, uint64(8), lfu.metrics().CostAdded()-lfu.metrics().CostEvicted())
 
 	// put t3, t1/t2 should be evicted all items. but t1/t2 still exists in the list
 	t3 := testutil.NewMockStatisticsTable(2, 1, true, false, false)
 	lfu.Put(int64(3), t3)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, lfu.Len(), 3)
 	require.Equal(t, uint64(12), lfu.metrics().CostAdded()-lfu.metrics().CostEvicted())
 }
@@ -133,7 +133,7 @@ func TestLFUCachePutGetWithManyConcurrency(t *testing.T) {
 		}(i)
 	}
 	wg.Wait()
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, lfu.Len(), 1000)
 	require.Equal(t, uint64(lfu.Cost()), lfu.metrics().CostAdded()-lfu.metrics().CostEvicted())
 	require.Equal(t, 1000, len(lfu.Values()))
@@ -164,7 +164,7 @@ func TestLFUCachePutGetWithManyConcurrency2(t *testing.T) {
 		}()
 	}
 	wg.Wait()
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, uint64(lfu.Cost()), lfu.metrics().CostAdded()-lfu.metrics().CostEvicted())
 	require.Equal(t, 1000, len(lfu.Values()))
 }
@@ -202,7 +202,7 @@ func TestLFUCachePutGetWithManyConcurrencyAndSmallConcurrency(t *testing.T) {
 		}()
 	}
 	wg.Wait()
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	v, ok := lfu.Get(rand.Int63n(50))
 	require.True(t, ok)
 	v.ForEachColumnImmutable(func(_ int64, c *statistics.Column) bool {
@@ -245,14 +245,14 @@ func TestLFUReject(t *testing.T) {
 	t1 := testutil.NewMockStatisticsTable(2, 1, true, false, false)
 	require.Equal(t, 2*mockCMSMemoryUsage+mockCMSMemoryUsage, t1.MemoryUsage().TotalTrackingMemUsage())
 	lfu.Put(1, t1)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, lfu.Cost(), 2*mockCMSMemoryUsage+mockCMSMemoryUsage)
 
 	lfu.SetCapacity(2*mockCMSMemoryUsage + mockCMSMemoryUsage - 1)
 
 	t2 := testutil.NewMockStatisticsTable(2, 1, true, false, false)
 	require.True(t, lfu.Put(2, t2))
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	time.Sleep(3 * time.Second)
 	require.Equal(t, int64(0), lfu.Cost())
 	require.Len(t, lfu.Values(), 2)
@@ -275,7 +275,7 @@ func TestMemoryControl(t *testing.T) {
 	t1 := testutil.NewMockStatisticsTable(2, 1, true, false, false)
 	require.Equal(t, 2*mockCMSMemoryUsage+mockCMSMemoryUsage, t1.MemoryUsage().TotalTrackingMemUsage())
 	lfu.Put(1, t1)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 
 	for i := 2; i <= 1000; i++ {
 		t1 := testutil.NewMockStatisticsTable(2, 1, true, false, false)
@@ -286,19 +286,19 @@ func TestMemoryControl(t *testing.T) {
 
 	for i := 1000; i > 990; i-- {
 		lfu.SetCapacity(int64(i-1) * (2*mockCMSMemoryUsage + mockCMSMemoryUsage))
-		lfu.wait()
+		lfu.WaitForAsyncUpdates()
 		require.Equal(t, int64(i-1)*(2*mockCMSMemoryUsage+mockCMSMemoryUsage), lfu.Cost())
 	}
 	for i := 990; i > 100; i = i - 100 {
 		lfu.SetCapacity(int64(i-1) * (2*mockCMSMemoryUsage + mockCMSMemoryUsage))
-		lfu.wait()
+		lfu.WaitForAsyncUpdates()
 		require.Equal(t, int64(i-1)*(2*mockCMSMemoryUsage+mockCMSMemoryUsage), lfu.Cost())
 	}
 	lfu.SetCapacity(int64(10) * (2*mockCMSMemoryUsage + mockCMSMemoryUsage))
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, int64(10)*(2*mockCMSMemoryUsage+mockCMSMemoryUsage), lfu.Cost())
 	lfu.SetCapacity(0)
-	lfu.wait()
+	lfu.WaitForAsyncUpdates()
 	require.Equal(t, int64(10)*(2*mockCMSMemoryUsage+mockCMSMemoryUsage), lfu.Cost())
 }
 

pkg/statistics/handle/cache/internal/mapcache/map_cache.go

@@ -134,3 +134,6 @@ func (*MapCache) Close() {}
 
 // TriggerEvict implements statsCacheInner
 func (*MapCache) TriggerEvict() {}
+
+// WaitForAsyncUpdates implements statsCacheInner.
+func (*MapCache) WaitForAsyncUpdates() {}

pkg/statistics/handle/cache/statscache.go

@@ -336,6 +336,11 @@ func (s *StatsCacheImpl) TriggerEvict() {
 	s.Load().TriggerEvict()
 }
 
+// WaitForAsyncUpdates blocks until buffered asynchronous cache writes are visible to later Get calls.
+func (s *StatsCacheImpl) WaitForAsyncUpdates() {
+	s.Load().WaitForAsyncUpdates()
+}
+
 // MaxTableStatsVersion returns the version of the current cache, which is defined as
 // the max table stats version the cache has in its lifecycle.
 func (s *StatsCacheImpl) MaxTableStatsVersion() uint64 {

pkg/statistics/handle/cache/statscache_test.go

@@ -196,3 +196,7 @@ func (*mockStatsCacheInner) Close() {
 func (*mockStatsCacheInner) TriggerEvict() {
 	panic("not implemented")
 }
+
+func (*mockStatsCacheInner) WaitForAsyncUpdates() {
+	panic("not implemented")
+}

pkg/statistics/handle/cache/statscacheinner.go

@@ -188,3 +188,8 @@ func (sc *StatsCache) Update(tables []*statistics.Table, deletedIDs []int64, ski
 func (sc *StatsCache) TriggerEvict() {
 	sc.c.TriggerEvict()
 }
+
+// WaitForAsyncUpdates blocks until buffered asynchronous cache writes are visible to later Get calls.
+func (sc *StatsCache) WaitForAsyncUpdates() {
+	sc.c.WaitForAsyncUpdates()
+}

pkg/statistics/handle/types/interfaces.go

@@ -264,6 +264,11 @@ type StatsCache interface {
 
 	// TriggerEvict triggers the cache to evict some items
 	TriggerEvict()
+
+	// WaitForAsyncUpdates blocks until buffered asynchronous cache writes are visible to later Get calls.
+	// Use it after adding new items when following reads depend on them. LFU/Ristretto admits
+	// non-resident items asynchronously; init stats calls this between load phases/chunks.
+	WaitForAsyncUpdates()
 }
 
 // StatsLockTable is the table info of which will be locked.