From 16f5d7c9d20a1fae6e254dd91d27f0ae1b55af7e Mon Sep 17 00:00:00 2001 From: hillium Date: Mon, 26 Dec 2022 18:35:19 +0800 Subject: [PATCH 1/3] advance ID to max of paritions Signed-off-by: hillium --- br/pkg/restore/db.go | 8 ++++++++ br/pkg/restore/prealloc_table_id/alloc.go | 8 ++++++++ 2 files changed, 16 insertions(+) diff --git a/br/pkg/restore/db.go b/br/pkg/restore/db.go index ae62162c3e890..97653b048b95c 100644 --- a/br/pkg/restore/db.go +++ b/br/pkg/restore/db.go @@ -285,6 +285,14 @@ func (db *DB) tableIDAllocFilter() ddl.AllocTableIDIf { return true } prealloced := db.preallocedIDs.Prealloced(ti.ID) + if ti.Partition != nil && ti.Partition.Definitions != nil { + for _, part := range ti.Partition.Definitions { + if !db.preallocedIDs.Prealloced(part.ID) { + prealloced = false + break + } + } + } if prealloced { log.Info("reusing table ID", zap.Stringer("table", ti.Name)) } diff --git a/br/pkg/restore/prealloc_table_id/alloc.go b/br/pkg/restore/prealloc_table_id/alloc.go index 9232ed84a8fc8..7503a3dffa024 100644 --- a/br/pkg/restore/prealloc_table_id/alloc.go +++ b/br/pkg/restore/prealloc_table_id/alloc.go @@ -48,6 +48,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, From 12af0ef4f85a4163bc28f9a9eecd215b481cfb15 Mon Sep 17 00:00:00 2001 From: hillium Date: Mon, 26 Dec 2022 18:54:12 +0800 Subject: [PATCH 2/3] added test Signed-off-by: hillium --- br/pkg/restore/db.go | 10 +---- br/pkg/restore/prealloc_table_id/alloc.go | 15 ++++++++ .../restore/prealloc_table_id/alloc_test.go | 38 ++++++++++++++++--- 3 files changed, 48 insertions(+), 15 deletions(-) diff --git a/br/pkg/restore/db.go b/br/pkg/restore/db.go index 97653b048b95c..132e3294f1617 100644 --- a/br/pkg/restore/db.go +++ b/br/pkg/restore/db.go @@ -284,15 +284,7 @@ func (db *DB) tableIDAllocFilter() ddl.AllocTableIDIf { if db.preallocedIDs == nil { return true } - prealloced := db.preallocedIDs.Prealloced(ti.ID) - if ti.Partition != nil && ti.Partition.Definitions != nil { - for _, part := range ti.Partition.Definitions { - if !db.preallocedIDs.Prealloced(part.ID) { - prealloced = false - break - } - } - } + 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/alloc.go b/br/pkg/restore/prealloc_table_id/alloc.go index 7503a3dffa024..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 ( @@ -94,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) From ca79a6c12c4dfc45f7547d490bf3c9cd31ac2428 Mon Sep 17 00:00:00 2001 From: hillium Date: Mon, 26 Dec 2022 19:35:01 +0800 Subject: [PATCH 3/3] bazel build Signed-off-by: hillium --- br/pkg/restore/prealloc_table_id/BUILD.bazel | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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(