forked from dgraph-io/dgraph
-
Notifications
You must be signed in to change notification settings - Fork 0
/
embedded_allocator.go
99 lines (86 loc) · 2.38 KB
/
embedded_allocator.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
/*
* Copyright (C) 2017 Dgraph Labs, Inc. and Contributors
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU Affero General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU Affero General Public License for more details.
*
* You should have received a copy of the GNU Affero General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package x
import (
"encoding/binary"
"sync"
"github.com/dgraph-io/badger"
"github.com/dgraph-io/dgraph/protos"
"github.com/pkg/errors"
"golang.org/x/net/context"
)
const (
leaseBandwidth = uint64(10000)
)
var (
emptyAssignedIds = &protos.AssignedIds{}
)
type EmbeddedUidAllocator struct {
sync.Mutex
nextLeaseId uint64
maxLeaseId uint64
pstore *badger.KV
}
// Start lease from 2, 1 is used by _lease_
func (e *EmbeddedUidAllocator) Init(kv *badger.KV) {
e.pstore = kv
var item badger.KVItem
// All keys start with 0x00 or 0x01 so shouldn't collide
e.pstore.Get([]byte("uid_lease"), &item)
e.maxLeaseId = 1
var n int
err := item.Value(func(val []byte) error {
if len(val) > 0 {
e.maxLeaseId, n = binary.Uvarint(val)
AssertTrue(n > 0)
}
return nil
})
Check(err)
e.nextLeaseId = e.maxLeaseId + 1
}
func (e *EmbeddedUidAllocator) AssignUids(ctx context.Context,
num *protos.Num) (*protos.AssignedIds, error) {
val := int(num.Val)
if val == 0 {
return emptyAssignedIds, Errorf("Nothing to be marked or assigned")
}
e.Lock()
defer e.Unlock()
howMany := leaseBandwidth
if num.Val > leaseBandwidth {
howMany = num.Val + leaseBandwidth
}
if e.nextLeaseId == 0 {
return nil, errors.New("Server not initialized.")
}
available := e.maxLeaseId - e.nextLeaseId + 1
if available < num.Val {
e.maxLeaseId += howMany
val := make([]byte, 10)
n := binary.PutUvarint(val, e.maxLeaseId)
err := e.pstore.Set([]byte("uid_lease"), val[:n], 0x01)
if err != nil {
return emptyAssignedIds, err
}
}
out := &protos.AssignedIds{}
out.StartId = e.nextLeaseId
out.EndId = out.StartId + num.Val - 1
e.nextLeaseId = out.EndId + 1
return out, nil
}