Skip to content

Commit c8f4b55

Browse files
committed
Update ristretto and refactor for use of generics
Update ristretto to the latest revision and refactor for the new use of generics in the cache types within. Signed-off-by: Christian Stewart <christian@aperture.us>
1 parent 6acc8e8 commit c8f4b55

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)