[bug] logtail/stats: do not unsubscribe table before it finished waited (#19682)

do not unsubscribe table before it finished waited, otherwise,
it will never stop waiting for the table's first logtail as it is subscribed.

Approved by: @triump2020, @sukki37

Author: volgariver6

pkg/vm/engine/disttae/engine.go

@@ -738,6 +738,10 @@ func (e *Engine) cleanMemoryTableWithTable(dbId, tblId uint64) {
 	logutil.Debugf("clean memory table of tbl[dbId: %d, tblId: %d]", dbId, tblId)
 }
 
+func (e *Engine) safeToUnsubscribe(tid uint64) bool {
+	return e.globalStats.safeToUnsubscribe(tid)
+}
+
 func (e *Engine) PushClient() *PushClient {
 	return &e.pClient
 }

pkg/vm/engine/disttae/logtail_consumer.go

@@ -947,6 +947,10 @@ func (c *PushClient) unusedTableGCTicker(ctx context.Context) {
 						// never unsubscribe the mo_databases, mo_tables, mo_columns.
 						continue
 					}
+					if !c.eng.safeToUnsubscribe(k) {
+						logutil.Infof("%s table [%d-%d] is not safe to unsubscribe", logTag, v.DBID, k)
+						continue
+					}
 					if !v.LatestTime.After(shouldClean) {
 						if v.SubState != Subscribed {
 							continue

pkg/vm/engine/disttae/logtail_consumer_test.go

@@ -32,6 +32,7 @@ func TestSubscribedTable(t *testing.T) {
 		partitions: make(map[[2]uint64]*logtailreplay.Partition),
 		globalStats: &GlobalStats{
 			logtailUpdate: newLogtailUpdate(),
+			waitKeeper:    newWaitKeeper(),
 		},
 	}
 	require.Equal(t, 0, len(subscribeRecord.m))

pkg/vm/engine/disttae/stats.go

@@ -44,6 +44,37 @@ var (
 	MinUpdateInterval = time.Second * 15
 )
 
+// waitKeeper is used to mark the table has finished waited,
+// only after which, the table can be unsubscribed.
+type waitKeeper struct {
+	sync.Mutex
+	records map[uint64]struct{}
+}
+
+func newWaitKeeper() *waitKeeper {
+	return &waitKeeper{
+		records: make(map[uint64]struct{}),
+	}
+}
+
+func (w *waitKeeper) reset() {
+	w.Lock()
+	defer w.Unlock()
+	w.records = make(map[uint64]struct{})
+}
+
+func (w *waitKeeper) add(tid uint64) {
+	w.Lock()
+	defer w.Unlock()
+	w.records[tid] = struct{}{}
+}
+
+func (w *waitKeeper) del(tid uint64) {
+	w.Lock()
+	defer w.Unlock()
+	delete(w.records, tid)
+}
+
 type updateStatsRequest struct {
 	// statsInfo is the field which is to update.
 	statsInfo *pb.StatsInfo
@@ -150,6 +181,10 @@ type GlobalStats struct {
 		statsInfoMap map[pb.StatsInfoKey]*pb.StatsInfo
 	}
 
+	// waitKeeper is used to make sure the table is safe to unsubscribe.
+	// Only when the table is finished waited, it can be unsubscribed safely.
+	waitKeeper *waitKeeper
+
 	// updateWorkerFactor is the times of CPU number of this node
 	// to start update worker. Default is 8.
 	updateWorkerFactor int
@@ -171,6 +206,7 @@ func NewGlobalStats(
 		logtailUpdate:       newLogtailUpdate(),
 		tableLogtailCounter: make(map[pb.StatsInfoKey]int64),
 		KeyRouter:           keyRouter,
+		waitKeeper:          newWaitKeeper(),
 	}
 	s.updatingMu.updating = make(map[pb.StatsInfoKey]*updateRecord)
 	s.mu.statsInfoMap = make(map[pb.StatsInfoKey]*pb.StatsInfo)
@@ -265,20 +301,34 @@ func (gs *GlobalStats) Get(ctx context.Context, key pb.StatsInfoKey, sync bool)
 }
 
 func (gs *GlobalStats) RemoveTid(tid uint64) {
+	gs.waitKeeper.del(tid)
+
 	gs.logtailUpdate.mu.Lock()
 	defer gs.logtailUpdate.mu.Unlock()
 	delete(gs.logtailUpdate.mu.updated, tid)
 }
 
 // clearTables clears the tables in the map if there are any tables in it.
 func (gs *GlobalStats) clearTables() {
+	// clear all the waiters in the keeper.
+	gs.waitKeeper.reset()
+
 	gs.logtailUpdate.mu.Lock()
 	defer gs.logtailUpdate.mu.Unlock()
 	if len(gs.logtailUpdate.mu.updated) > 0 {
 		gs.logtailUpdate.mu.updated = make(map[uint64]struct{})
 	}
 }
 
+func (gs *GlobalStats) safeToUnsubscribe(tid uint64) bool {
+	gs.waitKeeper.Lock()
+	defer gs.waitKeeper.Unlock()
+	if _, ok := gs.waitKeeper.records[tid]; ok {
+		return true
+	}
+	return false
+}
+
 func (gs *GlobalStats) enqueue(tail *logtail.TableLogtail) {
 	select {
 	case gs.tailC <- tail:
@@ -380,6 +430,8 @@ func (gs *GlobalStats) notifyLogtailUpdate(tid uint64) {
 }
 
 func (gs *GlobalStats) waitLogtailUpdated(tid uint64) {
+	defer gs.waitKeeper.add(tid)
+
 	// If the tid is less than reserved, return immediately.
 	if tid < catalog.MO_RESERVED_MAX {
 		return

pkg/vm/engine/disttae/stats_test.go

@@ -87,3 +87,32 @@ func TestGlobalStats_ClearTables(t *testing.T) {
 	gs.clearTables()
 	assert.Equal(t, 0, len(gs.logtailUpdate.mu.updated))
 }
+
+func TestWaitKeeper(t *testing.T) {
+	defer leaktest.AfterTest(t)()
+
+	var tid uint64 = 1
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+	gs := NewGlobalStats(ctx, nil, nil)
+	assert.False(t, gs.safeToUnsubscribe(tid))
+
+	w := gs.waitKeeper
+	w.add(tid)
+	_, ok := w.records[tid]
+	assert.True(t, ok)
+	assert.True(t, gs.safeToUnsubscribe(tid))
+
+	gs.RemoveTid(tid)
+	_, ok = w.records[tid]
+	assert.False(t, ok)
+	assert.False(t, gs.safeToUnsubscribe(tid))
+
+	w.add(tid)
+	_, ok = w.records[tid]
+	assert.True(t, ok)
+	gs.clearTables()
+	_, ok = w.records[tid]
+	assert.False(t, ok)
+	assert.False(t, gs.safeToUnsubscribe(tid))
+}