2021-10-14 19:22:43 +08:00
|
|
|
// Licensed to the LF AI & Data foundation under one
|
|
|
|
// or more contributor license agreements. See the NOTICE file
|
|
|
|
// distributed with this work for additional information
|
|
|
|
// regarding copyright ownership. The ASF licenses this file
|
|
|
|
// to you under the Apache License, Version 2.0 (the
|
|
|
|
// "License"); you may not use this file except in compliance
|
2021-04-19 15:15:33 +08:00
|
|
|
// with the License. You may obtain a copy of the License at
|
|
|
|
//
|
2021-10-14 19:22:43 +08:00
|
|
|
// http://www.apache.org/licenses/LICENSE-2.0
|
2021-04-19 15:15:33 +08:00
|
|
|
//
|
2021-10-14 19:22:43 +08:00
|
|
|
// Unless required by applicable law or agreed to in writing, software
|
|
|
|
// distributed under the License is distributed on an "AS IS" BASIS,
|
|
|
|
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
|
|
// See the License for the specific language governing permissions and
|
|
|
|
// limitations under the License.
|
2021-04-19 15:15:33 +08:00
|
|
|
|
2020-10-30 16:27:58 +08:00
|
|
|
package allocator
|
|
|
|
|
|
|
|
import (
|
2020-11-03 14:53:36 +08:00
|
|
|
"context"
|
2021-04-01 13:37:18 +08:00
|
|
|
"fmt"
|
2020-11-03 14:53:36 +08:00
|
|
|
"time"
|
|
|
|
|
2022-10-16 20:49:27 +08:00
|
|
|
"github.com/milvus-io/milvus-proto/go-api/commonpb"
|
2021-06-22 16:14:09 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/proto/rootcoordpb"
|
2022-10-21 15:57:28 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/util/commonpbutil"
|
2021-04-22 14:45:57 +08:00
|
|
|
"github.com/milvus-io/milvus/internal/util/typeutil"
|
2020-10-30 16:27:58 +08:00
|
|
|
)
|
|
|
|
|
2020-11-19 21:02:31 +08:00
|
|
|
const (
|
2021-10-04 08:36:30 +08:00
|
|
|
idCountPerRPC = 200000
|
2020-11-19 21:02:31 +08:00
|
|
|
)
|
|
|
|
|
2021-10-01 10:38:30 +08:00
|
|
|
// UniqueID is alias of typeutil.UniqueID
|
2020-11-04 17:58:43 +08:00
|
|
|
type UniqueID = typeutil.UniqueID
|
|
|
|
|
2021-10-24 10:45:10 +08:00
|
|
|
// IDAllocator allocates Unique and monotonically increasing IDs from Root Coord.
|
2021-09-14 09:54:31 +08:00
|
|
|
// It could also batch allocate for less root coord server access
|
2020-11-12 12:04:12 +08:00
|
|
|
type IDAllocator struct {
|
2022-10-09 10:06:58 +08:00
|
|
|
CachedAllocator
|
2020-10-30 16:27:58 +08:00
|
|
|
|
2022-10-09 10:06:58 +08:00
|
|
|
remoteAllocator remoteInterface
|
2021-01-29 09:27:26 +08:00
|
|
|
|
|
|
|
countPerRPC uint32
|
|
|
|
|
2020-11-04 17:58:43 +08:00
|
|
|
idStart UniqueID
|
|
|
|
idEnd UniqueID
|
2020-12-24 20:55:40 +08:00
|
|
|
|
|
|
|
PeerID UniqueID
|
2020-10-30 16:27:58 +08:00
|
|
|
}
|
|
|
|
|
2021-12-16 20:25:28 +08:00
|
|
|
// NewIDAllocator creates an ID Allocator allocate Unique and monotonically increasing IDs from RootCoord.
|
2022-10-09 10:06:58 +08:00
|
|
|
func NewIDAllocator(ctx context.Context, remoteAllocator remoteInterface, peerID UniqueID) (*IDAllocator, error) {
|
2020-11-03 14:53:36 +08:00
|
|
|
ctx1, cancel := context.WithCancel(ctx)
|
2020-11-12 12:04:12 +08:00
|
|
|
a := &IDAllocator{
|
2022-10-09 10:06:58 +08:00
|
|
|
CachedAllocator: CachedAllocator{
|
2021-01-29 09:27:26 +08:00
|
|
|
Ctx: ctx1,
|
|
|
|
CancelFunc: cancel,
|
2021-04-01 13:37:18 +08:00
|
|
|
Role: "IDAllocator",
|
2020-11-03 14:53:36 +08:00
|
|
|
},
|
2022-10-09 10:06:58 +08:00
|
|
|
countPerRPC: idCountPerRPC,
|
|
|
|
remoteAllocator: remoteAllocator,
|
|
|
|
PeerID: peerID,
|
2020-11-03 14:53:36 +08:00
|
|
|
}
|
2021-01-29 09:27:26 +08:00
|
|
|
a.TChan = &EmptyTicker{}
|
2022-10-09 10:06:58 +08:00
|
|
|
a.CachedAllocator.SyncFunc = a.syncID
|
|
|
|
a.CachedAllocator.ProcessFunc = a.processFunc
|
|
|
|
a.CachedAllocator.CheckSyncFunc = a.checkSyncFunc
|
|
|
|
a.CachedAllocator.PickCanDoFunc = a.pickCanDoFunc
|
2021-01-29 09:27:26 +08:00
|
|
|
a.Init()
|
2020-11-03 14:53:36 +08:00
|
|
|
return a, nil
|
2020-10-30 16:27:58 +08:00
|
|
|
}
|
|
|
|
|
2021-10-01 10:38:30 +08:00
|
|
|
// Start creates some working goroutines of IDAllocator.
|
2021-01-29 09:27:26 +08:00
|
|
|
func (ia *IDAllocator) Start() error {
|
2022-10-09 10:06:58 +08:00
|
|
|
return ia.CachedAllocator.Start()
|
2021-01-29 09:27:26 +08:00
|
|
|
}
|
|
|
|
|
2021-04-01 13:37:18 +08:00
|
|
|
func (ia *IDAllocator) gatherReqIDCount() uint32 {
|
|
|
|
need := uint32(0)
|
|
|
|
for _, req := range ia.ToDoReqs {
|
|
|
|
tReq := req.(*IDRequest)
|
|
|
|
need += tReq.count
|
|
|
|
}
|
|
|
|
return need
|
|
|
|
}
|
|
|
|
|
|
|
|
func (ia *IDAllocator) syncID() (bool, error) {
|
|
|
|
|
|
|
|
need := ia.gatherReqIDCount()
|
|
|
|
if need < ia.countPerRPC {
|
|
|
|
need = ia.countPerRPC
|
|
|
|
}
|
|
|
|
|
2020-11-03 14:53:36 +08:00
|
|
|
ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
2021-06-22 16:14:09 +08:00
|
|
|
req := &rootcoordpb.AllocIDRequest{
|
2022-10-21 15:57:28 +08:00
|
|
|
Base: commonpbutil.NewMsgBase(
|
|
|
|
commonpbutil.WithMsgType(commonpb.MsgType_RequestID),
|
|
|
|
commonpbutil.WithMsgID(0),
|
|
|
|
commonpbutil.WithSourceID(ia.PeerID),
|
|
|
|
),
|
2021-04-01 13:37:18 +08:00
|
|
|
Count: need,
|
2020-11-03 14:53:36 +08:00
|
|
|
}
|
2022-10-09 10:06:58 +08:00
|
|
|
resp, err := ia.remoteAllocator.AllocID(ctx, req)
|
2020-10-30 16:27:58 +08:00
|
|
|
|
2020-11-03 14:53:36 +08:00
|
|
|
cancel()
|
|
|
|
if err != nil {
|
2021-04-01 13:37:18 +08:00
|
|
|
return false, fmt.Errorf("syncID Failed:%w", err)
|
2020-11-03 14:53:36 +08:00
|
|
|
}
|
2020-11-19 21:02:31 +08:00
|
|
|
ia.idStart = resp.GetID()
|
|
|
|
ia.idEnd = ia.idStart + int64(resp.GetCount())
|
2021-04-01 13:37:18 +08:00
|
|
|
return true, nil
|
2020-12-24 16:53:31 +08:00
|
|
|
}
|
|
|
|
|
2020-12-24 20:55:40 +08:00
|
|
|
func (ia *IDAllocator) checkSyncFunc(timeout bool) bool {
|
2021-01-29 09:27:26 +08:00
|
|
|
return timeout || len(ia.ToDoReqs) > 0
|
2020-12-24 20:55:40 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
func (ia *IDAllocator) pickCanDoFunc() {
|
|
|
|
total := uint32(ia.idEnd - ia.idStart)
|
2020-11-19 21:02:31 +08:00
|
|
|
need := uint32(0)
|
2020-12-24 20:55:40 +08:00
|
|
|
idx := 0
|
2021-01-29 09:27:26 +08:00
|
|
|
for _, req := range ia.ToDoReqs {
|
|
|
|
iReq := req.(*IDRequest)
|
2020-11-19 21:02:31 +08:00
|
|
|
need += iReq.count
|
2020-12-24 20:55:40 +08:00
|
|
|
if need <= total {
|
2021-01-29 09:27:26 +08:00
|
|
|
ia.CanDoReqs = append(ia.CanDoReqs, req)
|
2020-12-24 20:55:40 +08:00
|
|
|
idx++
|
|
|
|
} else {
|
|
|
|
break
|
|
|
|
}
|
2020-11-19 21:02:31 +08:00
|
|
|
}
|
2021-01-29 09:27:26 +08:00
|
|
|
ia.ToDoReqs = ia.ToDoReqs[idx:]
|
2020-10-30 16:27:58 +08:00
|
|
|
}
|
|
|
|
|
2021-01-29 09:27:26 +08:00
|
|
|
func (ia *IDAllocator) processFunc(req Request) error {
|
|
|
|
idRequest := req.(*IDRequest)
|
2020-11-19 21:02:31 +08:00
|
|
|
idRequest.id = ia.idStart
|
2021-06-23 17:34:05 +08:00
|
|
|
ia.idStart += int64(idRequest.count)
|
2020-11-19 21:02:31 +08:00
|
|
|
return nil
|
2020-11-03 14:53:36 +08:00
|
|
|
}
|
2020-10-30 16:27:58 +08:00
|
|
|
|
2021-10-01 10:38:30 +08:00
|
|
|
// AllocOne allocates one id.
|
2020-11-19 21:02:31 +08:00
|
|
|
func (ia *IDAllocator) AllocOne() (UniqueID, error) {
|
|
|
|
ret, _, err := ia.Alloc(1)
|
2020-11-03 14:53:36 +08:00
|
|
|
if err != nil {
|
|
|
|
return 0, err
|
|
|
|
}
|
|
|
|
return ret, nil
|
2020-10-30 16:27:58 +08:00
|
|
|
}
|
|
|
|
|
2021-10-01 10:38:30 +08:00
|
|
|
// Alloc allocates the id of the count number.
|
2020-11-19 21:02:31 +08:00
|
|
|
func (ia *IDAllocator) Alloc(count uint32) (UniqueID, UniqueID, error) {
|
2021-01-29 09:27:26 +08:00
|
|
|
req := &IDRequest{BaseRequest: BaseRequest{Done: make(chan error), Valid: false}}
|
2020-11-03 14:53:36 +08:00
|
|
|
|
|
|
|
req.count = count
|
2021-01-29 09:27:26 +08:00
|
|
|
ia.Reqs <- req
|
2021-03-29 15:14:51 +08:00
|
|
|
if err := req.Wait(); err != nil {
|
|
|
|
return 0, 0, err
|
2020-11-03 14:53:36 +08:00
|
|
|
}
|
2021-03-29 15:14:51 +08:00
|
|
|
|
2020-11-04 17:58:43 +08:00
|
|
|
start, count := req.id, req.count
|
2020-11-03 14:53:36 +08:00
|
|
|
return start, start + int64(count), nil
|
|
|
|
}
|