Skip to content

Commit 8e08c43

Browse files
paralinparalin
authored
Update ristretto and refactor for use of generics (#2047)
Signed-off-by: Christian Stewart <[email protected]>
1 parent 6260e9f commit 8e08c43

File tree

7 files changed

+39
-44
lines changed

7 files changed

+39
-44
lines changed

‎db.go

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ import (
3434
humanize "github.com/dustin/go-humanize"
3535
"github.com/pkg/errors"
3636

37+
"github.com/dgraph-io/badger/v4/fb"
3738
"github.com/dgraph-io/badger/v4/options"
3839
"github.com/dgraph-io/badger/v4/pb"
3940
"github.com/dgraph-io/badger/v4/skl"
@@ -123,8 +124,8 @@ type DB struct {
123124

124125
pub *publisher
125126
registry *KeyRegistry
126-
blockCache *ristretto.Cache
127-
indexCache *ristretto.Cache
127+
blockCache *ristretto.Cache[[]byte, *table.Block]
128+
indexCache *ristretto.Cache[uint64, *fb.TableIndex]
128129
allocPool *z.AllocatorPool
129130
}
130131

@@ -274,14 +275,14 @@ func Open(opt Options) (*DB, error) {
274275
numInCache = 1
275276
}
276277

277-
config := ristretto.Config{
278+
config := ristretto.Config[[]byte, *table.Block]{
278279
NumCounters: numInCache * 8,
279280
MaxCost: opt.BlockCacheSize,
280281
BufferItems: 64,
281282
Metrics: true,
282283
OnExit: table.BlockEvictHandler,
283284
}
284-
db.blockCache, err = ristretto.NewCache(&config)
285+
db.blockCache, err = ristretto.NewCache[[]byte, *table.Block](&config)
285286
if err != nil {
286287
return nil, y.Wrap(err, "failed to create data cache")
287288
}
@@ -297,7 +298,7 @@ func Open(opt Options) (*DB, error) {
297298
numInCache = 1
298299
}
299300

300-
config := ristretto.Config{
301+
config := ristretto.Config[uint64, *fb.TableIndex]{
301302
NumCounters: numInCache * 8,
302303
MaxCost: opt.IndexCacheSize,
303304
BufferItems: 64,

‎go.mod

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -4,31 +4,30 @@ go 1.19
44

55
require (
66
github.com/cespare/xxhash/v2 v2.2.0
7-
github.com/dgraph-io/ristretto v0.1.1
8-
github.com/dustin/go-humanize v1.0.0
7+
github.com/dgraph-io/ristretto v0.1.2-0.20240116140435-c67e07994f91
8+
github.com/dustin/go-humanize v1.0.1
99
github.com/gogo/protobuf v1.3.2
1010
github.com/golang/protobuf v1.5.2
1111
github.com/google/flatbuffers v1.12.1
1212
github.com/klauspost/compress v1.15.15
1313
github.com/pkg/errors v0.9.1
1414
github.com/spf13/cobra v1.7.0
15-
github.com/stretchr/testify v1.4.0
15+
github.com/stretchr/testify v1.8.4
1616
go.opencensus.io v0.22.5
1717
golang.org/x/net v0.17.0
1818
golang.org/x/sys v0.13.0
1919
)
2020

2121
require (
2222
github.com/davecgh/go-spew v1.1.1 // indirect
23-
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b // indirect
2423
github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6 // indirect
2524
github.com/inconshreveable/mousetrap v1.1.0 // indirect
2625
github.com/pmezard/go-difflib v1.0.0 // indirect
2726
github.com/spf13/pflag v1.0.5 // indirect
2827
google.golang.org/genproto v0.0.0-20190425155659-357c62f0e4bb // indirect
2928
google.golang.org/grpc v1.20.1 // indirect
3029
google.golang.org/protobuf v1.31.0 // indirect
31-
gopkg.in/yaml.v2 v2.2.2 // indirect
30+
gopkg.in/yaml.v3 v3.0.1 // indirect
3231
)
3332

3433
retract v4.0.0 // see #1888 and #1889

‎go.sum

Lines changed: 8 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -1,22 +1,19 @@
11
cloud.google.com/go v0.26.0/go.mod h1:aQUYkXzVsufM+DwF1aE+0xfcU+56JwCaLick0ClmMTw=
22
github.com/BurntSushi/toml v0.3.1/go.mod h1:xHWCNGjB5oqiDr8zfno3MHue2Ht5sIBksp03qcyfWMU=
3-
github.com/cespare/xxhash/v2 v2.1.1/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
43
github.com/cespare/xxhash/v2 v2.2.0 h1:DC2CZ1Ep5Y4k3ZQ899DldepgrayRUGE6BBZ/cd9Cj44=
54
github.com/cespare/xxhash/v2 v2.2.0/go.mod h1:VGX0DQ3Q6kWi7AoAeZDth3/j3BFtOZR5XLFGgcrjCOs=
65
github.com/client9/misspell v0.3.4/go.mod h1:qj6jICC3Q7zFZvVWo7KLAzC3yx5G7kyvSDkc90ppPyw=
76
github.com/cpuguy83/go-md2man/v2 v2.0.2/go.mod h1:tgQtvFlXSQOSOSIRvRPT7W67SCa46tRHOmNcaadrF8o=
87
github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
98
github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
109
github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
11-
github.com/dgraph-io/ristretto v0.1.1 h1:6CWw5tJNgpegArSHpNHJKldNeq03FQCwYvfMVWajOK8=
12-
github.com/dgraph-io/ristretto v0.1.1/go.mod h1:S1GPSBCYCIhmVNfcth17y2zZtQT6wzkzgwUve0VDWWA=
13-
github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 h1:tdlZCpZ/P9DhczCTSixgIKmwPv6+wP5DGjqLYw5SUiA=
14-
github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
15-
github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=
16-
github.com/dustin/go-humanize v1.0.0/go.mod h1:HtrtbFcZ19U5GC7JDqmcUSB87Iq5E25KnS6fMYU6eOk=
10+
github.com/dgraph-io/ristretto v0.1.2-0.20240116140435-c67e07994f91 h1:Pux6+xANi0I7RRo5E1gflI4EZ2yx3BGZ75JkAIvGEOA=
11+
github.com/dgraph-io/ristretto v0.1.2-0.20240116140435-c67e07994f91/go.mod h1:swkazRqnUf1N62d0Nutz7KIj2UKqsm/H8tD0nBJAXqM=
12+
github.com/dgryski/go-farm v0.0.0-20200201041132-a6ae2369ad13 h1:fAjc9m62+UWV/WAFKLNi6ZS0675eEUC9y3AlwSbQu1Y=
13+
github.com/dustin/go-humanize v1.0.1 h1:GzkhY7T5VNhEkwH0PVJgjz+fX1rhBrR7pRT3mDkpeCY=
14+
github.com/dustin/go-humanize v1.0.1/go.mod h1:Mu1zIs6XwVuF/gI1OepvI0qD18qycQx+mFykh5fBlto=
1715
github.com/gogo/protobuf v1.3.2 h1:Ov1cvc58UF3b5XjBnZv7+opcTcQFZebYjWzi34vdm4Q=
1816
github.com/gogo/protobuf v1.3.2/go.mod h1:P1XiOD3dCwIKUDQYPy72D8LYyHL2YPYrpS2s69NZV8Q=
19-
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b h1:VKtxabqXZkF25pY9ekfRL6a582T4P37/31XEstQ5p58=
2017
github.com/golang/glog v0.0.0-20160126235308-23def4e6c14b/go.mod h1:SBH7ygxi8pfUlaOkMMuAQtPIUF8ecWP5IEl/CR7VP2Q=
2118
github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6 h1:ZgQEtGgCBiWRM39fZuwSd1LwSqqSW0hOdXCYYDX0R3I=
2219
github.com/golang/groupcache v0.0.0-20190702054246-869f871628b6/go.mod h1:cIg4eruTrX1D+g88fzRXU5OdNfaM+9IcxsU14FzY7Hc=
@@ -47,8 +44,9 @@ github.com/spf13/cobra v1.7.0/go.mod h1:uLxZILRyS/50WlhOIKD7W6V5bgeIt+4sICxh6uRM
4744
github.com/spf13/pflag v1.0.5 h1:iy+VFUOCP1a+8yFto/drg2CJ5u0yRoB7fZw3DKv/JXA=
4845
github.com/spf13/pflag v1.0.5/go.mod h1:McXfInJRrz4CZXVZOBLb0bTZqETkiAhM9Iw0y3An2Bg=
4946
github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
50-
github.com/stretchr/testify v1.4.0 h1:2E4SXV/wtOkTonXsotYi4li6zVWxYlZuYNCXe9XRJyk=
5147
github.com/stretchr/testify v1.4.0/go.mod h1:j7eGeouHqKxXV5pUuKE4zz7dFj8WfuZ+81PSLYec5m4=
48+
github.com/stretchr/testify v1.8.4 h1:CcVxjf3Q8PM0mHUKJCdn+eZZtm5yQwehR5yeSVQQcUk=
49+
github.com/stretchr/testify v1.8.4/go.mod h1:sz/lmYIOXD/1dqDmKjjqLyZ2RngseejIcXlSw2iwfAo=
5250
github.com/yuin/goldmark v1.1.27/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
5351
github.com/yuin/goldmark v1.2.1/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
5452
go.opencensus.io v0.22.5 h1:dntmOdLpSpHlVqbW5Eay97DelsZHe+55D+xC6i0dDS0=
@@ -84,7 +82,6 @@ golang.org/x/sys v0.0.0-20190215142949-d0b11bdaac8a/go.mod h1:STP8DvDyc/dI5b8T5h
8482
golang.org/x/sys v0.0.0-20190412213103-97732733099d/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
8583
golang.org/x/sys v0.0.0-20190502145724-3ef323f4f1fd/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
8684
golang.org/x/sys v0.0.0-20200930185726-fdedc70b468f/go.mod h1:h1NjWce9XRLGQEsW7wpKNCjG9DtNlClVuFLEZdDNbEs=
87-
golang.org/x/sys v0.0.0-20221010170243-090e33056c14/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
8885
golang.org/x/sys v0.13.0 h1:Af8nKPmuFypiUBjVoU9V20FiaFXOcuZI21p0ycVYYGE=
8986
golang.org/x/sys v0.13.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
9087
golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
@@ -116,7 +113,7 @@ google.golang.org/protobuf v1.31.0 h1:g0LDEJHgrBl9N9r17Ru3sqWhkIx2NB67okBHPwC7hs
116113
google.golang.org/protobuf v1.31.0/go.mod h1:HV8QOd/L58Z+nl8r43ehVNZIU/HEI6OcFqwMG9pJV4I=
117114
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+SDo693bJlVdllGtEeKM=
118115
gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
119-
gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
120116
gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
117+
gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA=
121118
gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
122119
honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=

‎table/builder_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ func TestTableIndex(t *testing.T) {
3838
key := make([]byte, 32)
3939
_, err := rand.Read(key)
4040
require.NoError(t, err)
41-
cache, err := ristretto.NewCache(&ristretto.Config{
41+
cache, err := ristretto.NewCache[uint64, *fb.TableIndex](&ristretto.Config[uint64, *fb.TableIndex]{
4242
NumCounters: 1000,
4343
MaxCost: 1 << 20,
4444
BufferItems: 64,
@@ -199,7 +199,7 @@ func BenchmarkBuilder(b *testing.B) {
199199
})
200200
b.Run("encryption", func(b *testing.B) {
201201
var opt Options
202-
cache, err := ristretto.NewCache(&ristretto.Config{
202+
cache, err := ristretto.NewCache(&ristretto.Config[uint64, *fb.TableIndex]{
203203
NumCounters: 1000,
204204
MaxCost: 1 << 20,
205205
BufferItems: 64,

‎table/iterator.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ type blockIterator struct {
3434
key []byte
3535
val []byte
3636
entryOffsets []uint32
37-
block *block
37+
block *Block
3838

3939
tableID uint64
4040
blockID int
@@ -43,7 +43,7 @@ type blockIterator struct {
4343
prevOverlap uint16
4444
}
4545

46-
func (itr *blockIterator) setBlock(b *block) {
46+
func (itr *blockIterator) setBlock(b *Block) {
4747
// Decrement the ref for the old block. If the old block was compressed, we
4848
// might be able to reuse it.
4949
itr.block.decrRef()

‎table/table.go

Lines changed: 15 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -77,8 +77,8 @@ type Options struct {
7777
Compression options.CompressionType
7878

7979
// Block cache is used to cache decompressed and decrypted blocks.
80-
BlockCache *ristretto.Cache
81-
IndexCache *ristretto.Cache
80+
BlockCache *ristretto.Cache[[]byte, *Block]
81+
IndexCache *ristretto.Cache[uint64, *fb.TableIndex]
8282

8383
AllocPool *z.AllocatorPool
8484

@@ -178,13 +178,11 @@ func (t *Table) DecrRef() error {
178178
}
179179

180180
// BlockEvictHandler is used to reuse the byte slice stored in the block on cache eviction.
181-
func BlockEvictHandler(value interface{}) {
182-
if b, ok := value.(*block); ok {
183-
b.decrRef()
184-
}
181+
func BlockEvictHandler(b *Block) {
182+
b.decrRef()
185183
}
186184

187-
type block struct {
185+
type Block struct {
188186
offset int
189187
data []byte
190188
checksum []byte
@@ -199,7 +197,7 @@ var NumBlocks atomic.Int32
199197

200198
// incrRef increments the ref of a block and return a bool indicating if the
201199
// increment was successful. A true value indicates that the block can be used.
202-
func (b *block) incrRef() bool {
200+
func (b *Block) incrRef() bool {
203201
for {
204202
// We can't blindly add 1 to ref. We need to check whether it has
205203
// reached zero first, because if it did, then we should absolutely not
@@ -222,7 +220,7 @@ func (b *block) incrRef() bool {
222220
}
223221
}
224222
}
225-
func (b *block) decrRef() {
223+
func (b *Block) decrRef() {
226224
if b == nil {
227225
return
228226
}
@@ -242,12 +240,12 @@ func (b *block) decrRef() {
242240
}
243241
y.AssertTrue(b.ref.Load() >= 0)
244242
}
245-
func (b *block) size() int64 {
243+
func (b *Block) size() int64 {
246244
return int64(3*intSize /* Size of the offset, entriesIndexStart and chkLen */ +
247245
cap(b.data) + cap(b.checksum) + cap(b.entryOffsets)*4)
248246
}
249247

250-
func (b *block) verifyCheckSum() error {
248+
func (b *Block) verifyCheckSum() error {
251249
cs := &pb.Checksum{}
252250
if err := proto.Unmarshal(b.checksum, cs); err != nil {
253251
return y.Wrapf(err, "unable to unmarshal checksum for block")
@@ -521,7 +519,7 @@ func (t *Table) fetchIndex() *fb.TableIndex {
521519
panic("Index Cache must be set for encrypted workloads")
522520
}
523521
if val, ok := t.opt.IndexCache.Get(t.indexKey()); ok && val != nil {
524-
return val.(*fb.TableIndex)
522+
return val
525523
}
526524

527525
index, err := t.readTableIndex()
@@ -537,7 +535,7 @@ func (t *Table) offsets(ko *fb.BlockOffset, i int) bool {
537535
// block function return a new block. Each block holds a ref and the byte
538536
// slice stored in the block will be reused when the ref becomes zero. The
539537
// caller should release the block by calling block.decrRef() on it.
540-
func (t *Table) block(idx int, useCache bool) (*block, error) {
538+
func (t *Table) block(idx int, useCache bool) (*Block, error) {
541539
y.AssertTruef(idx >= 0, "idx=%d", idx)
542540
if idx >= t.offsetsLength() {
543541
return nil, errors.New("block out of index")
@@ -549,15 +547,15 @@ func (t *Table) block(idx int, useCache bool) (*block, error) {
549547
// Use the block only if the increment was successful. The block
550548
// could get evicted from the cache between the Get() call and the
551549
// incrRef() call.
552-
if b := blk.(*block); b.incrRef() {
553-
return b, nil
550+
if blk.incrRef() {
551+
return blk, nil
554552
}
555553
}
556554
}
557555

558556
var ko fb.BlockOffset
559557
y.AssertTrue(t.offsets(&ko, idx))
560-
blk := &block{offset: int(ko.Offset())}
558+
blk := &Block{offset: int(ko.Offset())}
561559
blk.ref.Store(1)
562560
defer blk.decrRef() // Deal with any errors, where blk would not be returned.
563561
NumBlocks.Add(1)
@@ -795,7 +793,7 @@ func NewFilename(id uint64, dir string) string {
795793
}
796794

797795
// decompress decompresses the data stored in a block.
798-
func (t *Table) decompress(b *block) error {
796+
func (t *Table) decompress(b *Block) error {
799797
var dst []byte
800798
var err error
801799

‎table/table_test.go

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -705,7 +705,7 @@ func TestTableChecksum(t *testing.T) {
705705
})
706706
}
707707

708-
var cacheConfig = ristretto.Config{
708+
var cacheConfig = ristretto.Config[[]byte, *Block]{
709709
NumCounters: 1000000 * 10,
710710
MaxCost: 1000000,
711711
BufferItems: 64,
@@ -848,7 +848,7 @@ func BenchmarkRandomRead(b *testing.B) {
848848
}
849849
}
850850

851-
func getTableForBenchmarks(b *testing.B, count int, cache *ristretto.Cache) *Table {
851+
func getTableForBenchmarks(b *testing.B, count int, cache *ristretto.Cache[[]byte, *Block]) *Table {
852852
rand.Seed(time.Now().Unix())
853853
opts := Options{Compression: options.ZSTD, BlockSize: 4 * 1024, BloomFalsePositive: 0.01}
854854
if cache == nil {

0 commit comments

Comments
 (0)