feat(store): Add memory LRU cache by Anton-Kalpakchiev · Pull Request #647 · uber/kraken

Anton-Kalpakchiev · 2026-07-25T13:06:01Z

Read #633 first for more context.

thijmv · 2026-07-28T12:30:10Z

+// while [Store]'s APIs omit that arg (cleaner interface) and instead expose other APIs to scope the whole store.
+//
+// Check [Store]'s comments for details on functionality.
+type store struct {


Also briefly discussed this with @iawakimjan offline, but the implementations for the disk and memory stores are quite similar. I am wondering, did you consider abstracting the storage layer and using a single cache policy orchestrator?

I'm not sure if it's worth it as IMO there are subtle differences in the 2 storages that might not be easily abstractable. I am waiting for @iawakimjan 's review to see if he sees a way to abstract these 2 stores.

Ya I also feel that if there's a possibility of abstraction, we should implement it.
Atleast, I feel the methods are same so they should be an implementation of an interface.

The interface is mostly the same as the disk.Store with 1 big difference: entries that get evicted in the disk store but are still opened by users do not actually get deleted from disk, getting removed only once all clients close their handles. This is intentional and is handled by Linux. It leads to a bit of overreservation disk-wise, which is why we will leave a buffer when picking the store's capacity. However, the same behavior is much more dangerous when used in a memory store due to the danger of OOMs. Therefore, we have 2 options: 1. don't evict entries that are evictable but held by clients 2. evict the entries AND make sure the memory is actually freed. We pick option 2, as a crucial part of the mem store's design is to prioritize writes to the store over reads from it, as writes are origin's bottleneck.

- The APIs are very similar to the disk.Store's so I decided to implement them in a single commit. - Decided to change Stat's API to just return the size of a blob and not [os.FileInfo], as returning ps.FileInfo requires more bookkeeping (e.g. tracking the last access time of a blob) which would add redundant complexity, as clients never ever actually use Stat for anything else than 1) getting the blob's size and 2) getting the LAT of the blob, HOWEVER, they only need it for eviction, which after this refactor won't be necessary outside the store. - tests will be added in the next commit

- Follows the same pattern as disk.Store() - I extracted the common logic into lib/store/scope.go

I copypasted the tests for the disk store and adapted them for the memory store. I also added 2 extra tests that are important for the mem store. The first one is that after eviction/deletion of a blob, the handle returned to the client returns ErrEvicted. The second test was about a bug, where i was doing `data := make([]byte, size)` instead of `data := make([]byte, 0, size)` in Create. This meant that when clients overreport a blob's size (this happens) and they later read the blob, they would read trailing empty space at the end. Fixed the bug and added a test for under- and overreporting blob size. Added the test to the disk store as well. Finally, I added a log and a fail-open `if` case when releasing space (when bookkeeping the store's size) to ensure that we catch any bugs where adding and subsequently removing a blob changes store's size (used to be a bug in the past).

sambhav-jain-16 · 2026-07-29T12:53:51Z

