From dece884327d6211a4e8d9d3413580e8857abe1ac Mon Sep 17 00:00:00 2001
From: Rustin <tech@rustin.me>
Date: Thu, 7 Nov 2024 18:43:06 +0800
Subject: [PATCH] statistics: make sure the PQ can be re-initialized (#57194)

ref pingcap/tidb#55906
---
 .../handle/autoanalyze/priorityqueue/queue.go | 10 ++++++++++
 .../autoanalyze/priorityqueue/queue_test.go   | 20 +++++++++++++++++++
 2 files changed, 30 insertions(+)

diff --git a/pkg/statistics/handle/autoanalyze/priorityqueue/queue.go b/pkg/statistics/handle/autoanalyze/priorityqueue/queue.go
index daeb03496c1d4..a65d16c1f42c4 100644
--- a/pkg/statistics/handle/autoanalyze/priorityqueue/queue.go
+++ b/pkg/statistics/handle/autoanalyze/priorityqueue/queue.go
@@ -844,4 +844,14 @@ func (pq *AnalysisPriorityQueue) Close() {
 		pq.syncFields.cancel()
 	}
 	pq.wg.Wait()
+
+	// Reset the initialized flag to allow the priority queue to be closed and re-initialized.
+	pq.syncFields.initialized = false
+	// The rest fields will be reset when the priority queue is initialized.
+	// But we do it here for double safety.
+	pq.syncFields.inner = nil
+	pq.syncFields.runningJobs = nil
+	pq.syncFields.mustRetryJobs = nil
+	pq.syncFields.lastDMLUpdateFetchTimestamp = 0
+	pq.syncFields.cancel = nil
 }
diff --git a/pkg/statistics/handle/autoanalyze/priorityqueue/queue_test.go b/pkg/statistics/handle/autoanalyze/priorityqueue/queue_test.go
index af6fc37e93a6a..2cef81d87a761 100644
--- a/pkg/statistics/handle/autoanalyze/priorityqueue/queue_test.go
+++ b/pkg/statistics/handle/autoanalyze/priorityqueue/queue_test.go
@@ -582,3 +582,23 @@ func TestProcessDMLChangesWithLockedPartitionsAndStaticPruneMode(t *testing.T) {
 	pid = tbl.Meta().Partition.Definitions[0].ID
 	require.Equal(t, pid, job.GetTableID())
 }
+
+func TestPQCanBeClosedAndReInitialized(t *testing.T) {
+	_, dom := testkit.CreateMockStoreAndDomain(t)
+	handle := dom.StatsHandle()
+	pq := priorityqueue.NewAnalysisPriorityQueue(handle)
+	defer pq.Close()
+	require.NoError(t, pq.Initialize())
+
+	// Close the priority queue.
+	pq.Close()
+
+	// Check if the priority queue is closed.
+	require.False(t, pq.IsInitialized())
+
+	// Re-initialize the priority queue.
+	require.NoError(t, pq.Initialize())
+
+	// Check if the priority queue is initialized.
+	require.True(t, pq.IsInitialized())
+}