Refactor `store` package for performance and simplicity by Anton-Kalpakchiev · Pull Request #633 · uber/kraken

Anton-Kalpakchiev · 2026-06-06T14:38:32Z

TLDR

This PR is the first of several to refactor the store package to 1) be more performant (transition from a TTI to an LRU eviction policy). and 2) be much simpler (go from 8K to <3K LOC, remove unnecessary abstractions, etc.). Please review each commit and its description separately (the PR itself will become quite large).

NOTE - this PR is not plugged into production code, that will be done in a separate PR.

Justification

Why transition from a TTI to an LRU eviction policy (copied from this ERD)

Each origin replica in the hashring has a disk and a memory cache. Both caches are severely underutilized, bleeding performance. The disk cache has 30% util, while the mem cache has <10% util. By refactoring both caches to be LRU-based, we will achieve 100% utilization, significantly improving kraken-origin performance. There is absolutely 0 reason not to do this, as Kraken’s data is immutable, therefore it is optimal to evict data only when necessary, i.e. when making room to cache newer data.

Why refactor the code

The store package has balooned to over ~8k lines with several layers of abstraction. I believe this is way beyond the package's inherent complexity. Changing the package, as proposed in this PR is non-trivial. At this point, it's easier to rewrite it from scratch, reusing existing parts wherever possible. I expect this to 1) significantly simplify the package's interface and 2) reduce the package's internal complexity (from ~8K LOC to <3K).

Implementation (copied from the ERD)

Currently, the cache works as follows:

If the blob is in either caches, serve it to the user.
Otherwise, download it in the memory cache, serve it to the user, replicate it to the disk cache, and evict it from the mem cache.
Data is evicted from disk using a 1h TTI.
Data is evicted from memory as soon as it’s replicated to disk.
If the memory cache is full and a new download request comes in, we fallback to disk (the mem cache is not used for the download).

I propose the following changes:

We only evict data from memory using an LRU to free space for more data.
We only evict data from disk using an LRU to free space for more data.
We cannot evict blobs from the mem cache that are not yet replicated to the disk cache

The new flow is illustrated below:

sambhav-jain-16

I have a bit of nits and some questions. Overall the implementation looks good.
I do have an overall comment on the test files, Will it be possible to refactor them to table driven formats and avoid redundant code?

sambhav-jain-16 · 2026-07-27T10:52:21Z

+		err := os.RemoveAll(incompleteDirPath)
+		if err != nil {
+			return nil, fmt.Errorf("remove incomplete blobs left from a previous service run: %w", err)
+		}


Should we add retry logic here to avoid failing on a random flake?

My understanding is that disk APIs don't have transient errors, so if we get an error, it is most likely a real error, so a retry might only put extra load on the disk. Therefore, I decided to "fail fast and early" to surface the error/bug, considering this bug happens on application startup, when a rollback would still be possible. WDYT?

sambhav-jain-16 · 2026-07-27T11:27:56Z

