From 8ed1d9d4a7987aa87b586020795d2587286402bf Mon Sep 17 00:00:00 2001 From: djshow832 Date: Sat, 8 May 2021 16:36:41 +0800 Subject: [PATCH] meta: add an in-memory autoID allocator (#24438) --- meta/autoid/errors.go | 1 + meta/autoid/memid.go | 145 ++++++++++++++++++++++++++++++++++++++ meta/autoid/memid_test.go | 106 ++++++++++++++++++++++++++++ 3 files changed, 252 insertions(+) create mode 100644 meta/autoid/memid.go create mode 100644 meta/autoid/memid_test.go diff --git a/meta/autoid/errors.go b/meta/autoid/errors.go index 093e0abbdea53..777de466b7704 100644 --- a/meta/autoid/errors.go +++ b/meta/autoid/errors.go @@ -22,6 +22,7 @@ import ( var ( errInvalidTableID = dbterror.ClassAutoid.NewStd(mysql.ErrInvalidTableID) errInvalidIncrementAndOffset = dbterror.ClassAutoid.NewStd(mysql.ErrInvalidIncrementAndOffset) + errNotImplemented = dbterror.ClassAutoid.NewStd(mysql.ErrNotImplemented) ErrAutoincReadFailed = dbterror.ClassAutoid.NewStd(mysql.ErrAutoincReadFailed) ErrWrongAutoKey = dbterror.ClassAutoid.NewStd(mysql.ErrWrongAutoKey) ErrInvalidAllocatorType = dbterror.ClassAutoid.NewStd(mysql.ErrUnknownAllocatorType) diff --git a/meta/autoid/memid.go b/meta/autoid/memid.go new file mode 100644 index 0000000000000..703ffe1db4bf8 --- /dev/null +++ b/meta/autoid/memid.go @@ -0,0 +1,145 @@ +// Copyright 2021 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// See the License for the specific language governing permissions and +// limitations under the License. + +package autoid + +import ( + "context" + "math" + + "github.com/pingcap/parser/model" +) + +// NewAllocatorFromTempTblInfo creates an in-memory allocator from a temporary table info. +func NewAllocatorFromTempTblInfo(tblInfo *model.TableInfo) Allocator { + hasRowID := !tblInfo.PKIsHandle && !tblInfo.IsCommonHandle + hasAutoIncID := tblInfo.GetAutoIncrementColInfo() != nil + // Temporary tables don't support auto_random and sequence. + if hasRowID || hasAutoIncID { + return &inMemoryAllocator{ + isUnsigned: tblInfo.IsAutoIncColUnsigned(), + allocType: RowIDAllocType, + } + } + return nil +} + +// inMemoryAllocator is typically used for temporary tables. +// Some characteristics: +// - It allocates IDs from memory. +// - It's session-wide and thus won't be visited concurrently. +// - It doesn't support sequence. +// - The metrics are not reported. +type inMemoryAllocator struct { + base int64 + isUnsigned bool + allocType AllocatorType +} + +// Base implements autoid.Allocator Base interface. +func (alloc *inMemoryAllocator) Base() int64 { + return alloc.base +} + +// End implements autoid.Allocator End interface. +func (alloc *inMemoryAllocator) End() int64 { + // It doesn't matter. + return 0 +} + +// GetType implements autoid.Allocator GetType interface. +func (alloc *inMemoryAllocator) GetType() AllocatorType { + return alloc.allocType +} + +// NextGlobalAutoID implements autoid.Allocator NextGlobalAutoID interface. +func (alloc *inMemoryAllocator) NextGlobalAutoID(tableID int64) (int64, error) { + return 0, errNotImplemented.GenWithStackByArgs() +} + +func (alloc *inMemoryAllocator) Alloc(ctx context.Context, tableID int64, n uint64, increment, offset int64) (int64, int64, error) { + if n == 0 { + return 0, 0, nil + } + if alloc.allocType == AutoIncrementType || alloc.allocType == RowIDAllocType { + if !validIncrementAndOffset(increment, offset) { + return 0, 0, errInvalidIncrementAndOffset.GenWithStackByArgs(increment, offset) + } + } + if alloc.isUnsigned { + return alloc.alloc4Unsigned(n, increment, offset) + } + return alloc.alloc4Signed(n, increment, offset) +} + +// Rebase implements autoid.Allocator Rebase interface. +// The requiredBase is the minimum base value after Rebase. +// The real base may be greater than the required base. +func (alloc *inMemoryAllocator) Rebase(tableID, requiredBase int64, allocIDs bool) error { + if alloc.isUnsigned { + if uint64(requiredBase) > uint64(alloc.base) { + alloc.base = requiredBase + } + } else { + if requiredBase > alloc.base { + alloc.base = requiredBase + } + } + return nil +} + +func (alloc *inMemoryAllocator) alloc4Signed(n uint64, increment, offset int64) (int64, int64, error) { + // Check offset rebase if necessary. + if offset-1 > alloc.base { + alloc.base = offset - 1 + } + // CalcNeededBatchSize calculates the total batch size needed. + n1 := CalcNeededBatchSize(alloc.base, int64(n), increment, offset, alloc.isUnsigned) + + // Condition alloc.base+N1 > alloc.end will overflow when alloc.base + N1 > MaxInt64. So need this. + if math.MaxInt64-alloc.base <= n1 { + return 0, 0, ErrAutoincReadFailed + } + + min := alloc.base + alloc.base += n1 + return min, alloc.base, nil +} + +func (alloc *inMemoryAllocator) alloc4Unsigned(n uint64, increment, offset int64) (int64, int64, error) { + // Check offset rebase if necessary. + if uint64(offset)-1 > uint64(alloc.base) { + alloc.base = int64(uint64(offset) - 1) + } + + // CalcNeededBatchSize calculates the total batch size needed. + n1 := CalcNeededBatchSize(alloc.base, int64(n), increment, offset, alloc.isUnsigned) + + // Condition alloc.base+n1 > alloc.end will overflow when alloc.base + n1 > MaxInt64. So need this. + if math.MaxUint64-uint64(alloc.base) <= uint64(n1) { + return 0, 0, ErrAutoincReadFailed + } + + min := alloc.base + // Use uint64 n directly. + alloc.base = int64(uint64(alloc.base) + uint64(n1)) + return min, alloc.base, nil +} + +func (alloc *inMemoryAllocator) AllocSeqCache(tableID int64) (int64, int64, int64, error) { + return 0, 0, 0, errNotImplemented.GenWithStackByArgs() +} + +func (alloc *inMemoryAllocator) RebaseSeq(tableID, requiredBase int64) (int64, bool, error) { + return 0, false, errNotImplemented.GenWithStackByArgs() +} diff --git a/meta/autoid/memid_test.go b/meta/autoid/memid_test.go new file mode 100644 index 0000000000000..f4b1267ad8f99 --- /dev/null +++ b/meta/autoid/memid_test.go @@ -0,0 +1,106 @@ +// Copyright 2021 PingCAP, Inc. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// See the License for the specific language governing permissions and +// limitations under the License. + +package autoid_test + +import ( + "context" + "math" + + . "github.com/pingcap/check" + "github.com/pingcap/parser/model" + "github.com/pingcap/parser/mysql" + "github.com/pingcap/parser/terror" + "github.com/pingcap/parser/types" + "github.com/pingcap/tidb/meta/autoid" + "github.com/pingcap/tidb/store/mockstore" +) + +func (*testSuite) TestInMemoryAlloc(c *C) { + store, err := mockstore.NewMockStore() + c.Assert(err, IsNil) + defer func() { + err := store.Close() + c.Assert(err, IsNil) + }() + + columnInfo := &model.ColumnInfo{ + FieldType: types.FieldType{ + Flag: mysql.AutoIncrementFlag, + }, + } + tblInfo := &model.TableInfo{ + Columns: []*model.ColumnInfo{columnInfo}, + } + alloc := autoid.NewAllocatorFromTempTblInfo(tblInfo) + c.Assert(alloc, NotNil) + + // alloc 1 + ctx := context.Background() + _, id, err := alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(1)) + _, id, err = alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(2)) + + // alloc N + _, id, err = alloc.Alloc(ctx, 1, 10, 1, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(12)) + + // increment > N + _, id, err = alloc.Alloc(ctx, 1, 1, 10, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(21)) + + // offset + _, id, err = alloc.Alloc(ctx, 1, 1, 1, 30) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(30)) + + // rebase + err = alloc.Rebase(1, int64(40), true) + c.Assert(err, IsNil) + _, id, err = alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(41)) + err = alloc.Rebase(1, int64(10), true) + c.Assert(err, IsNil) + _, id, err = alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(42)) + + // maxInt64 + err = alloc.Rebase(1, int64(math.MaxInt64-2), true) + c.Assert(err, IsNil) + _, id, err = alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(math.MaxInt64-1)) + _, _, err = alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(terror.ErrorEqual(err, autoid.ErrAutoincReadFailed), IsTrue) + + // test unsigned + columnInfo.FieldType.Flag |= mysql.UnsignedFlag + alloc = autoid.NewAllocatorFromTempTblInfo(tblInfo) + c.Assert(alloc, NotNil) + + var n uint64 = math.MaxUint64 - 2 + err = alloc.Rebase(1, int64(n), true) + c.Assert(err, IsNil) + _, id, err = alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(err, IsNil) + c.Assert(id, Equals, int64(n+1)) + _, _, err = alloc.Alloc(ctx, 1, 1, 1, 1) + c.Assert(terror.ErrorEqual(err, autoid.ErrAutoincReadFailed), IsTrue) +}