Repository: dgraph-io/badger
Pull request #1551 (2020-10-01 22:56:05)

Allow sizes of block and index caches to be updated.


<!-- Reviewable:start -->
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/dgraph-io/badger/1551)
<!-- Reviewable:end -->



Patch:
diff --git a/db.go b/db.go
index f0686e7..7db5216 100644
--- a/db.go
+++ b/db.go
@@ -1801,3 +1801,42 @@ func (db *DB) StreamDB(outOptions Options) error {
 func (db *DB) Opts() Options {
 	return db.opt
 }
+
+type CacheType int
+
+const (
+	BlockCache CacheType = iota
+	IndexCache
+)
+
+// CacheMaxCost updates the max cost of the given cache (either block or index cache).
+// The call will have an effect only if the DB was created with the cache. Otherwise it is
+// a no-op. If you pass a negative value, the function will return the current value
+// without updating it.
+func (db *DB) CacheMaxCost(cache CacheType, maxCost int64) (int64, error) {
+	if db == nil {
+		return 0, nil
+	}
+
+	if maxCost < 0 {
+		switch cache {
+		case BlockCache:
+			return db.blockCache.MaxCost(), nil
+		case IndexCache:
+			return db.indexCache.MaxCost(), nil
+		default:
+			return 0, errors.Errorf("invalid cache type")
+		}
+	}
+
+	switch cache {
+	case BlockCache:
+		db.blockCache.UpdateMaxCost(maxCost)
+		return maxCost, nil
+	case IndexCache:
+		db.indexCache.UpdateMaxCost(maxCost)
+		return maxCost, nil
+	default:
+		return 0, errors.Errorf("invalid cache type")
+	}
+}
diff --git a/go.mod b/go.mod
index ec96847..a0a0fb9 100644
--- a/go.mod
+++ b/go.mod
@@ -7,7 +7,7 @@ go 1.12
 require (
 	github.com/DataDog/zstd v1.4.1
 	github.com/cespare/xxhash v1.1.0
-	github.com/dgraph-io/ristretto v0.0.4-0.20201012160933-079c5f0d0daf
+	github.com/dgraph-io/ristretto v0.0.4-0.20201007164332-9739cfa2564b
 	github.com/dustin/go-humanize v1.0.0
 	github.com/golang/protobuf v1.3.1
 	github.com/golang/snappy v0.0.1
diff --git a/go.sum b/go.sum
index 863a7fa..ccd9e6f 100644
--- a/go.sum
+++ b/go.sum
@@ -13,8 +13,8 @@ github.com/cpuguy83/go-md2man v1.0.10/go.mod h1:SmD6nW6nTyfqj6ABTjUi3V3JVMnlJmwc
 github.com/davecgh/go-spew v1.1.0/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c=
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
-github.com/dgraph-io/ristretto v0.0.4-0.20201012160933-079c5f0d0daf h1:FccUHvo6MMTYz+SrPSjgFuwp7W2376H34YVSzm8qjUY=
-github.com/dgraph-io/ristretto v0.0.4-0.20201012160933-079c5f0d0daf/go.mod h1:bDI4cDaalvYSji3vBVDKrn9ouDZrwN974u8ZO/AhYXs=
+github.com/dgraph-io/ristretto v0.0.4-0.20201007164332-9739cfa2564b h1:CfEwAh85rqfhhE/DrDCbFdcwVpZ8ESnFRhLatiKdDYM=
+github.com/dgraph-io/ristretto v0.0.4-0.20201007164332-9739cfa2564b/go.mod h1:bDI4cDaalvYSji3vBVDKrn9ouDZrwN974u8ZO/AhYXs=
 github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2 h1:tdlZCpZ/P9DhczCTSixgIKmwPv6+wP5DGjqLYw5SUiA=
 github.com/dgryski/go-farm v0.0.0-20190423205320-6a90982ecee2/go.mod h1:SqUrOPUnsFjfmXRMNPybcSiG0BgUW2AuFH8PAnS2iTw=
 github.com/dustin/go-humanize v1.0.0 h1:VSnTsYCnlFHaM2/igO1h6X3HA71jcobQuxemgkq4zYo=


Test patch:
diff --git a/db_test.go b/db_test.go
index b1b6c91..8ef6d35 100644
--- a/db_test.go
+++ b/db_test.go
@@ -2157,3 +2157,33 @@ func TestMinCacheSize(t *testing.T) {
 	}()
 
 }
+
+func TestUpdateMaxCost(t *testing.T) {
+	dir, err := ioutil.TempDir("", "badger-test")
+	require.NoError(t, err, "temp dir for badger count not be created")
+	defer os.RemoveAll(dir)
+
+	ops := getTestOptions(dir).
+		WithBlockCacheSize(1 << 20).
+		WithIndexCacheSize(2 << 20)
+	db, err := Open(ops)
+	require.NoError(t, err)
+
+	cost, err := db.CacheMaxCost(BlockCache, -1)
+	require.NoError(t, err)
+	require.Equal(t, int64(1<<20), cost)
+	cost, err = db.CacheMaxCost(IndexCache, -1)
+	require.NoError(t, err)
+	require.Equal(t, int64(2<<20), cost)
+
+	_, err = db.CacheMaxCost(BlockCache, 2<<20)
+	require.NoError(t, err)
+	cost, err = db.CacheMaxCost(BlockCache, -1)
+	require.NoError(t, err)
+	require.Equal(t, int64(2<<20), cost)
+	_, err = db.CacheMaxCost(IndexCache, 4<<20)
+	require.NoError(t, err)
+	cost, err = db.CacheMaxCost(IndexCache, -1)
+	require.NoError(t, err)
+	require.Equal(t, int64(4<<20), cost)
+}