+//   - Crash-resistant - all state is restored upon restart (check [newStore] for details).
+//
+//   - Supports directory sharding to speed up disk performance.
+type Store struct {


nit: maybe rename to ScopedStore?

This is the store that is exported to clients. So the import currently looks like disk.Store. I think disk.ScopedStore might be misleading, as it would imply there is another non-scoped store, but there isn't - clients have one single DiskStore to operate on and this is it, so I decided to just name it Store. The reason the file is called scoped_store.go is that internally within the disk package we have a non-scoped implementation, but it's not something we expose outside the package. WDYT?

Anton-Kalpakchiev · 2026-07-27T13:11:26Z

I have a bit of nits and some questions. Overall the implementation looks good. I do have an overall comment on the test files, Will it be possible to refactor them to table driven formats and avoid redundant code?

I agree there's a lot of repetition in the code, but am not sure if table tests can fix this - each test case does something a little different. For example, one test case might call MarkComplete, another might call BanEviction, a third might write 1KB of data instead of 2KB, etc. We could transition to table tests, but then the test code itself will have 4 if statements (using this example), which would make it much less readable:

Table tests can be difficult to read and maintain if the subtests contain conditional assertions or other branching logic. Table tests should NOT be used whenever there needs to be complex or conditional logic inside subtests (i.e. complex logic inside the for loop).

thijmv

Please also fix the lint errors before merging.

This commit designs the client-facing interface of the disk store and implements its core APIs. It establishes the core data structures used. It also adds tests.

- Implement the BanEviction and UnbanEviction APIs of the store that stop a blob from getting evicted - Finish other APIs that were half-implemented due to unevictable blobs not being implemented. - Change the API of Delete to now be able to delete evictable blobs. I looked through the current use of `Delete` for the ca_store and the only time it checked if it was trying to delete persisted blobs was during cleanup. Now that the store's APIs themselves do cleanup, there is no reason for Delete to respect whether a blob is evictable or not. - Add **extensive** testing for 1) the eviction logic and 2) other APIs. - Make some actions no-ops instead of returning errors (e.g. calling MarkComplete on an already complete blob) to make the interface more forgiving. The next commits will: - implement Metadata operations - add logic to recover state from disk after a crash.

Previously, the store used 3 maps to differentiate between different blob types: - complete blobs - unevictable blobs - incomplete blobs This was not a clean solution, as now all key-value APIs now needed to check all 3 maps for the existence of the key. Also APIs that transitioned a blob from one blob type to another required moving the blob between the maps. To remove this complexity, now we use a single map to manage all blobs. Instead of determining the blob's type by checking which map it belongs to, we create a `blob` struct with the blob's necessary data and store that in the map. I also fix issues in tests and improve the style.

As DiskStore is intended to be persistent, it should be able to survive crashes/restarts by rebooting any necessary state from disk to memory. This was already implemented, however, it was assumed that incomplete blobs do not need to be rebooted, as it was expected that the store users would forget about the incomplete blobs, resulting in a leak. This turned out to be false - while almost all Kraken services' storages drop incomplete blobs on system crash, the agent actually persists incomplete blobs, so it can resume downloading them after it comes back up. As DiskStore is intended to be a single DiskStore reused by all Kraken services that need disk storage, this commit adds support to reboot incomplete blobs after a crash. To do so, the client-provided size of the blob must be persisted on disk (currently implemented in a sidecar file), which this commit implements. The following alternatives were evaluated but discarded: 1. removing the agent's ability to resume downloads after a crash - will degrade p99 performance 2. truncating files to persist their size instead of using a sidecar file - client-provided blob sizes might differ by a few bytes from actual sizes 3. using `xattr` instead of a sidecar file - xattr is not supported on all filesystems 4. moving toward async eviction that measures the size of the directory directly - discarded, as measuring the size of a huge directory can be extremely slow (10s of minutes), which would block eviction and cause disk exhaustion AND could add extra disk IO pressure

Agent, origin, and build-index currently use CAStore (the disk cache implementation that DiskStore is replacing) and more specifically its functionality to scope an API to work on only a complete blob or an incomplete blob. Without this functionality, it is much harder to keep their code correct. Therefore, DiskStore must expose the same functionality in order to replace CAStore. This PR adds that logic. At first, I decided each API can take an extra arg for the scope, e.g. the following code would delete the blob with key `b7fe1643` only if it is complete. Otherwise, an error is returned. ```go key := "b7fe1643" err := store.Delete(key, ScopeCompleteBlobs) ``` This works, however, it adds complexity - each API call now requires an extra arg, while the user might not necessarily want to scope the API call. Therefore, I opted for the implementation in this PR, where we use constructor-like APIs like this: ```go key := "b7fe1643" err := store.ScopeComplete().Delete(key) // works only on complete blob err := store.ScopeIncomplete().Delete(key) // works only on incomplete blob err := store.Delete(key) // works on any blob ``` Now the APIs have a simpler interface (1 fewer arg) and scoping is optional.

- Add ListMetadata API (needed by proxy) - Add WriteAtMetadata API (needed by agent) - clean immovable metadata after marking a blob as complete. For context, some metadata is supposed to stay as long as the blob is alive, e.g. the MetaInfo metadata. However, other metadata is only needed while a blob is being downloaded into the DiskStore and can be discarded once the download is complete, e.g. `startedAtMetadata` (used to keep track of when an upload to the DiskStore started to enforce a TTL timeout). Up until now, we didn't clean up this metadata in MarkComplete, but after this commit we do. Also: - add more logs - add an extra test - emit an extra metric - fix typos

Build-index currently does not use directory sharding. It also does not store blobs, whose keys are blob digests (i.e. SHA256 strings). Therefore, to replace the existing SimpleStore used by build-index, we need to make sure sharding is configurable. This commit makes that change.

…d simplify naming - all DiskStore code is now under the `lib/store/disk` package. After I replace all storage code, I plan to have `lib/store/disk`, `lib/store/mem`, and `lib/store/tiered` for the 3 different cache implementations. - Since the package already has `disk` in its name, I made the naming more concise by renaming `DiskStore` to `Store` (the import will be `disk.Store`) - Added a `Config` struct to simplify the constructor interfaces and document the configurable parameters in a single place. - fix a bug where rebooting does not work when sharding is off. Add a test that catches the bug. - make `rebootKeys` take `complete bool` instead of `subDir string` to stay consistent with the other functions' signatures - make small improvements to tests and docs

Anton-Kalpakchiev requested review from iawakimjan and sambhav-jain-16 June 6, 2026 14:38

github-actions Bot added the size/m label Jun 6, 2026

Anton-Kalpakchiev changed the title ~~Refactor store package to be simpler and more performant~~ Refactor store package for performance and simplicity Jun 6, 2026

Anton-Kalpakchiev requested a review from thijmv June 6, 2026 14:45

github-actions Bot added the size/xl label Jun 6, 2026

Anton-Kalpakchiev force-pushed the lru-cache branch 3 times, most recently from d3ffa2a to 17bb3fc Compare June 8, 2026 14:01

thijmv reviewed Jun 11, 2026

View reviewed changes

Anton-Kalpakchiev force-pushed the lru-cache branch 5 times, most recently from 01cd3d2 to 2683553 Compare June 12, 2026 15:37

thijmv reviewed Jun 15, 2026

View reviewed changes

sambhav-jain-16 reviewed Jun 19, 2026

View reviewed changes

Comment thread lib/store/disk_store.go Outdated

Comment thread lib/store/disk/store.go

Comment thread lib/store/disk/store.go

Comment thread lib/store/disk_store.go Outdated

Anton-Kalpakchiev mentioned this pull request Jul 25, 2026

feat(store): Add memory LRU cache #647

Open

sambhav-jain-16 reviewed Jul 27, 2026

View reviewed changes

sambhav-jain-16 approved these changes Jul 27, 2026

View reviewed changes

thijmv reviewed Jul 27, 2026

View reviewed changes

thijmv approved these changes Jul 28, 2026

View reviewed changes

Comment thread lib/store/disk/pather.go

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

Design new LRU disk store and implement some APIs

24e18a8

This commit designs the client-facing interface of the disk store and implements its core APIs. It establishes the core data structures used. It also adds tests.

feat(DiskStore): implement storing blob metadata

c7bd730

feat(DiskStore): reboot state from disk after crash

cc3187d

Anton-Kalpakchiev added 3 commits July 28, 2026 17:12

Anton-Kalpakchiev force-pushed the lru-cache branch from 1bcdd96 to 85e4790 Compare July 28, 2026 15:22

Anton-Kalpakchiev removed the size/m label Jul 28, 2026

Anton-Kalpakchiev mentioned this pull request Jul 29, 2026

feat(store): Add tiered cache (disk + mem) #650

Open

Uh oh!

Conversation

Anton-Kalpakchiev commented Jun 6, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

TLDR

Justification

Why transition from a TTI to an LRU eviction policy (copied from this ERD)

Why refactor the code

Implementation (copied from the ERD)

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

sambhav-jain-16 left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sambhav-jain-16 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

Anton-Kalpakchiev Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

sambhav-jain-16 Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

Anton-Kalpakchiev Jul 27, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Anton-Kalpakchiev commented Jul 27, 2026

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

thijmv left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

Anton-Kalpakchiev commented Jun 6, 2026 •

edited

Loading