br/streamhelper: fix panic while removing task (#50869) (#51054)

close #50839

br/pkg/streamhelper/BUILD.bazel

@@ -68,7 +68,7 @@ go_test(
     ],
     flaky = True,
     race = "on",
-    shard_count = 21,
+    shard_count = 22,
     deps = [
         ":streamhelper",
         "//br/pkg/errors",

br/pkg/streamhelper/advancer.go

@@ -78,6 +78,22 @@ type CheckpointAdvancer struct {
 	subscriberMu sync.Mutex
 }
 
+// HasTask returns whether the advancer has been bound to a task.
+func (c *CheckpointAdvancer) HasTask() bool {
+	c.taskMu.Lock()
+	defer c.taskMu.Unlock()
+
+	return c.task != nil
+}
+
+// HasSubscriber returns whether the advancer is associated with a subscriber.
+func (c *CheckpointAdvancer) HasSubscribion() bool {
+	c.subscriberMu.Lock()
+	defer c.subscriberMu.Unlock()
+
+	return c.subscriber != nil && len(c.subscriber.subscriptions) > 0
+}
+
 // checkpoint represents the TS with specific range.
 // it's only used in advancer.go.
 type checkpoint struct {
@@ -363,6 +379,12 @@ func (c *CheckpointAdvancer) StartTaskListener(ctx context.Context) {
 	}()
 }
 
+func (c *CheckpointAdvancer) setCheckpoints(cps *spans.ValueSortedFull) {
+	c.checkpointsMu.Lock()
+	c.checkpoints = cps
+	c.checkpointsMu.Unlock()
+}
+
 func (c *CheckpointAdvancer) onTaskEvent(ctx context.Context, e TaskEvent) error {
 	c.taskMu.Lock()
 	defer c.taskMu.Unlock()
@@ -371,20 +393,20 @@ func (c *CheckpointAdvancer) onTaskEvent(ctx context.Context, e TaskEvent) error
 		utils.LogBackupTaskCountInc()
 		c.task = e.Info
 		c.taskRange = spans.Collapse(len(e.Ranges), func(i int) kv.KeyRange { return e.Ranges[i] })
-		c.checkpoints = spans.Sorted(spans.NewFullWith(e.Ranges, 0))
+		c.setCheckpoints(spans.Sorted(spans.NewFullWith(e.Ranges, 0)))
 		c.lastCheckpoint = newCheckpointWithTS(e.Info.StartTs)
 		log.Info("added event", zap.Stringer("task", e.Info),
 			zap.Stringer("ranges", logutil.StringifyKeys(c.taskRange)))
 	case EventDel:
 		utils.LogBackupTaskCountDec()
 		c.task = nil
 		c.taskRange = nil
-		c.checkpoints = nil
 		// This would be synced by `taskMu`, perhaps we'd better rename that to `tickMu`.
 		// Do the null check because some of test cases won't equip the advancer with subscriber.
 		if c.subscriber != nil {
 			c.subscriber.Clear()
 		}
+		c.setCheckpoints(nil)
 		if err := c.env.ClearV3GlobalCheckpointForTask(ctx, e.Name); err != nil {
 			log.Warn("failed to clear global checkpoint", logutil.ShortError(err))
 		}
@@ -444,8 +466,10 @@ func (c *CheckpointAdvancer) SpawnSubscriptionHandler(ctx context.Context) {
 	defer c.subscriberMu.Unlock()
 	c.subscriber = NewSubscriber(c.env, c.env, WithMasterContext(ctx))
 	es := c.subscriber.Events()
+	log.Info("Subscription handler spawned.", zap.String("category", "log backup subscription manager"))
 
 	go func() {
+		defer utils.CatchAndLogPanic()
 		for {
 			select {
 			case <-ctx.Done():
@@ -454,12 +478,18 @@ func (c *CheckpointAdvancer) SpawnSubscriptionHandler(ctx context.Context) {
 				if !ok {
 					return
 				}
-				c.checkpointsMu.Lock()
-				log.Debug("Accepting region flush event.",
-					zap.Stringer("range", logutil.StringifyRange(event.Key)),
-					zap.Uint64("checkpoint", event.Value))
-				c.checkpoints.Merge(event)
-				c.checkpointsMu.Unlock()
+				failpoint.Inject("subscription-handler-loop", func() {})
+				c.WithCheckpoints(func(vsf *spans.ValueSortedFull) {
+					if vsf == nil {
+						log.Warn("Span tree not found, perhaps stale event of removed tasks.",
+							zap.String("category", "log backup subscription manager"))
+						return
+					}
+					log.Debug("Accepting region flush event.",
+						zap.Stringer("range", logutil.StringifyRange(event.Key)),
+						zap.Uint64("checkpoint", event.Value))
+					vsf.Merge(event)
+				})
 			}
 		}
 	}()

br/pkg/streamhelper/advancer_test.go

@@ -414,3 +414,29 @@ func TestOwnerDropped(t *testing.T) {
 		require.Equal(t, vsf.MinValue(), cp)
 	})
 }
+
+// TestRemoveTaskAndFlush tests the bug has been described in #50839.
+func TestRemoveTaskAndFlush(t *testing.T) {
+	log.SetLevel(zapcore.DebugLevel)
+	ctx := context.Background()
+	c := createFakeCluster(t, 4, true)
+	installSubscribeSupport(c)
+	env := &testEnv{
+		fakeCluster: c,
+		testCtx:     t,
+	}
+	adv := streamhelper.NewCheckpointAdvancer(env)
+	adv.StartTaskListener(ctx)
+	adv.SpawnSubscriptionHandler(ctx)
+	require.NoError(t, adv.OnTick(ctx))
+	require.NoError(t, failpoint.Enable("github.com/pingcap/tidb/br/pkg/streamhelper/subscription-handler-loop", "pause"))
+	c.flushAll()
+	env.unregisterTask()
+	require.Eventually(t, func() bool {
+		return !adv.HasTask()
+	}, 10*time.Second, 100*time.Millisecond)
+	require.NoError(t, failpoint.Disable("github.com/pingcap/tidb/br/pkg/streamhelper/subscription-handler-loop"))
+	require.Eventually(t, func() bool {
+		return !adv.HasSubscribion()
+	}, 10*time.Second, 100*time.Millisecond)
+}

br/pkg/streamhelper/basic_lib_for_test.go

@@ -590,6 +590,7 @@ type testEnv struct {
 	checkpoint uint64
 	testCtx    *testing.T
 	ranges     []kv.KeyRange
+	taskCh     chan<- streamhelper.TaskEvent
 
 	resolveLocks func([]*txnlock.Lock, *tikv.KeyLocation) (*tikv.KeyLocation, error)
 
@@ -611,6 +612,7 @@ func (t *testEnv) Begin(ctx context.Context, ch chan<- streamhelper.TaskEvent) e
 		Ranges: rngs,
 	}
 	ch <- tsk
+	t.taskCh = ch
 	return nil
 }
 
@@ -640,6 +642,13 @@ func (t *testEnv) getCheckpoint() uint64 {
 	return t.checkpoint
 }
 
+func (t *testEnv) unregisterTask() {
+	t.taskCh <- streamhelper.TaskEvent{
+		Type: streamhelper.EventDel,
+		Name: "whole",
+	}
+}
+
 func (t *testEnv) ScanLocksInOneRegion(bo *tikv.Backoffer, key []byte, maxVersion uint64, limit uint32) ([]*txnlock.Lock, *tikv.KeyLocation, error) {
 	for _, r := range t.regions {
 		if len(r.locks) != 0 {

br/pkg/utils/worker.go

@@ -159,3 +159,17 @@ func BuildWorkerTokenChannel(size uint) chan struct{} {
 	}
 	return ch
 }
+
+// CatchAndLogPanic recovers when the execution get panicked, and log the panic.
+// generally, this would be used with `defer`, like:
+//
+//	func foo() {
+//	  defer utils.CatchAndLogPanic()
+//	  maybePanic()
+//	}
+func CatchAndLogPanic() {
+	item := recover()
+	if item != nil {
+		log.Warn("CatchAndLogPanic: panicked, but ignored.", zap.StackSkip("stack", 1), zap.Any("panic", item))
+	}
+}

tablecodec/BUILD.bazel

@@ -32,6 +32,7 @@ go_test(
         "main_test.go",
         "tablecodec_test.go",
     ],
+    data = glob(["testdata/**"]),
     embed = [":tablecodec"],
     flaky = True,
     deps = [