From 787cb27265a01ce5194a4210636278c3a9d746e0 Mon Sep 17 00:00:00 2001
From: Jk Xu <54522439+Dousir9@users.noreply.github.com>
Date: Wed, 18 Jan 2023 00:43:59 +0800
Subject: [PATCH] meta: add lock at autoid.(*allocator).Base() (#40588)

close pingcap/tidb#40584
---
 meta/autoid/autoid.go      |  4 +++
 meta/autoid/autoid_test.go | 55 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 59 insertions(+)

diff --git a/meta/autoid/autoid.go b/meta/autoid/autoid.go
index 1f5ffeb2fd094..8c5e1d7bc58d8 100644
--- a/meta/autoid/autoid.go
+++ b/meta/autoid/autoid.go
@@ -285,11 +285,15 @@ func SetStep(s int64) {
 
 // Base implements autoid.Allocator Base interface.
 func (alloc *allocator) Base() int64 {
+	alloc.mu.Lock()
+	defer alloc.mu.Unlock()
 	return alloc.base
 }
 
 // End implements autoid.Allocator End interface.
 func (alloc *allocator) End() int64 {
+	alloc.mu.Lock()
+	defer alloc.mu.Unlock()
 	return alloc.end
 }
 
diff --git a/meta/autoid/autoid_test.go b/meta/autoid/autoid_test.go
index 0b8cd60257cf4..912b1e173bc8c 100644
--- a/meta/autoid/autoid_test.go
+++ b/meta/autoid/autoid_test.go
@@ -20,6 +20,7 @@ import (
 	"math"
 	"math/rand"
 	"sync"
+	"sync/atomic"
 	"testing"
 	"time"
 
@@ -594,3 +595,57 @@ func TestAllocComputationIssue(t *testing.T) {
 	require.Equal(t, int64(7), min)
 	require.Equal(t, int64(13), max)
 }
+
+func TestIssue40584(t *testing.T) {
+	store, err := mockstore.NewMockStore()
+	require.NoError(t, err)
+	defer func() {
+		err := store.Close()
+		require.NoError(t, err)
+	}()
+
+	ctx := kv.WithInternalSourceType(context.Background(), kv.InternalTxnMeta)
+	err = kv.RunInNewTxn(ctx, store, false, func(ctx context.Context, txn kv.Transaction) error {
+		m := meta.NewMeta(txn)
+		err = m.CreateDatabase(&model.DBInfo{ID: 1, Name: model.NewCIStr("a")})
+		require.NoError(t, err)
+		err = m.CreateTableOrView(1, &model.TableInfo{ID: 1, Name: model.NewCIStr("t")})
+		require.NoError(t, err)
+		return nil
+	})
+	require.NoError(t, err)
+
+	alloc := autoid.NewAllocator(store, 1, 1, false, autoid.RowIDAllocType)
+	require.NotNil(t, alloc)
+
+	finishAlloc := make(chan bool)
+	finishBase := make(chan bool)
+	var done int32 = 0
+
+	// call allocator.Alloc and allocator.Base in parallel for 3 seconds to detect data race
+	go func() {
+		for {
+			alloc.Alloc(ctx, 1, 1, 1)
+			if atomic.LoadInt32(&done) > 0 {
+				break
+			}
+		}
+		finishAlloc <- true
+	}()
+
+	go func() {
+		for {
+			alloc.Base()
+			if atomic.LoadInt32(&done) > 0 {
+				break
+			}
+		}
+		finishBase <- true
+	}()
+
+	runTime := time.NewTimer(time.Second * 3)
+	<-runTime.C
+	atomic.AddInt32(&done, 1)
+	<-finishAlloc
+	<-finishBase
+}