statistics: handle deleted tables correctly in the PQ

Signed-off-by: Rustin170506 <[email protected]>

Author: 0xPoe

pkg/statistics/handle/autoanalyze/priorityqueue/analysis_job_factory.go

@@ -57,6 +57,9 @@ func (f *AnalysisJobFactory) CreateNonPartitionedTableAnalysisJob(
 	tblInfo *model.TableInfo,
 	tblStats *statistics.Table,
 ) AnalysisJob {
+	if tblStats == nil {
+		return nil
+	}
 	if !tblStats.IsEligibleForAnalysis() {
 		return nil
 	}
@@ -92,6 +95,9 @@ func (f *AnalysisJobFactory) CreateStaticPartitionAnalysisJob(
 	partitionID int64,
 	partitionStats *statistics.Table,
 ) AnalysisJob {
+	if partitionStats == nil {
+		return nil
+	}
 	if !partitionStats.IsEligibleForAnalysis() {
 		return nil
 	}
@@ -128,6 +134,9 @@ func (f *AnalysisJobFactory) CreateDynamicPartitionedTableAnalysisJob(
 	globalTblStats *statistics.Table,
 	partitionStats map[PartitionIDAndName]*statistics.Table,
 ) AnalysisJob {
+	if globalTblStats == nil {
+		return nil
+	}
 	if !globalTblStats.IsEligibleForAnalysis() {
 		return nil
 	}

pkg/statistics/handle/autoanalyze/priorityqueue/queue.go

@@ -644,14 +644,16 @@ func (pq *AnalysisPriorityQueue) RefreshLastAnalysisDuration() {
 					zap.Int64("tableID", job.GetTableID()),
 					zap.String("job", job.String()),
 				)
-				// TODO: Remove this after handling the DDL event.
+				// Delete the job from the queue since its table is missing. This is a safeguard -
+				// DDL events should have already cleaned up jobs for dropped tables.
 				err := pq.syncFields.inner.delete(job)
 				if err != nil {
 					statslogutil.StatsLogger().Error("Failed to delete job from priority queue",
 						zap.Error(err),
 						zap.String("job", job.String()),
 					)
 				}
+				continue
 			}
 			indicators.LastAnalysisDuration = jobFactory.GetTableLastAnalyzeDuration(tableStats)
 			job.SetIndicators(indicators)

pkg/statistics/handle/autoanalyze/priorityqueue/queue_test.go

@@ -19,6 +19,7 @@ import (
 	"testing"
 	"time"
 
+	"github.com/pingcap/tidb/pkg/meta/model"
 	pmodel "github.com/pingcap/tidb/pkg/parser/model"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/statistics"
@@ -608,3 +609,47 @@ func TestPQCanBeClosedAndReInitialized(t *testing.T) {
 	// Check if the priority queue is initialized.
 	require.True(t, pq.IsInitialized())
 }
+
+func TestPQHandlesTableDeletionGracefully(t *testing.T) {
+	store, dom := testkit.CreateMockStoreAndDomain(t)
+	handle := dom.StatsHandle()
+
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t1 (a int)")
+	tk.MustExec("insert into t1 values (1)")
+	statistics.AutoAnalyzeMinCnt = 0
+	defer func() {
+		statistics.AutoAnalyzeMinCnt = 1000
+	}()
+
+	ctx := context.Background()
+	require.NoError(t, handle.DumpStatsDeltaToKV(true))
+	require.NoError(t, handle.Update(ctx, dom.InfoSchema()))
+	pq := priorityqueue.NewAnalysisPriorityQueue(handle)
+	defer pq.Close()
+	require.NoError(t, pq.Initialize())
+
+	// Check the priority queue is not empty.
+	l, err := pq.Len()
+	require.NoError(t, err)
+	require.NotEqual(t, 0, l)
+
+	tbl, err := dom.InfoSchema().TableByName(ctx, pmodel.NewCIStr("test"), pmodel.NewCIStr("t1"))
+	require.NoError(t, err)
+
+	// Drop the table and mock the table stats is removed from the cache.
+	tk.MustExec("drop table t1")
+	deleteEvent := findEvent(handle.DDLEventCh(), model.ActionDropTable)
+	require.NotNil(t, deleteEvent)
+	require.NoError(t, handle.HandleDDLEvent(deleteEvent))
+	require.NoError(t, handle.Update(ctx, dom.InfoSchema()))
+
+	// Make sure handle.Get() returns false.
+	_, ok := handle.Get(tbl.Meta().ID)
+	require.False(t, ok)
+
+	require.NotPanics(t, func() {
+		pq.RefreshLastAnalysisDuration()
+	})
+}

pkg/statistics/handle/autoanalyze/refresher/refresher.go

@@ -262,5 +262,8 @@ func (*Refresher) OnBecomeOwner() {
 // OnRetireOwner is used to handle the event when the current TiDB instance retires from being the stats owner.
 func (r *Refresher) OnRetireOwner() {
 	// Stop the worker and close the queue.
+	// Note: we have to guarantee that the worker is stopped before closing the queue.
+	// Otherwise, the worker may still access the queue after it is closed.
+	r.worker.Stop()
 	r.jobs.Close()
 }