+var _ storelib.FileReadWriter = &handle{}
+
+// handle is the struct returned to clients when they create/open a blob.
+// As soon as the blob's data is evicted, the handle no longer has a reference to it,
+// ensuring GC can clean it.
+type handle struct {
+	data    *atomic.Pointer[[]byte]
+	sliceMu *sync.Mutex // A potential optimization if contention is too high: currently all writes are not parallelized. However, writes that only mutate the array but not the slice are parallelizable with each other. Thus, we could transition to a RWMutex to enable that.
+	off     int64
+}
+
+func newHandle(data *atomic.Pointer[[]byte], sliceMu *sync.Mutex) *handle {
+	return &handle{
+		data:    data,
+		sliceMu: sliceMu,
+		off:     0,
+	}
+}
+
+func (h *handle) getData() (data []byte, evicted bool) {
+	buf := h.data.Load()
+	if buf == nil {
+		return nil, true
+	}
+	return *buf, false
+}
+
+func (h *handle) Read(p []byte) (n int, err error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	buf, evicted := h.getData()
+	if evicted {
+		return 0, ErrEvicted
+	}
+
+	if h.off >= int64(len(buf)) {
+		return 0, io.EOF
+	}
+	n = copy(p, buf[h.off:])
+	h.off += int64(n)
+	return n, nil
+}
+
+// ReadAt implements [io.ReaderAt]. Thread-safe.
+func (h *handle) ReadAt(p []byte, off int64) (n int, err error) {
+	if len(p) == 0 {
+		return 0, nil
+	}
+	if off < 0 {
+		return 0, errors.New("negative offset")
+	}
+
+	buf, evicted := h.getData()
+	if evicted {
+		return 0, ErrEvicted
+	}
+	if off >= int64(len(buf)) {
+		return 0, io.EOF
+	}
+
+	n = copy(p, buf[off:])
+	if n < len(p) {
+		return n, io.EOF
+	}
+	return n, nil
+}
+
+// Seek implements [io.Seeker]. Not thread-safe.
+func (h *handle) Seek(off int64, whence int) (int64, error) {
+	buf, evicted := h.getData()
+	if evicted {
+		return 0, ErrEvicted
+	}
+
+	var newOff int64
+	switch whence {
+	case io.SeekStart:
+		newOff = off
+	case io.SeekCurrent:
+		newOff = h.off + off
+	case io.SeekEnd:
+		newOff = int64(len(buf)) + off
+	default:
+		return 0, errors.New("invalid whence")
+	}
+
+	if newOff < 0 || newOff > int64(len(buf)) {
+		return 0, errors.New("invalid seek location")
+	}
+	h.off = newOff
+	return newOff, nil
+}
+
+// Stat returns the blob's actual size, even if it differs from the size reported during Create.
+func (h *handle) Size() int64 {
+	buf, evicted := h.getData()
+	if evicted {
+		// TODO - consider whether this is ok or whether we need to store the user-provided blob size in [*handle].
+		return 0
+	}
+	return int64(len(buf))
+}
+
+// WriteAt implements [io.WriterAt]. It is fully thread-safe.
+func (h *handle) WriteAt(p []byte, off int64) (n int, err error) {
+	if off < 0 {
+		return 0, errors.New("negative offset")
+	}
+
+	h.sliceMu.Lock()
+	defer h.sliceMu.Unlock()
+
+	buf, evicted := h.getData()
+	if evicted {
+		return 0, ErrEvicted
+	}
+
+	end := int(off) + len(p)
+	buf, resized := resizeSliceIfNecessary(buf, end)
+	n = copy(buf[off:], p)
+	if resized {
+		h.data.Store(&buf)
+	}
+	return n, nil
+}
+
+func resizeSliceIfNecessary(buf []byte, end int) ([]byte, bool) {
+	resized := false
+	if len(buf) < end {
+		if cap(buf) < end {
+			newBuf := make([]byte, end)
+			copy(newBuf, buf)
+			buf = newBuf
+		}
+		buf = buf[:end]
+		resized = true
+	}
+	return buf, resized
+}
+
+// Write implements io.Writer.
+func (h *handle) Write(p []byte) (n int, err error) {
+	h.sliceMu.Lock() // We need to lock in case we update the pointer to `data`.
+	defer h.sliceMu.Unlock()
+
+	buf, evicted := h.getData()
+	if evicted {
+		return 0, ErrEvicted
+	}
+
+	end := int(h.off) + len(p)
+	buf, resized := resizeSliceIfNecessary(buf, end)
+
+	n = copy(buf[h.off:], p)
+	if resized {
+		h.data.Store(&buf)
+	}
+	h.off += int64(n)
+	return n, nil
+}
+
+func (h *handle) Cancel() error { return nil } // no-op
+func (h *handle) Close() error  { return nil } // no-op
+func (h *handle) Commit() error { return nil } // no-op


I'm curious: shouldn't we just reuse lib/store/base/BufferReadWriter since the implementation is almost similar?

I used BufferReadWriter as inspiration, but it doesn't support the ErrEvicted use case, so I decided to write a replacement, which I don't think is avoidable during the migration period, where we have both CAStore and disk.Store + mem.Store in the codebase. WDYT?

sambhav-jain-16 · 2026-07-29T12:57:59Z

