cache: mark cache entry as referenced when multiple readers

Author: sumeerbhola

internal/cache/cache.go

@@ -322,9 +322,14 @@ func (c *Handle) GetWithReadHandle(
 // existing value if present. The value must have been allocated by Cache.Alloc.
 //
 // The cache takes a reference on the Value and holds it until it gets evicted.
+//
+// REQUIRES: value.refs() == 1
 func (c *Handle) Set(fileNum base.DiskFileNum, offset uint64, value *Value) {
+	if n := value.refs(); n != 1 {
+		panic(fmt.Sprintf("pebble: Value has already been added to the cache: refs=%d", n))
+	}
 	k := makeKey(c.id, fileNum, offset)
-	c.cache.getShard(k).set(k, value)
+	c.cache.getShard(k).set(k, value, false /*markAccessed*/)
 }
 
 // Delete deletes the cached value for the specified file and offset.

internal/cache/clockpro.go

@@ -158,11 +158,7 @@ func (c *shard) getWithMaybeReadEntry(k key, desireReadEntry bool) (*Value, *rea
 	return value, re
 }
 
-func (c *shard) set(k key, value *Value) {
-	if n := value.refs(); n != 1 {
-		panic(fmt.Sprintf("pebble: Value has already been added to the cache: refs=%d", n))
-	}
-
+func (c *shard) set(k key, value *Value, markAccessed bool) {
 	c.mu.Lock()
 	defer c.mu.Unlock()
 
@@ -173,6 +169,9 @@ func (c *shard) set(k key, value *Value) {
 		// no cache entry? add it
 		e = newEntry(k, int64(len(value.buf)))
 		e.setValue(value)
+		if markAccessed {
+			e.referenced.Store(true)
+		}
 		if c.metaAdd(k, e) {
 			value.ref.trace("add-cold")
 			c.sizeCold += e.size
@@ -219,7 +218,7 @@ func (c *shard) set(k key, value *Value) {
 			c.coldTarget = c.targetSize()
 		}
 
-		e.referenced.Store(false)
+		e.referenced.Store(markAccessed)
 		e.setValue(value)
 		e.ptype = etHot
 		if c.metaAdd(k, e) {

internal/cache/read_shard.go

@@ -6,6 +6,7 @@ package cache
 
 import (
 	"context"
+	"fmt"
 	"sync"
 	"time"
 
@@ -293,15 +294,10 @@ func (e *readEntry) unrefAndTryRemoveFromMap() {
 }
 
 func (e *readEntry) setReadValue(v *Value) {
-	// Add to the cache before taking another ref for readEntry, since the cache
-	// expects ref=1 when it is called.
-	//
-	// TODO(sumeer): if e.refCount > 1, we should consider overriding to ensure
-	// that it is added as etHot. The common case will be e.refCount = 1, and we
-	// don't want to acquire e.mu twice, so one way to do this would be relax
-	// the invariant in shard.Set that requires Value.refs() == 1. Then we can
-	// do the work under e.mu before calling shard.Set.
-	e.readShard.shard.set(e.key, v)
+	if n := v.refs(); n != 1 {
+		panic(fmt.Sprintf("pebble: Value has already been added to the cache: refs=%d", n))
+	}
+	concurrentRequesters := false
 	e.mu.Lock()
 	// Acquire a ref for readEntry, since we are going to remember it in e.mu.v.
 	v.acquire()
@@ -318,8 +314,12 @@ func (e *readEntry) setReadValue(v *Value) {
 		// readEntry.waitForReadPermissionOrHandle, and those will also use
 		// e.mu.v.
 		close(e.mu.ch)
+		// e.mu.ch is non-nil only when there were concurrent requesters. NB: we
+		// can't read e.refCount here since it is protected by e.readShard.mu.
+		concurrentRequesters = true
 	}
 	e.mu.Unlock()
+	e.readShard.shard.set(e.key, v, concurrentRequesters)
 	e.unrefAndTryRemoveFromMap()
 }
 
@@ -368,6 +368,8 @@ func (rh ReadHandle) Valid() bool {
 //
 // The cache takes a reference on the Value and holds it until it is evicted and
 // no longer needed by other readers.
+//
+// REQUIRES: v.refs() == 1
 func (rh ReadHandle) SetReadValue(v *Value) {
 	rh.entry.setReadValue(v)
 }

internal/cache/read_shard_test.go

@@ -5,10 +5,13 @@
 package cache
 
 import (
+	"cmp"
 	"context"
 	crand "crypto/rand"
 	"fmt"
 	"math/rand"
+	"slices"
+	"strings"
 	"sync"
 	"testing"
 	"time"
@@ -190,6 +193,34 @@ func TestReadShard(t *testing.T) {
 				time.Sleep(10 * time.Millisecond)
 				return fmt.Sprintf("map-len: %d", c.readShard.lenForTesting())
 
+			case "print-shard":
+				return func() string {
+					c.mu.RLock()
+					defer c.mu.RUnlock()
+					type shardEntry struct {
+						k          key
+						hasValue   bool
+						referenced bool
+					}
+					var entries []shardEntry
+					c.blocks.All(func(k key, e *entry) bool {
+						entries = append(entries,
+							shardEntry{k: k, hasValue: e.val != nil, referenced: e.referenced.Load()})
+						return true
+					})
+					slices.SortFunc(entries, func(a, b shardEntry) int {
+						return cmp.Or(
+							cmp.Compare(a.k.id, b.k.id), cmp.Compare(a.k.fileNum, b.k.fileNum),
+							cmp.Compare(a.k.offset, b.k.offset))
+					})
+					var b strings.Builder
+					for _, e := range entries {
+						fmt.Fprintf(&b, "id=%d file=%d offset=%d hasValue=%t referenced=%t\n",
+							e.k.id, e.k.fileNum, e.k.offset, e.hasValue, e.referenced)
+					}
+					return b.String()
+				}()
+
 			default:
 				return fmt.Sprintf("unknown command: %s", td.Cmd)
 

internal/cache/testdata/read_shard

@@ -153,6 +153,12 @@ set-read-value name=s1_f2_off3 val=brinjal
 ----
 map-len: 0
 
+# The added entry is marked as referenced, since there is a concurrent reader
+# (that is still waiting).
+print-shard
+----
+id=1 file=2 offset=3 hasValue=true referenced=true
+
 wait name=s1_f2_off3_2
 ----
 val: brinjal
@@ -161,3 +167,24 @@ map-len: 0
 get name=s1_f2_off3_3 id=1 file-num=2 offset=3
 ----
 val: brinjal
+
+# Get misses the cache and has turn to read.
+get name=s1_f2_off20 id=1 file-num=2 offset=20
+----
+waiting
+map-len: 1
+
+wait name=s1_f2_off20
+----
+turn to read
+map-len: 1
+
+set-read-value name=s1_f2_off20 val=okra
+----
+map-len: 0
+
+# The added entry is not referenced, since there was not a concurrent reader.
+print-shard
+----
+id=1 file=2 offset=3 hasValue=true referenced=true
+id=1 file=2 offset=20 hasValue=true referenced=false