Skip to content

Commit

Permalink
restore: Fix prealloc (#40176)
Browse files Browse the repository at this point in the history
close #40177
  • Loading branch information
YuJuncen authored Dec 26, 2022
1 parent fc714e2 commit 79ab2b9
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 8 deletions.
2 changes: 1 addition & 1 deletion br/pkg/restore/db.go
Original file line number Diff line number Diff line change
Expand Up @@ -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))
}
Expand Down
5 changes: 4 additions & 1 deletion br/pkg/restore/prealloc_table_id/BUILD.bazel
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand Down
23 changes: 23 additions & 0 deletions br/pkg/restore/prealloc_table_id/alloc.go
Original file line number Diff line number Diff line change
Expand Up @@ -7,6 +7,7 @@ import (
"math"

"github.com/pingcap/tidb/br/pkg/metautil"
"github.com/pingcap/tidb/parser/model"
)

const (
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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
}
38 changes: 32 additions & 6 deletions br/pkg/restore/prealloc_table_id/alloc_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -57,26 +58,51 @@ 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)
allocator := testAllocator(c.hasAllocatedTo)
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)
Expand Down

0 comments on commit 79ab2b9

Please sign in to comment.