From 79ab2b92a046bba03f957c4bf421eb9fb55bc053 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E5=B1=B1=E5=B2=9A?=
 <36239017+YuJuncen@users.noreply.github.com>
Date: Mon, 26 Dec 2022 22:38:15 +0800
Subject: [PATCH] restore: Fix prealloc (#40176)

close pingcap/tidb#40177
---
 br/pkg/restore/db.go                          |  2 +-
 br/pkg/restore/prealloc_table_id/BUILD.bazel  |  5 ++-
 br/pkg/restore/prealloc_table_id/alloc.go     | 23 +++++++++++
 .../restore/prealloc_table_id/alloc_test.go   | 38 ++++++++++++++++---
 4 files changed, 60 insertions(+), 8 deletions(-)

diff --git a/br/pkg/restore/db.go b/br/pkg/restore/db.go
index ae62162c3e890..132e3294f1617 100644
--- a/br/pkg/restore/db.go
+++ b/br/pkg/restore/db.go
@@ -284,7 +284,7 @@ func (db *DB) tableIDAllocFilter() ddl.AllocTableIDIf {
 		if db.preallocedIDs == nil {
 			return true
 		}
-		prealloced := db.preallocedIDs.Prealloced(ti.ID)
+		prealloced := db.preallocedIDs.PreallocedFor(ti)
 		if prealloced {
 			log.Info("reusing table ID", zap.Stringer("table", ti.Name))
 		}
diff --git a/br/pkg/restore/prealloc_table_id/BUILD.bazel b/br/pkg/restore/prealloc_table_id/BUILD.bazel
index 8ce80b039178a..cfdb0432fd446 100644
--- a/br/pkg/restore/prealloc_table_id/BUILD.bazel
+++ b/br/pkg/restore/prealloc_table_id/BUILD.bazel
@@ -5,7 +5,10 @@ go_library(
     srcs = ["alloc.go"],
     importpath = "github.com/pingcap/tidb/br/pkg/restore/prealloc_table_id",
     visibility = ["//visibility:public"],
-    deps = ["//br/pkg/metautil"],
+    deps = [
+        "//br/pkg/metautil",
+        "//parser/model",
+    ],
 )
 
 go_test(
diff --git a/br/pkg/restore/prealloc_table_id/alloc.go b/br/pkg/restore/prealloc_table_id/alloc.go
index 9232ed84a8fc8..8554de5e9891b 100644
--- a/br/pkg/restore/prealloc_table_id/alloc.go
+++ b/br/pkg/restore/prealloc_table_id/alloc.go
@@ -7,6 +7,7 @@ import (
 	"math"
 
 	"github.com/pingcap/tidb/br/pkg/metautil"
+	"github.com/pingcap/tidb/parser/model"
 )
 
 const (
@@ -48,6 +49,14 @@ func New(tables []*metautil.Table) *PreallocIDs {
 		if t.Info.ID > max && t.Info.ID < insaneTableIDThreshold {
 			max = t.Info.ID
 		}
+
+		if t.Info.Partition != nil && t.Info.Partition.Definitions != nil {
+			for _, part := range t.Info.Partition.Definitions {
+				if part.ID > max && part.ID < insaneTableIDThreshold {
+					max = part.ID
+				}
+			}
+		}
 	}
 	return &PreallocIDs{
 		end: max + 1,
@@ -86,3 +95,17 @@ func (p *PreallocIDs) Alloc(m Allocator) error {
 func (p *PreallocIDs) Prealloced(tid int64) bool {
 	return p.allocedFrom <= tid && tid < p.end
 }
+
+func (p *PreallocIDs) PreallocedFor(ti *model.TableInfo) bool {
+	if !p.Prealloced(ti.ID) {
+		return false
+	}
+	if ti.Partition != nil && ti.Partition.Definitions != nil {
+		for _, part := range ti.Partition.Definitions {
+			if !p.Prealloced(part.ID) {
+				return false
+			}
+		}
+	}
+	return true
+}
diff --git a/br/pkg/restore/prealloc_table_id/alloc_test.go b/br/pkg/restore/prealloc_table_id/alloc_test.go
index 8cf6b95fb070e..c1c3f018a2de8 100644
--- a/br/pkg/restore/prealloc_table_id/alloc_test.go
+++ b/br/pkg/restore/prealloc_table_id/alloc_test.go
@@ -27,6 +27,7 @@ func (t *testAllocator) AdvanceGlobalIDs(n int) (int64, error) {
 func TestAllocator(t *testing.T) {
 	type Case struct {
 		tableIDs              []int64
+		partitions            map[int64][]int64
 		hasAllocatedTo        int64
 		successfullyAllocated []int64
 		shouldAllocatedTo     int64
@@ -57,16 +58,41 @@ func TestAllocator(t *testing.T) {
 			successfullyAllocated: []int64{5, 6},
 			shouldAllocatedTo:     7,
 		},
+		{
+			tableIDs:              []int64{1, 2, 5, 6, 7},
+			hasAllocatedTo:        6,
+			successfullyAllocated: []int64{6, 7},
+			shouldAllocatedTo:     13,
+			partitions: map[int64][]int64{
+				7: {8, 9, 10, 11, 12},
+			},
+		},
+		{
+			tableIDs:              []int64{1, 2, 5, 6, 7, 13},
+			hasAllocatedTo:        9,
+			successfullyAllocated: []int64{13},
+			shouldAllocatedTo:     14,
+			partitions: map[int64][]int64{
+				7: {8, 9, 10, 11, 12},
+			},
+		},
 	}
 
 	run := func(t *testing.T, c Case) {
 		tables := make([]*metautil.Table, 0, len(c.tableIDs))
 		for _, id := range c.tableIDs {
-			tables = append(tables, &metautil.Table{
+			table := metautil.Table{
 				Info: &model.TableInfo{
-					ID: id,
+					ID:        id,
+					Partition: &model.PartitionInfo{},
 				},
-			})
+			}
+			if c.partitions != nil {
+				for _, part := range c.partitions[id] {
+					table.Info.Partition.Definitions = append(table.Info.Partition.Definitions, model.PartitionDefinition{ID: part})
+				}
+			}
+			tables = append(tables, &table)
 		}
 
 		ids := prealloctableid.New(tables)
@@ -74,9 +100,9 @@ func TestAllocator(t *testing.T) {
 		require.NoError(t, ids.Alloc(&allocator))
 
 		allocated := make([]int64, 0, len(c.successfullyAllocated))
-		for _, t := range c.tableIDs {
-			if ids.Prealloced(t) {
-				allocated = append(allocated, t)
+		for _, t := range tables {
+			if ids.PreallocedFor(t.Info) {
+				allocated = append(allocated, t.Info.ID)
 			}
 		}
 		require.ElementsMatch(t, allocated, c.successfullyAllocated)