+// while [Store]'s APIs omit that arg (cleaner interface) and instead expose other APIs to scope the whole store.
+//
+// Check [Store]'s comments for details on functionality.
+type store struct {


Ya I also feel that if there's a possibility of abstraction, we should implement it.
Atleast, I feel the methods are same so they should be an implementation of an interface.

sambhav-jain-16 · 2026-07-29T13:00:37Z

+	off     int64
+}
+
+func newHandle(data *atomic.Pointer[[]byte], sliceMu *sync.Mutex) *handle {


Any reason to pass a mutex from the caller? This can cause unnecessary locking in case the caller accidentally uses the same locks for the handle and other work

The only caller is the memory.Store which needs to synchronize with the handle as to prevent a race on the blob's data. If the store misuses the mutex, that's just a bug in the store. The mutex is encapsulated within the memory package and therefore unreachable from clients of the package.

Check the comment on sliceMu in the blob struct for more info.

sambhav-jain-16 · 2026-07-29T13:03:05Z

+}
+
+type blob struct {
+	data           atomic.Pointer[[]byte] // set to nil upon eviction/deletion.


Any reason to use atomic?

I considered doing data *[]byte instead of data *atomic.Pointer[[]byte] and having eviction be done with *data = nil instead of data.Store(nil), however, this means that now reads must call sliceMu.RLock, which protects again a race between buf := *h.data in read APIs and *data = nil when evicting the data and *data = newBuf (not an atomic operation, this actually mutates len, cap, and the slice's underlying array non-atomically).

Having said that, perhaps it would be simpler to just use a RWMutex, hmm... The original reason I went for an atomic is because I wanted writes to be parallelizable, but we pivoted away from that idea.

Yeah. let's try to be consistent. Either we move fully to optimistic locking or fully to mutexes.
In the current implementation, it is better to use RLock rather than using atomic

sambhav-jain-16 · 2026-07-29T13:05:49Z

+	b.sliceMu.Lock()
+	b.data.Store(nil) // Ensure the byte slice is not referenced by clients outside the store holding [*handle], so GC can evict the memory.
+	b.sliceMu.Unlock()


Locking inside a lock is a potential deadlock scenario, IMO. Is there a way we can avoid this?

Hmm, at first I also considered doing CAS, but I decided it would be simpler to have sliceMu be responsible for synchronizing mutations to the data slice and not share ownership of that between sync primitives.

Also, right now there can't be a deadlock, as no operation acquires sliceMu first and then store.mu.

WDYT?

Yeah, let's not consider CAS for this. I also don't see any cases, but I wonder if it might happen in the future if code is added to the store.

- add Has method for disk.Store and mem.Store (needed by tiered.Store) - rename each store's `handle`` to `File` and return the File struct instead of the store.FileReadWriter interface in the Create and Open APIs of the stores (in accordance with Golang's principle of returning) a struct and accepting an interface. - Add method `Off` to memory.File that returns the current offset. ( needed in tiered.Store) - move the disk handle implementation from lib/store/file.go to lib/store/disk/file.go - add tests for the new behavior

github-actions Bot added size/m size/xl labels Jul 25, 2026

Anton-Kalpakchiev force-pushed the lru-cache-mem branch 2 times, most recently from e247a4a to 1186526 Compare July 27, 2026 13:15

Anton-Kalpakchiev requested review from sambhav-jain-16 and thijmv July 27, 2026 19:44

Anton-Kalpakchiev force-pushed the lru-cache-mem branch 2 times, most recently from c5fe6f7 to aa1df50 Compare July 28, 2026 14:55

thijmv reviewed Jul 28, 2026

View reviewed changes

Anton-Kalpakchiev added 7 commits July 28, 2026 17:18

feat(memory.Store): Implement memory blob handle

adae87b

feat(memory.Store): add Open and Delete methods

e7833af

feat(memory.Store): Implement scoping

4b9703f

- Follows the same pattern as disk.Store() - I extracted the common logic into lib/store/scope.go

feat: Emit size and num_blobs in disk.Store and mem.Store

4f50aed

Anton-Kalpakchiev force-pushed the lru-cache-mem branch from aa1df50 to 80c0230 Compare July 28, 2026 15:23

Anton-Kalpakchiev mentioned this pull request Jul 28, 2026

Refactor store package for performance and simplicity #633

Open

perf(memory.Store): optimize handle.WriteAt

b69f5d8

Anton-Kalpakchiev force-pushed the lru-cache-mem branch from 8f92134 to b69f5d8 Compare July 28, 2026 17:39

sambhav-jain-16 reviewed Jul 29, 2026

View reviewed changes

Anton-Kalpakchiev requested review from sambhav-jain-16 and thijmv July 29, 2026 17:08

Uh oh!

Conversation

Anton-Kalpakchiev commented Jul 25, 2026

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

sambhav-jain-16 Jul 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

sambhav-jain-16 Jul 30, 2026 •

edited

Loading