ddl notifier: use pessimistic txn and fix updating memory state too early

pkg/ddl/notifier/subscribe.go

@@ -268,15 +268,20 @@ func (n *DDLNotifier) processEventForHandler(
 	if (change.processedByFlag & (1 << handlerID)) != 0 {
 		return nil
 	}
+	newFlag := change.processedByFlag | (1 << handlerID)
 
-	if err = session.Begin(ctx); err != nil {
+	if err = session.BeginPessimistic(ctx); err != nil {
 		return errors.Trace(err)
 	}
 	defer func() {
-		if err == nil {
-			err = errors.Trace(session.Commit(ctx))
-		} else {
+		if err != nil {
 			session.Rollback()
+			return
+		}
+
+		err = errors.Trace(session.Commit(ctx))
+		if err == nil {
+			change.processedByFlag = newFlag
 		}
 	}()
 
@@ -293,7 +298,6 @@ func (n *DDLNotifier) processEventForHandler(
 			zap.Duration("duration", time.Since(now)))
 	}
 
-	newFlag := change.processedByFlag | (1 << handlerID)
 	if err = n.store.UpdateProcessed(
 		ctx,
 		session,
@@ -303,7 +307,6 @@ func (n *DDLNotifier) processEventForHandler(
 	); err != nil {
 		return errors.Trace(err)
 	}
-	change.processedByFlag = newFlag
 
 	return nil
 }

pkg/ddl/notifier/testkit_test.go

@@ -489,3 +489,101 @@ func TestBeginTwice(t *testing.T) {
 	require.NoError(t, err)
 	require.NotContains(t, string(content), "context provider not set")
 }
+
+func TestHandlersSeePessimisticTxnError(t *testing.T) {
+	// 1. One always fails
+	// 2. One always succeeds
+	// Make sure events don't get lost after the second handler succeeds.
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("USE test")
+	tk.MustExec("DROP TABLE IF EXISTS " + ddl.NotifierTableName)
+	tk.MustExec(ddl.NotifierTableSQL)
+	ctx := context.Background()
+	s := notifier.OpenTableStore("test", ddl.NotifierTableName)
+	sessionPool := util.NewSessionPool(
+		4,
+		func() (pools.Resource, error) {
+			return testkit.NewTestKit(t, store).Session(), nil
+		},
+		nil,
+		nil,
+	)
+	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
+	// Always fails
+	failHandler := func(_ context.Context, sctx sessionctx.Context, _ *notifier.SchemaChangeEvent) error {
+		// Mock a duplicate key error
+		_, err := sctx.GetSQLExecutor().Execute(ctx, "INSERT INTO test."+ddl.NotifierTableName+" VALUES(1, -1, 'some', 0)")
+		return err
+	}
+	// Always succeeds
+	successHandler := func(context.Context, sessionctx.Context, *notifier.SchemaChangeEvent) error {
+		return nil
+	}
+	n.RegisterHandler(2, successHandler)
+	n.RegisterHandler(1, failHandler)
+	n.OnBecomeOwner()
+	tk2 := testkit.NewTestKit(t, store)
+	se := sess.NewSession(tk2.Session())
+	event1 := notifier.NewCreateTableEvent(&model.TableInfo{ID: 1000, Name: ast.NewCIStr("t1")})
+	err := notifier.PubSchemeChangeToStore(ctx, se, 1, -1, event1, s)
+	require.NoError(t, err)
+	require.Never(t, func() bool {
+		changes := make([]*notifier.SchemaChange, 8)
+		result, closeFn := s.List(ctx, se)
+		count, err2 := result.Read(changes)
+		require.NoError(t, err2)
+		closeFn()
+		return count == 0
+	}, time.Second, 50*time.Millisecond)
+}
+
+func TestCommitFailed(t *testing.T) {
+	// Make sure events don't get lost if internal txn commit failed.
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("USE test")
+	tk.MustExec("set global tidb_enable_metadata_lock=0")
+	tk.MustExec("DROP TABLE IF EXISTS " + ddl.NotifierTableName)
+	tk.MustExec(ddl.NotifierTableSQL)
+	tk.MustExec("CREATE TABLE subscribe_table (id INT PRIMARY KEY, c INT)")
+	tk.MustExec("INSERT INTO subscribe_table VALUES (1, 1)")
+
+	ctx := context.Background()
+	s := notifier.OpenTableStore("test", ddl.NotifierTableName)
+	sessionPool := util.NewSessionPool(
+		4,
+		func() (pools.Resource, error) {
+			return testkit.NewTestKit(t, store).Session(), nil
+		},
+		nil,
+		nil,
+	)
+	n := notifier.NewDDLNotifier(sessionPool, s, 50*time.Millisecond)
+	handler := func(_ context.Context, sctx sessionctx.Context, _ *notifier.SchemaChangeEvent) error {
+		// pessimistic + DDL will cause an "infoschema is changed" error at commit time.
+		_, err := sctx.GetSQLExecutor().Execute(
+			ctx, "UPDATE test.subscribe_table SET c = c + 1 WHERE id = 1",
+		)
+		require.NoError(t, err)
+
+		tk.MustExec("TRUNCATE test.subscribe_table")
+		return nil
+	}
+	n.RegisterHandler(notifier.TestHandlerID, handler)
+	n.OnBecomeOwner()
+	tk2 := testkit.NewTestKit(t, store)
+	se := sess.NewSession(tk2.Session())
+	event1 := notifier.NewCreateTableEvent(&model.TableInfo{ID: 1000, Name: ast.NewCIStr("t1")})
+	err := notifier.PubSchemeChangeToStore(ctx, se, 1, -1, event1, s)
+	require.NoError(t, err)
+	require.Never(t, func() bool {
+		changes := make([]*notifier.SchemaChange, 8)
+		result, closeFn := s.List(ctx, se)
+		count, err2 := result.Read(changes)
+		require.NoError(t, err2)
+		closeFn()
+		return count == 0
+	}, time.Second, 50*time.Millisecond)
+	n.OnRetireOwner()
+}

pkg/ddl/session/session.go

@@ -22,6 +22,7 @@ import (
 	"github.com/pingcap/failpoint"
 	"github.com/pingcap/tidb/pkg/kv"
 	"github.com/pingcap/tidb/pkg/metrics"
+	"github.com/pingcap/tidb/pkg/parser/ast"
 	"github.com/pingcap/tidb/pkg/parser/terror"
 	"github.com/pingcap/tidb/pkg/sessionctx"
 	"github.com/pingcap/tidb/pkg/sessiontxn"
@@ -49,6 +50,19 @@ func (s *Session) Begin(ctx context.Context) error {
 	return nil
 }
 
+// BeginPessimistic starts a pessimistic transaction.
+func (s *Session) BeginPessimistic(ctx context.Context) error {
+	err := sessiontxn.GetTxnManager(s.Context).EnterNewTxn(ctx, &sessiontxn.EnterNewTxnRequest{
+		Type:    sessiontxn.EnterNewTxnDefault,
+		TxnMode: ast.Pessimistic,
+	})
+	if err != nil {
+		return err
+	}
+	s.GetSessionVars().SetInTxn(true)
+	return nil
+}
+
 // Commit commits the transaction.
 func (s *Session) Commit(ctx context.Context) error {
 	s.StmtCommit(ctx)

pkg/ddl/session/session_pool_test.go

@@ -17,6 +17,7 @@ package session_test
 import (
 	"context"
 	"testing"
+	"time"
 
 	"github.com/ngaut/pools"
 	"github.com/pingcap/tidb/pkg/ddl/session"
@@ -66,3 +67,52 @@ func TestSessionPool(t *testing.T) {
 	}
 	require.Equal(t, uint64(0), targetTS)
 }
+
+func TestPessimisticTxn(t *testing.T) {
+	store := testkit.CreateMockStore(t)
+	tk := testkit.NewTestKit(t, store)
+	tk.MustExec("use test")
+	tk.MustExec("create table t (a int primary key, b int)")
+	tk.MustExec("insert into t values (1, 1)")
+
+	resourcePool := pools.NewResourcePool(func() (pools.Resource, error) {
+		newTk := testkit.NewTestKit(t, store)
+		return newTk.Session(), nil
+	}, 4, 4, 0)
+	pool := session.NewSessionPool(resourcePool)
+	ctx := context.Background()
+
+	sessCtx, err := pool.Get()
+	require.NoError(t, err)
+	se := session.NewSession(sessCtx)
+	sessCtx2, err := pool.Get()
+	require.NoError(t, err)
+	se2 := session.NewSession(sessCtx2)
+
+	err = se.BeginPessimistic(ctx)
+	require.NoError(t, err)
+	err = se2.BeginPessimistic(ctx)
+	require.NoError(t, err)
+	_, err = se.Execute(ctx, "update test.t set b = b + 1 where a = 1", "ut")
+	require.NoError(t, err)
+	done := make(chan struct{}, 1)
+	go func() {
+		_, err := se2.Execute(ctx, "update test.t set b = b + 1 where a = 1", "ut")
+		require.NoError(t, err)
+		done <- struct{}{}
+		err = se2.Commit(ctx)
+		require.NoError(t, err)
+		close(done)
+	}()
+
+	time.Sleep(100 * time.Millisecond)
+	// because this is a pessimistic transaction, the second transaction should be blocked
+	require.Len(t, done, 0)
+	err = se.Commit(ctx)
+	require.NoError(t, err)
+	<-done
+	_, ok := <-done
+	require.False(t, ok)
+	pool.Put(sessCtx)
+	pool.Put(sessCtx2)
+}