github.com/vescale/zgraph@v0.0.0-20230410094002-959c02d50f95/stmtctx/id_range.go (about) 1 // Copyright 2022 zGraph Authors. All rights reserved. 2 // 3 // Licensed under the Apache License, Version 2.0 (the "License"); 4 // you may not use this file except in compliance with the License. 5 // You may obtain a copy of the License at 6 // 7 // http://www.apache.org/licenses/LICENSE-2.0 8 // 9 // Unless required by applicable law or agreed to in writing, software 10 // distributed under the License is distributed on an "AS IS" BASIS, 11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 // See the License for the specific language governing permissions and 13 // limitations under the License. 14 15 package stmtctx 16 17 import ( 18 "sync/atomic" 19 20 "github.com/vescale/zgraph/catalog" 21 "github.com/vescale/zgraph/meta" 22 "github.com/vescale/zgraph/storage/kv" 23 ) 24 25 // IDRange represents an ID range. The ID range will be (base, max] 26 type IDRange struct { 27 base atomic.Int64 28 max int64 29 } 30 31 // NewIDRange returns a new ID range. 32 func NewIDRange(base, max int64) *IDRange { 33 idr := &IDRange{ 34 max: max, 35 } 36 idr.base.Store(base) 37 return idr 38 } 39 40 // Next retrieves the next available ID. 41 func (r *IDRange) Next() (int64, error) { 42 next := r.base.Add(1) 43 if next > r.max { 44 return 0, ErrIDExhaust 45 } 46 return next, nil 47 } 48 49 // AllocID allocates n IDs. 50 func (sc *Context) AllocID(graph *catalog.Graph, n int) (*IDRange, error) { 51 graph.MDLock() 52 defer graph.MDUnlock() 53 54 var idRange *IDRange 55 err := kv.Txn(sc.store, func(txn kv.Transaction) error { 56 meta := meta.New(txn) 57 base, err := meta.AdvanceID(graph.Meta().ID, n) 58 if err != nil { 59 return err 60 } 61 idRange = NewIDRange(base, base+int64(n)) 62 return nil 63 }) 64 65 return idRange, err 66 }