From 24e18a86d27a21fb473ff4cc4d978f216be6a802 Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Fri, 5 Jun 2026 21:23:29 +0200 Subject: [PATCH 01/10] Design new LRU disk store and implement some APIs 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. --- lib/store/disk_store.go | 348 ++++++++++++++++ lib/store/disk_store_test.go | 372 ++++++++++++++++++ lib/store/file.go | 38 +- .../storage/originstorage/torrent_archive.go | 2 +- 4 files changed, 758 insertions(+), 2 deletions(-) create mode 100644 lib/store/disk_store.go create mode 100644 lib/store/disk_store_test.go diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go new file mode 100644 index 000000000..e4ee35ad9 --- /dev/null +++ b/lib/store/disk_store.go @@ -0,0 +1,348 @@ +package store + +import ( + "container/list" + "errors" + "fmt" + "os" + "path/filepath" + "sync" + + "github.com/uber/kraken/lib/store/base" + "github.com/uber/kraken/lib/store/metadata" +) + +const ( + _defaultFilePerm = 0775 + // _defaultShardIDLength is the number of bytes of file digest to be used for shard ID. + // For every byte (2 HEX char), one more level of directories will be created. + _defaultShardIDLength = 2 + _incompleteSubDir = "incomplete" + _completeSubDir = "complete" + _blobFileName = "data" +) + +// DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. +// +// - Supports pagination of blobs during reading/writing, such that blobs don't need to be fully loaded into memory. +// +// - New blobs are considered 'incomplete', which unlists them from LRU eviction. Read APIs may filter out incomplete blobs. +// +// - All APIs are thread-safe. Parallel access to a single file is allowed but clients must ensure they don't intervene with one another. +// +// - Supports (un-)marking blobs as non-evictable (may be needed when that data must be written back to remote storage). +// +// - Crash-resistant - all state is restored upon restart (LRU order is approximated through file `ctime`). +// +// - Uses directory sharding to speed up disk performance. +type DiskStore struct { + capacity uint64 + size uint64 // includes both used and reserved space + dir string + // synchronizes mem state access and syscalls to the fs in the APIs (opening, moving files, etc.) + mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. + // complete, evictable blobs. + blobs map[string]*list.Element // value of [list.Element] is [el]. + // Back is most recently used, front is the next to evict. + evictQueue *list.List + incompleteBlobs map[string]uint64 + // complete blobs that cannot be evicted. + unevictableBlobs map[string]uint64 +} + +type el struct { + key string + size uint64 +} + +// NewDiskStore creates a [DiskStore]. +func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { + // TODO - create a Config struct. + // TODO - recover persisted state in case of crash. + return &DiskStore{ + dir: rootDir, + capacity: capacityBytes, + blobs: make(map[string]*list.Element), + evictQueue: list.New(), + incompleteBlobs: make(map[string]uint64), + unevictableBlobs: make(map[string]uint64), + size: 0, + }, nil +} + +// Open returns an FD to a file in the store. [os.ErrNotExists] is returned on missing entry. +// The blob cannot be evicted before the client calls Close() on it. +func (s *DiskStore) Open(key string, ignoreIncomplete bool) (FileReadWriter, error) { + s.mu.Lock() + defer s.mu.Unlock() + + unevictableBlobSize, okUnevictable := s.unevictableBlobs[key] + if okUnevictable { + complete := true + return s.open(key, unevictableBlobSize, complete) + } + + incompleteBlobSize, okIncomplete := s.incompleteBlobs[key] + if okIncomplete && !ignoreIncomplete { + complete := false + return s.open(key, incompleteBlobSize, complete) + } + + node, okBlob := s.blobs[key] + if !okBlob { + return nil, os.ErrNotExist + } + + size := node.Value.(el).size + s.evictQueue.MoveToBack(node) + complete := true + return s.open(key, size, complete) +} + +func (s *DiskStore) open(key string, size uint64, complete bool) (FileReadWriter, error) { + path := s.blobPath(key, complete) + f, err := os.OpenFile(path, os.O_RDWR, _defaultFilePerm) + if err != nil { + return nil, fmt.Errorf("open: %w", err) + } + return newReadWriter(f, size), nil +} + +// Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExists] if the blob is not found. +func (s *DiskStore) Stat(key string, ignoreIncomplete bool) (os.FileInfo, error) { + // We **could** avoid locking the mutex by just statting the file directly. However, the current implementation + // prefers mutex contention over extra disk usage, as origin is bottlenecked by disk IO. + s.mu.Lock() + defer s.mu.Unlock() + + _, okBlob := s.blobs[key] + _, okUnevictable := s.unevictableBlobs[key] + _, okIncomplete := s.incompleteBlobs[key] + + if !okBlob && !okUnevictable && (ignoreIncomplete || !okIncomplete) { + return nil, os.ErrNotExist + } + complete := okBlob || okUnevictable + blobPath := s.blobPath(key, complete) + return os.Stat(blobPath) +} + +// Create adds a new, incomplete blob to the store. Incomplete entries cannot be automatically +// evicted. MarkComplete must be called once the blob is complete. +func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) { + // TODO - we might want some TTI on uploads to the store, after which we cancel the upload, e.g. 1min without the client uploading more data. + s.mu.Lock() + defer s.mu.Unlock() + + if _, ok := s.blobs[key]; ok { + // TODO - consider whether we need public errors for these cases. + return nil, errors.New("blob is already in store") + } + if _, ok := s.unevictableBlobs[key]; ok { + return nil, errors.New("blob is already in store") + } + if _, ok := s.incompleteBlobs[key]; ok { + return nil, errors.New("blob is already in store (it is incomplete)") + } + + if err := s.reserveSpace(sizeBytes); err != nil { + return nil, fmt.Errorf("reserve space: %w", err) + } + + complete := false + dirName := s.dirPath(key, complete) + err := os.MkdirAll(dirName, _defaultFilePerm) + if err != nil { + s.releaseSpace(sizeBytes) + return nil, fmt.Errorf("ensure dir: %w", err) + } + blobPath := s.blobPath(key, complete) + flag := os.O_RDWR | os.O_CREATE | os.O_EXCL + f, err := os.OpenFile(blobPath, flag, _defaultFilePerm) + if err != nil { + s.releaseSpace(sizeBytes) + return nil, fmt.Errorf("open file: %w", err) + } + + s.incompleteBlobs[key] = sizeBytes + + return newReadWriter(f, sizeBytes), nil +} + +func (s *DiskStore) reserveSpace(space uint64) error { + // TODO - emit latency to reserve space for a blob. + for s.size+space > s.capacity { + if s.evictQueue.Len() == 0 { + return errors.New("cannot evict enough, the unevictable/incomplete blobs are using up all the space") + } + + toEvictNode := s.evictQueue.Front() + toEvictEl := toEvictNode.Value.(el) + + complete := true + err := s.deleteFromDisk(toEvictEl.key, complete) + if err != nil { + // TODO - consider whether we want to fail-open by doing `continue` here. + return fmt.Errorf("delete from disk: %w", err) + } + s.evictQueue.Remove(toEvictNode) + delete(s.blobs, toEvictEl.key) + s.releaseSpace(toEvictEl.size) + } + + s.size += space + return nil +} + +func (s *DiskStore) releaseSpace(space uint64) { + // TODO - if space > s.size, emit an error log for an invariant violation + s.size -= space +} + +// fully deletes the state of a blob, including metadata. Works both on persisted and non-persisted blobs. +func (s *DiskStore) deleteFromDisk(key string, complete bool) error { + dir := s.dirPath(key, complete) + return os.RemoveAll(dir) +} + +// MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless ForbidEviction has been called). +// Additionally, read APIs may optionally filter out incomplete blobs. +func (s *DiskStore) MarkComplete(key string) error { + s.mu.Lock() + defer s.mu.Unlock() + + size, ok := s.incompleteBlobs[key] + if !ok { + return fmt.Errorf("blob is not in incomplete state") + } + + oldPathDir := s.dirPath(key, false) + newPathDir := s.dirPath(key, true) + err := os.MkdirAll(filepath.Dir(newPathDir), _defaultFilePerm) + if err != nil { + return fmt.Errorf("mkdirall: %w", err) + } + err = os.Rename(oldPathDir, newPathDir) // atomic + if err != nil { + return fmt.Errorf("move dir: %w", err) + } + + delete(s.incompleteBlobs, key) + // TODO - handle unevictable blobs correctly (they must transition to s.unevictableBlobs, not s.blobs and s.evictQueue) + node := s.evictQueue.PushBack(el{key: key, size: size}) + s.blobs[key] = node + return nil +} + +// Delete removes a blob and its [metadata.Metadata] from the store. Works on incomplete blobs. +// Does NOT work on unevictable blobs and returns [base.ErrFilePersisted]. +func (s *DiskStore) Delete(key string) error { + s.mu.Lock() + defer s.mu.Unlock() + + if _, ok := s.unevictableBlobs[key]; ok { + return base.ErrFilePersisted + } + + size, ok := s.incompleteBlobs[key] + if ok { + return s.cancelWrite(key, size) + } + + node, ok := s.blobs[key] + if !ok { + return os.ErrNotExist + } + + complete := true + err := s.deleteFromDisk(key, complete) + if err != nil { + return fmt.Errorf("delete from disk: %w", err) + } + delete(s.blobs, key) + s.evictQueue.Remove(node) + s.releaseSpace(node.Value.(el).size) + return nil +} + +func (s *DiskStore) cancelWrite(key string, size uint64) error { + complete := false + err := s.deleteFromDisk(key, complete) + if err != nil { + return fmt.Errorf("delete from disk: %w", err) + } + delete(s.incompleteBlobs, key) + s.releaseSpace(size) + return nil +} + +// List returns the blobs' keys. +func (s *DiskStore) List(ignoreIncomplete bool) []string { + s.mu.RLock() + defer s.mu.RUnlock() + + l := len(s.blobs) + len(s.unevictableBlobs) + if !ignoreIncomplete { + l += len(s.incompleteBlobs) + } + res := make([]string, l) + i := 0 + for key := range s.blobs { + res[i] = key + i++ + } + for key := range s.unevictableBlobs { + res[i] = key + i++ + } + if !ignoreIncomplete { + for key := range s.incompleteBlobs { + res[i] = key + i++ + } + } + return res +} + +// ForbidEviction unlists a blob from LRU eviction. +// Needed when e.g. blobs must be written back to GCS/S3 and eviction before that is unacceptable. +func (s *DiskStore) ForbidEviction(key string) error { return errors.New("not implemented") } + +// AllowEviction removes the effect of ForbidEviction for a blob. +func (s *DiskStore) AllowEviction(key string) error { return errors.New("not implemented") } + +// WriteMetadata atomically stores `md` on disk. Can be called on both complete and incomplete blobs. +func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { + return errors.New("not implemented") +} + +// GetMetadata populates `md` if the metadata is present. Returns [os.ErrNotExists] if key is not in store. +func (s *DiskStore) GetMetadata(key string, md metadata.Metadata, ignoreIncomplete bool) error { + return errors.New("not implemented") +} + +// DeleteMetadata removes the respective metadata, if present. +func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { + return errors.New("not implemented") +} + +func (s *DiskStore) blobPath(key string, complete bool) string { + dirName := s.dirPath(key, complete) + return filepath.Join(dirName, _blobFileName) +} + +func (s *DiskStore) dirPath(key string, complete bool) string { + // TODO - allow config to specify whether to shard or not. Allow no sharding, so we can replace [SimpleStore]. + subDirName := _incompleteSubDir + if complete { + subDirName = _completeSubDir + } + dirPath := filepath.Join(s.dir, subDirName) + for i := 0; i < int(_defaultShardIDLength) && i < len(key)/2; i++ { + // (1 byte = 2 char of file name assumming file name is in HEX) + dirName := key[i*2 : i*2+2] + dirPath = filepath.Join(dirPath, dirName) + } + + return filepath.Join(dirPath, key) +} diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go new file mode 100644 index 000000000..896619b3d --- /dev/null +++ b/lib/store/disk_store_test.go @@ -0,0 +1,372 @@ +package store + +import ( + "bytes" + "io" + "os" + "sync" + "testing" + "testing/iotest" + "time" + + "github.com/stretchr/testify/require" + "github.com/uber/kraken/core" + "github.com/uber/kraken/utils/memsize" +) + +func testStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string) { + rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") + require.NoError(t, err) + t.Cleanup(func() { os.RemoveAll(rootDir) }) + + store, err := NewDiskStore(capacity, rootDir) + require.NoError(t, err) + return store, rootDir +} + +func TestOpen(t *testing.T) { + t.Skip("TODO") + // Assert that reading a file before it's committed doesn't work. +} + +func TestDiskStore(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + + digests := []core.Digest{} + for i := range 10 { + digest := core.DigestFixture() + digests = append(digests, digest) + writer, err := store.Create(digest.Hex(), memsize.KB) + require.NoError(err) + + data := make([]byte, memsize.KB) + for k := range data { + data[k] = byte(i + 1) + } + n, err := io.Copy(writer, bytes.NewReader(make([]byte, memsize.KB))) + require.Equal(n, int64(memsize.KB)) + require.NoError(err) + } + require.Equal(store.size, 10*memsize.KB) + + digest := core.DigestFixture() + writer, err := store.Create(digest.Hex(), memsize.B) + require.EqualError(err, "reserve space: cannot evict enough, the unevictable/incomplete blobs are using up all the space") + require.Nil(writer) + + reader, err := store.Open(digests[0].Hex(), true) + require.Equal(err, os.ErrNotExist) + require.Nil(reader) + + require.NoError(store.MarkComplete(digests[0].Hex())) + reader, err = store.Open(digests[0].Hex(), true) + require.NoError(err) + wantData := make([]byte, memsize.KB) + for k := range wantData { + wantData[k] = byte(1) + } + require.NoError(iotest.TestReader(reader, make([]byte, memsize.KB))) + + // now test that LRU logic works - make sure that the 1 that is committed gets evicted. + digest = core.DigestFixture() + writer, err = store.Create(digest.Hex(), memsize.KB) + require.NoError(err) + reader, err = store.Open(digests[0].Hex(), true) + require.Equal(err, os.ErrNotExist) + require.Nil(reader) +} + +func TestEviction(t *testing.T) { + // order - when 1) complete evictable, 2) incomplete evictable, 3) incomplete unevictable, and 4) complete unevictable blobs are in store. + // order - make sure that marking a blob as complete considers it as most recently used. same for `EnableEviction`. + // correctness - does not evict unevictable/incomplete blobs no matter what. + // new test - eviction works under parallel access + // assert actual size of disk directory changes as expected throuhgout the whole process + t.Skip("TODO") +} + +func TestCrashRecovery(t *testing.T) { + t.Skip("TODO - implement") +} + +func TestParallelAccessToSingleFile(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + + key := core.DigestFixture().Hex() + f, err := store.Create(key, 1*memsize.KB) + require.NoError(err) + require.NoError(f.Close()) + // Spawn 5 routines in parallel that write and read to different parts of the file. + var wg sync.WaitGroup + wg.Add(5) + + for idx := range 5 { + go func(idx int64) { + defer wg.Done() + + ignoreIncomplete := false + f, err := store.Open(key, ignoreIncomplete) + require.NoError(err) + pos := idx * 10 + writtenData := make([]byte, 10) + for k := range writtenData { + writtenData[k] = byte(idx) + } + n, err := f.WriteAt(writtenData, pos) + require.NoError(err) + require.Equal(10, n) + + defer func() { require.NoError(f.Close()) }() + readData := make([]byte, 10) + n, err = f.ReadAt(readData, pos) + require.NoError(err) + require.Equal(10, n) + require.Equal(writtenData, readData) + }(int64(idx)) + } + + wg.Wait() + require.NoError(store.MarkComplete(key)) + + ignoreIncomplete := true + f, err = store.Open(key, ignoreIncomplete) + require.NoError(err) + defer func() { require.NoError(f.Close()) }() + + wantFileData := make([]byte, 50) + for i := range 50 { + wantFileData[i] = byte(i / 10) + } + fData, err := io.ReadAll(f) + require.NoError(err) + require.Equal(wantFileData, fData) +} + +func TestOpenedFileAccessibleAfterMarkedComplete(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + + key := core.DigestFixture().Hex() + f, err := store.Create(key, 1*memsize.KB) + require.NoError(err) + _, err = io.Copy(f, bytes.NewReader([]byte("Hello World"))) + require.NoError(err) + require.NoError(f.Close()) + + ignoreIncomplete := false + incompleteFile, err := store.Open(key, ignoreIncomplete) + require.NoError(err) + defer func() { require.NoError(incompleteFile.Close()) }() + + require.NoError(store.MarkComplete(key)) + + ignoreIncomplete = true + completeFile, err := store.Open(key, ignoreIncomplete) + require.NoError(err) + defer func() { require.NoError(completeFile.Close()) }() + + incompleteFileData, err := io.ReadAll(incompleteFile) + require.NoError(err) + completeFileData, err := io.ReadAll(completeFile) + require.NoError(err) + + require.Equal([]byte("Hello World"), incompleteFileData) + require.Equal([]byte("Hello World"), completeFileData) +} + +func TestOpenedFileAccessibleAfterEviction(t *testing.T) { + // when a file is created, then received with a Get, then deleted, check that the store doesn't count its size but it is still on disk and accessible. then check that when the client closes the file, it is not accessible anymore and not on disk anymore. + // test after both "Delete" and natural LRU eviction. + // what happens when we try to close the fd? i assume it works? + t.Skip("TODO") +} + +func TestDelete(t *testing.T) { + t.Run("uncommitted blob", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + w, err := store.Create(key, 100*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + + require.NoError(store.Delete(key)) + + require.Empty(store.List(true)) + require.Equal(uint64(0), store.size) + }) + t.Run("uncommitted blob with forbidden eviction", func(t *testing.T) { + t.Skip("TODO - test this once ForbidEviction is implemented") + }) + t.Run("committed, evictable blob", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + w, err := store.Create(key, 100*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + require.NoError(store.MarkComplete(key)) + + require.NoError(store.Delete(key)) + + require.Empty(store.List(true)) + require.Equal(uint64(0), store.size) + _, err = store.Open(key, true) + require.Equal(os.ErrNotExist, err) + }) + t.Run("committed, unevictable blob", func(t *testing.T) { + t.Skip("TODO - test this once unevictable blobs are implemented") + }) + t.Run("not found", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + + err := store.Delete(key) + require.Equal(os.ErrNotExist, err) + }) +} + +func TestMarkComplete(t *testing.T) { + t.Run("uncommitted blob", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + w, err := store.Create(key, 100*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + + require.NoError(store.MarkComplete(key)) + + require.Equal([]string{key}, store.List(true)) + require.Equal(uint64(100), store.size) + _, err = store.Open(key, true) + require.NoError(err) + }) + t.Run("uncommitted blob with forbidden eviction", func(t *testing.T) { + t.Skip("TODO - test this once ForbidEviction is implemented") + }) + t.Run("already committed blob", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + w, err := store.Create(key, 100*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + require.NoError(store.MarkComplete(key)) + + err = store.MarkComplete(key) + require.EqualError(err, "blob is not in incomplete state") + }) + t.Run("already committed blob with forbidden eviction", func(t *testing.T) { + t.Skip("TODO - test this once ForbidEviction is implemented") + }) + t.Run("not found", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + + err := store.MarkComplete(key) + require.EqualError(err, "blob is not in incomplete state") + }) +} + +func TestStat(t *testing.T) { + t.Run("committed, evictable blob", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + w, err := store.Create(key, 10*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + require.NoError(store.MarkComplete(key)) + + fInfo, err := store.Stat(key, true) + require.NoError(err) + + require.False(fInfo.IsDir()) + require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) + require.Equal(_blobFileName, fInfo.Name()) + require.Equal(int64(10), fInfo.Size()) + }) + t.Run("committed, unevictable blob", func(t *testing.T) { + t.Skip("TODO - test this once unevictable blobs are implemented") + }) + t.Run("uncommitted blob", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + w, err := store.Create(key, 10*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + + fInfo, err := store.Stat(key, true) + require.Equal(os.ErrNotExist, err) + require.Nil(fInfo) + }) + t.Run("non-existent blob", func(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + + _, err := store.Stat(key, true) + require.Equal(os.ErrNotExist, err) + }) +} + +func TestList(t *testing.T) { + require := require.New(t) + store, _ := testStore(t, 10*memsize.KB) + + require.Empty(store.List(true)) + + uncommittedBlobKey := core.DigestFixture().Hex() + w, err := store.Create(uncommittedBlobKey, 10*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + committedBlobKey := core.DigestFixture().Hex() + w, err = store.Create(committedBlobKey, 10*memsize.B) + require.NoError(err) + _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + require.NoError(store.MarkComplete(committedBlobKey)) + // TODO - add an unevictable blob here once that's implemented + + require.Equal([]string{committedBlobKey}, store.List(true)) +} + +func TestPathing(t *testing.T) { + require := require.New(t) + store, rootDir := testStore(t, 10*memsize.KB) + + key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + committed := false + dirPath := store.dirPath(key, committed) + wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + require.Equal(wantDirPath, dirPath) + blobPath := store.blobPath(key, committed) + wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" + require.Equal(wantBlobPath, blobPath) + + committed = true + dirPath = store.dirPath(key, committed) + wantDirPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + require.Equal(wantDirPath, dirPath) + blobPath = store.blobPath(key, committed) + wantBlobPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" + require.Equal(wantBlobPath, blobPath) +} + +func TestForbidEviction(t *testing.T) { + t.Skip("TODO") + // unevictable blobs can neither be evicted nor deleted. remarking as evictable works but resets access time. +} diff --git a/lib/store/file.go b/lib/store/file.go index 15bb43733..852dae1d4 100644 --- a/lib/store/file.go +++ b/lib/store/file.go @@ -13,10 +13,46 @@ // limitations under the License. package store -import "github.com/uber/kraken/lib/store/base" +import ( + "os" + + "github.com/uber/kraken/lib/store/base" +) // FileReadWriter is a readable, writable file. type FileReadWriter = base.FileReadWriter // FileReader is a read-only file. type FileReader = base.FileReader + +func newReadWriter(f *os.File, size uint64) FileReadWriter { + return &rwImpl{ + File: f, + size: int64(size), + } +} + +var _ FileReadWriter = &rwImpl{} + +type rwImpl struct { + *os.File + size int64 +} + +// Size returns the full size of the blob in bytes, even if the blob is not fully written yet. +func (i *rwImpl) Size() int64 { + return i.size +} + +// Cancel is supposed to remove any written content. +// In this implementation file is not actually removed, and it's fine since there won't be name +// collision between upload files. +func (i *rwImpl) Cancel() error { + return i.Close() +} + +// Commit is supposed to flush all content for buffered writer. +// In this implementation all writes write to the file directly through syscall. +func (i *rwImpl) Commit() error { + return i.Close() +} diff --git a/lib/torrent/storage/originstorage/torrent_archive.go b/lib/torrent/storage/originstorage/torrent_archive.go index d190ef3d1..bfcb8d8e3 100644 --- a/lib/torrent/storage/originstorage/torrent_archive.go +++ b/lib/torrent/storage/originstorage/torrent_archive.go @@ -74,7 +74,7 @@ func (a *TorrentArchive) CreateTorrent(namespace string, d core.Digest) (storage } // GetTorrent returns a Torrent for an existing file on disk. If the file does -// not exist, attempts to re-fetch the file from the storae backend configured +// not exist, attempts to re-fetch the file from the storage backend configured // for namespace in a background goroutine, and returns os.ErrNotExist. func (a *TorrentArchive) GetTorrent(namespace string, d core.Digest) (storage.Torrent, error) { mi, err := a.getMetaInfo(namespace, d) From 1f5e5330a5f27b037fab5b109d2c38197cfbe000 Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Tue, 9 Jun 2026 22:59:17 +0200 Subject: [PATCH 02/10] feat(DiskStore): Implement unevictable blobs - 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. --- lib/store/disk_store.go | 217 +++++++++++++----- lib/store/disk_store_test.go | 422 ++++++++++++++++++++++++++++------- 2 files changed, 500 insertions(+), 139 deletions(-) diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index e4ee35ad9..c4e4ceb85 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -8,18 +8,19 @@ import ( "path/filepath" "sync" - "github.com/uber/kraken/lib/store/base" "github.com/uber/kraken/lib/store/metadata" + "github.com/uber/kraken/utils/closers" ) const ( - _defaultFilePerm = 0775 // _defaultShardIDLength is the number of bytes of file digest to be used for shard ID. // For every byte (2 HEX char), one more level of directories will be created. - _defaultShardIDLength = 2 - _incompleteSubDir = "incomplete" - _completeSubDir = "complete" - _blobFileName = "data" + _defaultFilePerm = 0775 + _defaultShardIDLength = 2 + _incompleteSubDir = "incomplete" + _completeSubDir = "complete" + _blobFileName = "data" + _evictionBannedFileName = "_eviction_banned" ) // DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. @@ -44,7 +45,8 @@ type DiskStore struct { // complete, evictable blobs. blobs map[string]*list.Element // value of [list.Element] is [el]. // Back is most recently used, front is the next to evict. - evictQueue *list.List + evictQueue *list.List + // incomplete blobs that may or may not be evictable. incompleteBlobs map[string]uint64 // complete blobs that cannot be evicted. unevictableBlobs map[string]uint64 @@ -58,7 +60,7 @@ type el struct { // NewDiskStore creates a [DiskStore]. func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { // TODO - create a Config struct. - // TODO - recover persisted state in case of crash. + // TODO - recover persisted state in case of crash. might need to evict. emit log if so. return &DiskStore{ dir: rootDir, capacity: capacityBytes, @@ -127,8 +129,9 @@ func (s *DiskStore) Stat(key string, ignoreIncomplete bool) (os.FileInfo, error) return os.Stat(blobPath) } -// Create adds a new, incomplete blob to the store. Incomplete entries cannot be automatically -// evicted. MarkComplete must be called once the blob is complete. +// Create adds a new, incomplete blob to the store and reserves space for it. +// Incomplete entries cannot be automatically evicted. MarkComplete must be called once the blob is complete. +// DiskStore does not ever check/use the real size of the blob and only uses `sizeBytes` for its eviction logic. func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) { // TODO - we might want some TTI on uploads to the store, after which we cancel the upload, e.g. 1min without the client uploading more data. s.mu.Lock() @@ -199,26 +202,38 @@ func (s *DiskStore) releaseSpace(space uint64) { s.size -= space } -// fully deletes the state of a blob, including metadata. Works both on persisted and non-persisted blobs. +// fully deletes the disk state of a blob, including metadata. Works on any blob. func (s *DiskStore) deleteFromDisk(key string, complete bool) error { dir := s.dirPath(key, complete) return os.RemoveAll(dir) } -// MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless ForbidEviction has been called). +// MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless BanEviction has been called). // Additionally, read APIs may optionally filter out incomplete blobs. func (s *DiskStore) MarkComplete(key string) error { s.mu.Lock() defer s.mu.Unlock() - size, ok := s.incompleteBlobs[key] - if !ok { - return fmt.Errorf("blob is not in incomplete state") + size, okIncomplete := s.incompleteBlobs[key] + _, okBlob := s.blobs[key] + _, okUnevictable := s.unevictableBlobs[key] + if !okIncomplete && !okBlob && !okUnevictable { + return os.ErrNotExist + } + + if !okIncomplete { + // no-op + return nil } oldPathDir := s.dirPath(key, false) newPathDir := s.dirPath(key, true) - err := os.MkdirAll(filepath.Dir(newPathDir), _defaultFilePerm) + complete := false + isUnevictable, err := s.checkDiskIfUnevictable(key, complete) + if err != nil { + return fmt.Errorf("check if blob is evictable or not: %w", err) + } + err = os.MkdirAll(filepath.Dir(newPathDir), _defaultFilePerm) if err != nil { return fmt.Errorf("mkdirall: %w", err) } @@ -226,54 +241,68 @@ func (s *DiskStore) MarkComplete(key string) error { if err != nil { return fmt.Errorf("move dir: %w", err) } - delete(s.incompleteBlobs, key) - // TODO - handle unevictable blobs correctly (they must transition to s.unevictableBlobs, not s.blobs and s.evictQueue) + if isUnevictable { + s.unevictableBlobs[key] = size + return nil + } + node := s.evictQueue.PushBack(el{key: key, size: size}) s.blobs[key] = node return nil } -// Delete removes a blob and its [metadata.Metadata] from the store. Works on incomplete blobs. -// Does NOT work on unevictable blobs and returns [base.ErrFilePersisted]. +func (s *DiskStore) checkDiskIfUnevictable(key string, complete bool) (bool, error) { + dirPath := s.dirPath(key, complete) + flagBlobPath := filepath.Join(dirPath, _evictionBannedFileName) + _, err := os.Stat(flagBlobPath) + if err == nil { + return true, nil + } + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, fmt.Errorf("stat: %w", err) +} + +// Delete removes a blob and its [metadata.Metadata] from the store. Works on any blob. func (s *DiskStore) Delete(key string) error { s.mu.Lock() defer s.mu.Unlock() - if _, ok := s.unevictableBlobs[key]; ok { - return base.ErrFilePersisted - } - - size, ok := s.incompleteBlobs[key] - if ok { - return s.cancelWrite(key, size) + if size, ok := s.unevictableBlobs[key]; ok { + complete := true + err := s.deleteFromDisk(key, complete) + if err != nil { + return fmt.Errorf("delete from disk: %w", err) + } + delete(s.unevictableBlobs, key) + s.releaseSpace(size) + return nil } - - node, ok := s.blobs[key] - if !ok { - return os.ErrNotExist + if size, ok := s.incompleteBlobs[key]; ok { + complete := false + err := s.deleteFromDisk(key, complete) + if err != nil { + return fmt.Errorf("delete from disk: %w", err) + } + delete(s.incompleteBlobs, key) + s.releaseSpace(size) + return nil } - - complete := true - err := s.deleteFromDisk(key, complete) - if err != nil { - return fmt.Errorf("delete from disk: %w", err) + if node, ok := s.blobs[key]; ok { + complete := true + err := s.deleteFromDisk(key, complete) + if err != nil { + return fmt.Errorf("delete from disk: %w", err) + } + delete(s.blobs, key) + s.evictQueue.Remove(node) + s.releaseSpace(node.Value.(el).size) + return nil } - delete(s.blobs, key) - s.evictQueue.Remove(node) - s.releaseSpace(node.Value.(el).size) - return nil -} -func (s *DiskStore) cancelWrite(key string, size uint64) error { - complete := false - err := s.deleteFromDisk(key, complete) - if err != nil { - return fmt.Errorf("delete from disk: %w", err) - } - delete(s.incompleteBlobs, key) - s.releaseSpace(size) - return nil + return os.ErrNotExist } // List returns the blobs' keys. @@ -304,12 +333,92 @@ func (s *DiskStore) List(ignoreIncomplete bool) []string { return res } -// ForbidEviction unlists a blob from LRU eviction. +// BanEviction marks a blob as unevictable by LRU eviction. It is idempotent. // Needed when e.g. blobs must be written back to GCS/S3 and eviction before that is unacceptable. -func (s *DiskStore) ForbidEviction(key string) error { return errors.New("not implemented") } +func (s *DiskStore) BanEviction(key string) error { + s.mu.Lock() + defer s.mu.Unlock() + + node, okBlob := s.blobs[key] + _, okUnevictable := s.unevictableBlobs[key] + _, okIncomplete := s.incompleteBlobs[key] + + if !okBlob && !okUnevictable && !okIncomplete { + return os.ErrNotExist + } + if okUnevictable { + // no-op + return nil + } -// AllowEviction removes the effect of ForbidEviction for a blob. -func (s *DiskStore) AllowEviction(key string) error { return errors.New("not implemented") } + complete := okBlob + dirPath := s.dirPath(key, complete) + flagBlobPath := filepath.Join(dirPath, _evictionBannedFileName) + // We persist the ban as a flag file on disk, such that after + // a system crash, we can recover the ban. + f, err := os.OpenFile(flagBlobPath, os.O_RDONLY|os.O_CREATE, _defaultFilePerm) + if err != nil { + return fmt.Errorf("create file that flags eviction as banned: %w", err) + } + closers.Close(f) + + if okIncomplete { + return nil + } + + nodeEl := node.Value.(el) + delete(s.blobs, key) + s.evictQueue.Remove(node) + s.unevictableBlobs[key] = nodeEl.size + return nil +} + +// UnbanDeletion removes the effect of BanDeletion for a blob. It is idempotent. +func (s *DiskStore) UnbanEviction(key string) error { + s.mu.Lock() + defer s.mu.Unlock() + + _, okBlob := s.blobs[key] + sizeUnevictable, okUnevictable := s.unevictableBlobs[key] + _, okIncomplete := s.incompleteBlobs[key] + + if !okBlob && !okUnevictable && !okIncomplete { + return os.ErrNotExist + } + + if okBlob { + // no-op + return nil + } + + complete := okUnevictable + isUnevictable, err := s.checkDiskIfUnevictable(key, complete) + if err != nil { + return fmt.Errorf("check on disk if file is evictable: %w", err) + } + if !isUnevictable { + // no-op + if okUnevictable { + // TODO - log invariant violation + } + return nil + } + // TODO - create a helper function that finds a file with a specific suffix and returns its path + dirPath := s.dirPath(key, complete) + flagBlobPath := filepath.Join(dirPath, _evictionBannedFileName) + err = os.Remove(flagBlobPath) + if err != nil { + return fmt.Errorf("remove file that flags eviction as banned: %w", err) + } + if okIncomplete { + return nil + } + + delete(s.unevictableBlobs, key) + node := s.evictQueue.PushBack(el{key: key, size: sizeUnevictable}) + s.blobs[key] = node + return nil +} // WriteMetadata atomically stores `md` on disk. Can be called on both complete and incomplete blobs. func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index 896619b3d..9f74661ed 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -3,7 +3,10 @@ package store import ( "bytes" "io" + "io/fs" "os" + "path/filepath" + "strings" "sync" "testing" "testing/iotest" @@ -14,7 +17,7 @@ import ( "github.com/uber/kraken/utils/memsize" ) -func testStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string) { +func newTestStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) t.Cleanup(func() { os.RemoveAll(rootDir) }) @@ -24,27 +27,46 @@ func testStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string) { return store, rootDir } -func TestOpen(t *testing.T) { - t.Skip("TODO") - // Assert that reading a file before it's committed doesn't work. +func newTestFile(t *testing.T, store *DiskStore, size uint64) (f FileReadWriter, key string) { + require := require.New(t) + key = core.DigestFixture().Hex() + var err error + f, err = store.Create(key, size) + require.NoError(err) + return f, key +} + +// does not count 1) the directories for sharding, 2) metadata files, and 3) the _eviction_banned flag file. +func numBlobsOnDisk(t *testing.T, store *DiskStore) int { + numBlobs := 0 + err := filepath.Walk(store.dir, func(path string, info fs.FileInfo, err error) error { + if !strings.HasSuffix(path, _blobFileName) { + return nil + } + + numBlobs++ + return nil + }) + require.NoError(t, err) + return numBlobs } func TestDiskStore(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) digests := []core.Digest{} for i := range 10 { digest := core.DigestFixture() digests = append(digests, digest) - writer, err := store.Create(digest.Hex(), memsize.KB) + f, err := store.Create(digest.Hex(), memsize.KB) require.NoError(err) data := make([]byte, memsize.KB) for k := range data { data[k] = byte(i + 1) } - n, err := io.Copy(writer, bytes.NewReader(make([]byte, memsize.KB))) + n, err := io.Copy(f, bytes.NewReader(make([]byte, memsize.KB))) require.Equal(n, int64(memsize.KB)) require.NoError(err) } @@ -68,7 +90,7 @@ func TestDiskStore(t *testing.T) { } require.NoError(iotest.TestReader(reader, make([]byte, memsize.KB))) - // now test that LRU logic works - make sure that the 1 that is committed gets evicted. + // now test that LRU logic works - make sure that the 1 that is complete gets evicted. digest = core.DigestFixture() writer, err = store.Create(digest.Hex(), memsize.KB) require.NoError(err) @@ -78,21 +100,128 @@ func TestDiskStore(t *testing.T) { } func TestEviction(t *testing.T) { - // order - when 1) complete evictable, 2) incomplete evictable, 3) incomplete unevictable, and 4) complete unevictable blobs are in store. - // order - make sure that marking a blob as complete considers it as most recently used. same for `EnableEviction`. - // correctness - does not evict unevictable/incomplete blobs no matter what. - // new test - eviction works under parallel access - // assert actual size of disk directory changes as expected throuhgout the whole process - t.Skip("TODO") + const _dontIgnoreIncompleteFiles = false + + require := require.New(t) + store, _ := newTestStore(t, 25*memsize.KB) + // create 5 blobs - a, b, c, d, e with different sizes. + a, aKey := newTestFile(t, store, 10*memsize.KB) + require.NoError(a.Close()) + b, bKey := newTestFile(t, store, 5*memsize.KB) + require.NoError(b.Close()) + c, cKey := newTestFile(t, store, 5*memsize.KB) + require.NoError(c.Close()) + d, dKey := newTestFile(t, store, 3*memsize.KB) + require.NoError(d.Close()) + e, eKey := newTestFile(t, store, 1*memsize.KB) + require.NoError(e.Close()) + + require.Equal(24*memsize.KB, store.size) + require.Equal(5, numBlobsOnDisk(t, store)) + // incomplete files cannot be evicted and adding 2KB would result in overreservation. + _, err := store.Create(core.DigestFixture().Hex(), 2*memsize.KB) + require.EqualError(err, "reserve space: cannot evict enough, the unevictable/incomplete blobs are using up all the space") + // start marking as complete in this specific order - c, b, a (MarkComplete resets access time). + require.NoError(store.MarkComplete(cKey)) + require.NoError(store.MarkComplete(bKey)) + require.NoError(store.MarkComplete(aKey)) + // d is complete but its eviction is banned. + require.NoError(store.MarkComplete(dKey)) + require.NoError(store.BanEviction(dKey)) + // e is banned from eviction before it even becomes complete. + require.NoError(store.BanEviction(eKey)) + require.Equal(24*memsize.KB, store.size) + // Add f (4KB) which should evict c to make space, as d and e are unevictable and c was accessed last (the MarkComplete call). + f, fKey := newTestFile(t, store, 4*memsize.KB) + require.NoError(f.Close()) + require.NoError(store.MarkComplete(fKey)) + keys := store.List(_dontIgnoreIncompleteFiles) + require.NotContains(keys, cKey) + // new size == 23KB == 24KB - 5KB (c) + 4KB (f) + require.Equal(23*memsize.KB, store.size) + require.Equal(5, numBlobsOnDisk(t, store)) + + // Add g (1KB), which will not evict anything + g, gKey := newTestFile(t, store, 1*memsize.KB) + require.NoError(g.Close()) + require.NoError(store.MarkComplete(gKey)) + require.Equal(24*memsize.KB, store.size) + require.Equal(6, numBlobsOnDisk(t, store)) + + // Add h (15KB), which evicts b and a: + h, hKey := newTestFile(t, store, 15*memsize.KB) + require.NoError(h.Close()) + require.NoError(store.MarkComplete(hKey)) + // size == 24KB == 24KB + 15KB (h) - 5KB (b) - 10KB (a) + require.Equal(24*memsize.KB, store.size) + require.Equal(5, numBlobsOnDisk(t, store)) + keys = store.List(_dontIgnoreIncompleteFiles) + require.NotContains(keys, bKey) + require.NotContains(keys, aKey) + + // allow e to be evicted. + require.NoError(store.MarkComplete(eKey)) + require.NoError(store.UnbanEviction(eKey)) + // eviction order (left-most is next to evict): f(4KB), g(1KB), h(15KB), e(1KB); d(3KB) is unevictable + // we open g to change the order to f, h, e, g + g, err = store.Open(gKey, _dontIgnoreIncompleteFiles) + require.NoError(err) + require.NoError(g.Close()) + + i, iKey := newTestFile(t, store, 5*memsize.KB) + require.NoError(store.MarkComplete(iKey)) + require.NoError(i.Close()) + keys = store.List(_dontIgnoreIncompleteFiles) + require.NotContains(keys, fKey) + require.Equal(25*memsize.KB, store.size) + require.Equal(5, numBlobsOnDisk(t, store)) + // eviction order: h(15KB), e(1KB), g(1KB), i(5KB); d(3KB) is unevictable + + j, jKey := newTestFile(t, store, 14*memsize.KB) + require.NoError(j.Close()) + require.NoError(store.MarkComplete(jKey)) + keys = store.List(_dontIgnoreIncompleteFiles) + require.NotContains(keys, hKey) + require.Equal(24*memsize.KB, store.size) + require.Equal(5, numBlobsOnDisk(t, store)) + // eviction order: e(1KB), g(1KB), i(5KB), j(14KB); d(3KB) is unevictable + + k, kKey := newTestFile(t, store, 2*memsize.KB) + require.NoError(k.Close()) + require.NoError(store.MarkComplete(kKey)) + keys = store.List(_dontIgnoreIncompleteFiles) + require.NotContains(keys, eKey) + require.Equal(25*memsize.KB, store.size) + require.Equal(5, numBlobsOnDisk(t, store)) + // eviction order: g(1KB), i(5KB), j(14KB), k(2KB); d(3KB) is unevictable + + l, lKey := newTestFile(t, store, 1*memsize.KB) + require.NoError(store.MarkComplete(lKey)) + require.NoError(l.Close()) + keys = store.List(_dontIgnoreIncompleteFiles) + require.NotContains(keys, gKey) + require.Equal(25*memsize.KB, store.size) + require.Equal(5, numBlobsOnDisk(t, store)) + // eviction order: i(5KB), j(14KB), k(2KB), l(1KB); d(3KB) is unevictable + + require.NoError(store.Delete(iKey)) + require.NoError(store.Delete(jKey)) + require.NoError(store.Delete(lKey)) + // evictionOrder: k(2KB); d(3KB) + require.Equal(5*memsize.KB, store.size) + require.Equal(2, numBlobsOnDisk(t, store)) } func TestCrashRecovery(t *testing.T) { t.Skip("TODO - implement") + + // test each blob type is as expected in the store's internal state + // test LRU order is approximated correctly after recovery (ctime shd be different than access time) } func TestParallelAccessToSingleFile(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() f, err := store.Create(key, 1*memsize.KB) @@ -146,7 +275,7 @@ func TestParallelAccessToSingleFile(t *testing.T) { func TestOpenedFileAccessibleAfterMarkedComplete(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() f, err := store.Create(key, 1*memsize.KB) @@ -184,46 +313,72 @@ func TestOpenedFileAccessibleAfterEviction(t *testing.T) { } func TestDelete(t *testing.T) { - t.Run("uncommitted blob", func(t *testing.T) { + t.Run("incomplete blob", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() - w, err := store.Create(key, 100*memsize.B) + f, err := store.Create(key, 100*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) require.NoError(err) - + require.NoError(f.Close()) require.NoError(store.Delete(key)) - require.Empty(store.List(true)) + require.Empty(store.List(false)) require.Equal(uint64(0), store.size) }) - t.Run("uncommitted blob with forbidden eviction", func(t *testing.T) { - t.Skip("TODO - test this once ForbidEviction is implemented") + t.Run("incomplete, unevictable blob", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 100*memsize.B) + require.NoError(err) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + require.NoError(f.Close()) + require.NoError(store.BanEviction(key)) + + require.NoError(store.Delete(key)) + + require.Empty(store.List(false)) + require.Equal(uint64(0), store.size) }) - t.Run("committed, evictable blob", func(t *testing.T) { + t.Run("complete blob", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() - w, err := store.Create(key, 100*memsize.B) + f, err := store.Create(key, 100*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) require.NoError(err) + require.NoError(f.Close()) require.NoError(store.MarkComplete(key)) require.NoError(store.Delete(key)) - require.Empty(store.List(true)) + require.Empty(store.List(false)) require.Equal(uint64(0), store.size) - _, err = store.Open(key, true) - require.Equal(os.ErrNotExist, err) }) - t.Run("committed, unevictable blob", func(t *testing.T) { - t.Skip("TODO - test this once unevictable blobs are implemented") + t.Run("complete, unevictable blob", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 100*memsize.B) + require.NoError(err) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + require.NoError(f.Close()) + require.NoError(store.MarkComplete(key)) + require.NoError(store.BanEviction(key)) + + require.NoError(store.Delete(key)) + + require.Empty(store.List(false)) + require.Equal(uint64(0), store.size) }) t.Run("not found", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() err := store.Delete(key) @@ -232,13 +387,14 @@ func TestDelete(t *testing.T) { } func TestMarkComplete(t *testing.T) { - t.Run("uncommitted blob", func(t *testing.T) { + t.Run("incomplete blob", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() - w, err := store.Create(key, 100*memsize.B) + f, err := store.Create(key, 100*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) require.NoError(err) require.NoError(store.MarkComplete(key)) @@ -248,125 +404,221 @@ func TestMarkComplete(t *testing.T) { _, err = store.Open(key, true) require.NoError(err) }) - t.Run("uncommitted blob with forbidden eviction", func(t *testing.T) { - t.Skip("TODO - test this once ForbidEviction is implemented") + t.Run("incomplete blob with forbidden eviction", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 100*memsize.B) + require.NoError(err) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + require.NoError(store.BanEviction(key)) + + require.NoError(store.MarkComplete(key)) + + require.Equal([]string{key}, store.List(true)) + require.Equal(uint64(100), store.size) + _, err = store.Open(key, true) + require.NoError(err) }) - t.Run("already committed blob", func(t *testing.T) { + t.Run("already complete blob", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() - w, err := store.Create(key, 100*memsize.B) + f, err := store.Create(key, 100*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 100))) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) require.NoError(err) require.NoError(store.MarkComplete(key)) - err = store.MarkComplete(key) - require.EqualError(err, "blob is not in incomplete state") + require.NoError(store.MarkComplete(key)) }) - t.Run("already committed blob with forbidden eviction", func(t *testing.T) { - t.Skip("TODO - test this once ForbidEviction is implemented") + t.Run("already complete blob with forbidden eviction", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 100*memsize.B) + require.NoError(err) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 100))) + require.NoError(err) + require.NoError(store.MarkComplete(key)) + require.NoError(store.BanEviction(key)) + + require.NoError(store.MarkComplete(key)) }) t.Run("not found", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() err := store.MarkComplete(key) - require.EqualError(err, "blob is not in incomplete state") + require.Equal(os.ErrNotExist, err) }) } func TestStat(t *testing.T) { - t.Run("committed, evictable blob", func(t *testing.T) { + t.Run("complete blob", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() - w, err := store.Create(key, 10*memsize.B) + f, err := store.Create(key, 10*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) require.NoError(err) require.NoError(store.MarkComplete(key)) fInfo, err := store.Stat(key, true) require.NoError(err) + _, err = store.Stat(key, false) + require.NoError(err) require.False(fInfo.IsDir()) require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) require.Equal(_blobFileName, fInfo.Name()) require.Equal(int64(10), fInfo.Size()) }) - t.Run("committed, unevictable blob", func(t *testing.T) { - t.Skip("TODO - test this once unevictable blobs are implemented") - }) - t.Run("uncommitted blob", func(t *testing.T) { + t.Run("complete, unevictable blob", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() - w, err := store.Create(key, 10*memsize.B) + f, err := store.Create(key, 10*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) require.NoError(err) + require.NoError(store.MarkComplete(key)) + require.NoError(store.BanEviction(key)) fInfo, err := store.Stat(key, true) + require.NoError(err) + _, err = store.Stat(key, false) + require.NoError(err) + + require.False(fInfo.IsDir()) + require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) + require.Equal(_blobFileName, fInfo.Name()) + require.Equal(int64(10), fInfo.Size()) + }) + t.Run("incomplete blob", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 10*memsize.B) + require.NoError(err) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + + _, err = store.Stat(key, true) + require.Equal(os.ErrNotExist, err) + fInfo, err := store.Stat(key, false) + require.NoError(err) + + require.False(fInfo.IsDir()) + require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) + require.Equal(_blobFileName, fInfo.Name()) + require.Equal(int64(10), fInfo.Size()) + }) + + t.Run("incomplete, unevictable blob", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 10*memsize.B) + require.NoError(err) + defer func() { require.NoError(f.Close()) }() + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + require.NoError(store.BanEviction(key)) + + _, err = store.Stat(key, true) require.Equal(os.ErrNotExist, err) - require.Nil(fInfo) + fInfo, err := store.Stat(key, false) + require.NoError(err) + + require.False(fInfo.IsDir()) + require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) + require.Equal(_blobFileName, fInfo.Name()) + require.Equal(int64(10), fInfo.Size()) }) t.Run("non-existent blob", func(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() _, err := store.Stat(key, true) require.Equal(os.ErrNotExist, err) + _, err = store.Stat(key, false) + require.Equal(os.ErrNotExist, err) }) } func TestList(t *testing.T) { require := require.New(t) - store, _ := testStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB) + require.Empty(store.List(false)) require.Empty(store.List(true)) - uncommittedBlobKey := core.DigestFixture().Hex() - w, err := store.Create(uncommittedBlobKey, 10*memsize.B) + incompleteBlobKey := core.DigestFixture().Hex() + f, err := store.Create(incompleteBlobKey, 10*memsize.B) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + completeBlobKey := core.DigestFixture().Hex() + f, err = store.Create(completeBlobKey, 10*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) require.NoError(err) - committedBlobKey := core.DigestFixture().Hex() - w, err = store.Create(committedBlobKey, 10*memsize.B) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + require.NoError(store.MarkComplete(completeBlobKey)) + unevictableIncompleteBlobKey := core.DigestFixture().Hex() + f, err = store.Create(unevictableIncompleteBlobKey, 10*memsize.B) require.NoError(err) - _, err = io.Copy(w, bytes.NewReader(make([]byte, 10))) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) require.NoError(err) - require.NoError(store.MarkComplete(committedBlobKey)) - // TODO - add an unevictable blob here once that's implemented + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + unevictableCompleteBlobKey := core.DigestFixture().Hex() + f, err = store.Create(unevictableCompleteBlobKey, 10*memsize.B) + require.NoError(err) + _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) + require.NoError(err) + defer func() { require.NoError(f.Close()) }() + require.NoError(store.MarkComplete(unevictableCompleteBlobKey)) - require.Equal([]string{committedBlobKey}, store.List(true)) + require.Equal([]string{completeBlobKey, unevictableCompleteBlobKey}, store.List(true)) + require.Equal([]string{completeBlobKey, unevictableCompleteBlobKey, incompleteBlobKey, unevictableIncompleteBlobKey}, store.List(false)) +} + +func TestMetadata(t *testing.T) { + t.Skip("TODO - implement") + // in-place + // after move works as intended } func TestPathing(t *testing.T) { require := require.New(t) - store, rootDir := testStore(t, 10*memsize.KB) + store, rootDir := newTestStore(t, 10*memsize.KB) key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - committed := false - dirPath := store.dirPath(key, committed) + complete := false + dirPath := store.dirPath(key, complete) wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" require.Equal(wantDirPath, dirPath) - blobPath := store.blobPath(key, committed) + blobPath := store.blobPath(key, complete) wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" require.Equal(wantBlobPath, blobPath) - committed = true - dirPath = store.dirPath(key, committed) + complete = true + dirPath = store.dirPath(key, complete) wantDirPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" require.Equal(wantDirPath, dirPath) - blobPath = store.blobPath(key, committed) + blobPath = store.blobPath(key, complete) wantBlobPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" require.Equal(wantBlobPath, blobPath) } - -func TestForbidEviction(t *testing.T) { - t.Skip("TODO") - // unevictable blobs can neither be evicted nor deleted. remarking as evictable works but resets access time. -} From c7bd73033067b7caddddd891df7ea9f592d6420f Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Wed, 10 Jun 2026 17:49:09 +0200 Subject: [PATCH 03/10] feat(DiskStore): implement storing blob metadata --- lib/store/disk_store.go | 113 ++++++++++++++++++++++++---- lib/store/disk_store_test.go | 138 +++++++++++++++++++++++++++++++++-- 2 files changed, 231 insertions(+), 20 deletions(-) diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index c4e4ceb85..7b79f535a 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -4,6 +4,7 @@ import ( "container/list" "errors" "fmt" + "io" "os" "path/filepath" "sync" @@ -57,10 +58,11 @@ type el struct { size uint64 } -// NewDiskStore creates a [DiskStore]. +// NewDiskStore creates a [*DiskStore]. func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { // TODO - create a Config struct. // TODO - recover persisted state in case of crash. might need to evict. emit log if so. + // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. return &DiskStore{ dir: rootDir, capacity: capacityBytes, @@ -73,7 +75,7 @@ func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { } // Open returns an FD to a file in the store. [os.ErrNotExists] is returned on missing entry. -// The blob cannot be evicted before the client calls Close() on it. +// The blob cannot be evicted before the client calls Close on it. func (s *DiskStore) Open(key string, ignoreIncomplete bool) (FileReadWriter, error) { s.mu.Lock() defer s.mu.Unlock() @@ -211,6 +213,8 @@ func (s *DiskStore) deleteFromDisk(key string, complete bool) error { // MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless BanEviction has been called). // Additionally, read APIs may optionally filter out incomplete blobs. func (s *DiskStore) MarkComplete(key string) error { + // TODO - check if we can derive when a blob is considered complete (e.g. when client calls Close on file (although that depends on the + // assumption that Close means the file is complete which may not be true if the client that created the file expects another client to continue mutating it)). s.mu.Lock() defer s.mu.Unlock() @@ -253,8 +257,7 @@ func (s *DiskStore) MarkComplete(key string) error { } func (s *DiskStore) checkDiskIfUnevictable(key string, complete bool) (bool, error) { - dirPath := s.dirPath(key, complete) - flagBlobPath := filepath.Join(dirPath, _evictionBannedFileName) + flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) _, err := os.Stat(flagBlobPath) if err == nil { return true, nil @@ -352,8 +355,7 @@ func (s *DiskStore) BanEviction(key string) error { } complete := okBlob - dirPath := s.dirPath(key, complete) - flagBlobPath := filepath.Join(dirPath, _evictionBannedFileName) + flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) // We persist the ban as a flag file on disk, such that after // a system crash, we can recover the ban. f, err := os.OpenFile(flagBlobPath, os.O_RDONLY|os.O_CREATE, _defaultFilePerm) @@ -403,9 +405,7 @@ func (s *DiskStore) UnbanEviction(key string) error { } return nil } - // TODO - create a helper function that finds a file with a specific suffix and returns its path - dirPath := s.dirPath(key, complete) - flagBlobPath := filepath.Join(dirPath, _evictionBannedFileName) + flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) err = os.Remove(flagBlobPath) if err != nil { return fmt.Errorf("remove file that flags eviction as banned: %w", err) @@ -420,19 +420,97 @@ func (s *DiskStore) UnbanEviction(key string) error { return nil } -// WriteMetadata atomically stores `md` on disk. Can be called on both complete and incomplete blobs. +// SetMetadata atomically sets the respective metadata for a blob. Works on any blob. func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { - return errors.New("not implemented") + s.mu.Lock() + defer s.mu.Unlock() + + _, okBlob := s.blobs[key] + _, okUnevictable := s.unevictableBlobs[key] + _, okIncomplete := s.incompleteBlobs[key] + if !okBlob && !okUnevictable && !okIncomplete { + return os.ErrNotExist + } + + mdData, err := md.Serialize() + if err != nil { + return fmt.Errorf("serialize metadata: %w", err) + } + complete := !okIncomplete + mdFilePath := s.sidecarFilePath(key, complete, md.GetSuffix()) + // We use a tmp file to ensure atomicity. + tmpFilePath := mdFilePath + "-tmp" + tmpFile, err := os.OpenFile(tmpFilePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, _defaultFilePerm) + if err != nil { + return fmt.Errorf("create tmp file for md: %w", err) + } + _, err = tmpFile.Write(mdData) + if err != nil { + return fmt.Errorf("write to tmp file: %w", err) + } + err = tmpFile.Close() + if err != nil { + return fmt.Errorf("close tmp file: %w", err) + } + err = os.Rename(tmpFile.Name(), mdFilePath) + if err != nil { + return fmt.Errorf("rename tmp file: %w", err) + } + return nil } // GetMetadata populates `md` if the metadata is present. Returns [os.ErrNotExists] if key is not in store. -func (s *DiskStore) GetMetadata(key string, md metadata.Metadata, ignoreIncomplete bool) error { - return errors.New("not implemented") +func (s *DiskStore) GetMetadata(key string, md metadata.Metadata, ignoreIncomplete bool) (ok bool, err error) { + s.mu.RLock() + defer s.mu.RUnlock() + + _, okBlob := s.blobs[key] + _, okUnevictable := s.unevictableBlobs[key] + _, okIncomplete := s.incompleteBlobs[key] + if !okBlob && !okUnevictable && (ignoreIncomplete || !okIncomplete) { + return false, os.ErrNotExist + } + + complete := !okIncomplete + mdFilePath := s.sidecarFilePath(key, complete, md.GetSuffix()) + mdFile, err := os.OpenFile(mdFilePath, os.O_RDONLY, _defaultFilePerm) + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + defer closers.Close(mdFile) + data, err := io.ReadAll(mdFile) + if err != nil { + return false, fmt.Errorf("read from metadata file: %w", err) + } + err = md.Deserialize(data) + if err != nil { + return false, fmt.Errorf("deserialize into metadata: %w", err) + } + return true, nil } -// DeleteMetadata removes the respective metadata, if present. +// DeleteMetadata removes any metadata of a blob with `md`'s suffix, if present. func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { - return errors.New("not implemented") + s.mu.Lock() + defer s.mu.Unlock() + + _, okBlob := s.blobs[key] + _, okUnevictable := s.unevictableBlobs[key] + _, okIncomplete := s.incompleteBlobs[key] + if !okBlob && !okUnevictable && !okIncomplete { + return os.ErrNotExist + } + complete := !okIncomplete + mdFilePath := s.sidecarFilePath(key, complete, md.GetSuffix()) + err := os.Remove(mdFilePath) + if errors.Is(err, os.ErrNotExist) { + // no-op + return nil + } + if err != nil { + return fmt.Errorf("remove metadata file: %w", err) + } + return nil } func (s *DiskStore) blobPath(key string, complete bool) string { @@ -455,3 +533,8 @@ func (s *DiskStore) dirPath(key string, complete bool) string { return filepath.Join(dirPath, key) } + +func (s *DiskStore) sidecarFilePath(key string, complete bool, sidecarFilePath string) string { + dirPath := s.dirPath(key, complete) + return filepath.Join(dirPath, sidecarFilePath) +} diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index 9f74661ed..88433c926 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -2,6 +2,7 @@ package store import ( "bytes" + "errors" "io" "io/fs" "os" @@ -14,9 +15,15 @@ import ( "github.com/stretchr/testify/require" "github.com/uber/kraken/core" + "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/memsize" ) +const ( + _dontIgnoreIncompleteFiles = false + _ignoreIncompleteFiles = true +) + func newTestStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) @@ -100,8 +107,6 @@ func TestDiskStore(t *testing.T) { } func TestEviction(t *testing.T) { - const _dontIgnoreIncompleteFiles = false - require := require.New(t) store, _ := newTestStore(t, 25*memsize.KB) // create 5 blobs - a, b, c, d, e with different sizes. @@ -596,14 +601,131 @@ func TestList(t *testing.T) { } func TestMetadata(t *testing.T) { - t.Skip("TODO - implement") - // in-place - // after move works as intended + t.Run("basic functionality", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 10*memsize.KB) + require.NoError(err) + require.NoError(f.Close()) + + mdStruct := core.MetaInfoFixture() + writtenMd := metadata.NewTorrentMeta(mdStruct) + err = store.SetMetadata(key, writtenMd) + // ensure metadata is not included in LRU eviction calculation. + require.NoError(err) + + var readMd metadata.TorrentMeta + ok, err := store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + require.NoError(err) + require.True(ok) + require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + + require.NoError(store.DeleteMetadata(key, &readMd)) + ok, err = store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + require.NoError(err) + require.False(ok) + mdFilePath := store.sidecarFilePath(key, false, readMd.GetSuffix()) + // ensure the metadata file is deleted from disk + _, err = os.Stat(mdFilePath) + require.True(errors.Is(err, os.ErrNotExist)) + // deleting a second time should be a no-op. + require.NoError(store.DeleteMetadata(key, &readMd)) + }) + + t.Run("non-existant blob", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + nonExistentKey := core.DigestFixture().Hex() + mdStruct := core.MetaInfoFixture() + md := metadata.NewTorrentMeta(mdStruct) + + err := store.SetMetadata(nonExistentKey, md) + require.Equal(os.ErrNotExist, err) + + ok, err := store.GetMetadata(nonExistentKey, md, _dontIgnoreIncompleteFiles) + require.Equal(os.ErrNotExist, err) + require.False(ok) + + err = store.DeleteMetadata(nonExistentKey, md) + require.Equal(os.ErrNotExist, err) + }) + + t.Run("metadata does not change after marking a file as complete and/or evictable/unevictable", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 1*memsize.KB) + require.NoError(err) + require.NoError(f.Close()) + + mdStruct := core.MetaInfoFixture() + writtenMd := metadata.NewTorrentMeta(mdStruct) + require.NoError(store.SetMetadata(key, writtenMd)) + + var readMd metadata.TorrentMeta + ok, err := store.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + require.Equal(os.ErrNotExist, err) + // incomplete files are ignored + require.False(ok) + + ok, err = store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + require.NoError(err) + require.True(ok) + require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + + // Repeat the tests above for an evictable file + require.NoError(store.BanEviction(key)) + ok, err = store.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + require.Equal(os.ErrNotExist, err) + // incomplete files are ignored + require.False(ok) + + ok, err = store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + require.NoError(err) + require.True(ok) + require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + + require.NoError(store.MarkComplete(key)) + ok, err = store.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + require.NoError(err) + require.True(ok) + require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + }) + t.Run("metadata fully gone after blob is evicted", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB) + keyA := core.DigestFixture().Hex() + fA, err := store.Create(keyA, 10*memsize.KB) + require.NoError(err) + require.NoError(fA.Close()) + require.NoError(store.MarkComplete(keyA)) + + md := metadata.NewTorrentMeta(core.MetaInfoFixture()) + err = store.SetMetadata(keyA, md) + store.SetMetadata(keyA, md) + complete := true + mdFilePath := store.sidecarFilePath(keyA, complete, md.GetSuffix()) + _, err = os.Stat(mdFilePath) + require.NoError(err) + + keyB := core.DigestFixture().Hex() + fB, err := store.Create(keyB, 10*memsize.KB) + require.NoError(err) + defer func() { require.NoError(fB.Close()) }() + + ok, err := store.GetMetadata(keyA, md, _dontIgnoreIncompleteFiles) + require.Equal(os.ErrNotExist, err) + require.False(ok) + _, err = os.Stat(mdFilePath) + require.True(errors.Is(err, os.ErrNotExist)) + }) } func TestPathing(t *testing.T) { require := require.New(t) store, rootDir := newTestStore(t, 10*memsize.KB) + md := metadata.NewTorrentMeta(core.MetaInfoFixture()) key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" complete := false @@ -613,6 +735,9 @@ func TestPathing(t *testing.T) { blobPath := store.blobPath(key, complete) wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" require.Equal(wantBlobPath, blobPath) + sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) + wantSidecarFilePath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" + require.Equal(wantSidecarFilePath, sidecarFilePath) complete = true dirPath = store.dirPath(key, complete) @@ -621,4 +746,7 @@ func TestPathing(t *testing.T) { blobPath = store.blobPath(key, complete) wantBlobPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" require.Equal(wantBlobPath, blobPath) + sidecarFilePath = store.sidecarFilePath(key, complete, md.GetSuffix()) + wantSidecarFilePath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" + require.Equal(wantSidecarFilePath, sidecarFilePath) } From cc3187dd02b698064bb10d500694ffae271888fa Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Wed, 10 Jun 2026 21:10:28 +0200 Subject: [PATCH 04/10] feat(DiskStore): reboot state from disk after crash --- lib/store/crash_recovery.go | 167 ++++++++++++++++++++++++++++++ lib/store/crash_recovery_test.go | 171 +++++++++++++++++++++++++++++++ lib/store/disk_store.go | 106 +++++++++++-------- lib/store/disk_store_test.go | 43 ++------ lib/store/pather.go | 44 ++++++++ lib/store/pather_test.go | 39 +++++++ 6 files changed, 489 insertions(+), 81 deletions(-) create mode 100644 lib/store/crash_recovery.go create mode 100644 lib/store/crash_recovery_test.go create mode 100644 lib/store/pather.go create mode 100644 lib/store/pather_test.go diff --git a/lib/store/crash_recovery.go b/lib/store/crash_recovery.go new file mode 100644 index 000000000..3f19227e3 --- /dev/null +++ b/lib/store/crash_recovery.go @@ -0,0 +1,167 @@ +package store + +import ( + "container/list" + "errors" + "fmt" + "io/fs" + "os" + "path/filepath" + "slices" + "strings" + "time" + + "go.uber.org/zap" +) + +type blobState struct { + key string + size uint64 + mTime time.Time + evictable bool +} + +func rebootPersistedStateAfterCrash(capacityBytes uint64, rootDir string, log *zap.SugaredLogger) (*DiskStore, error) { + completeDirPath, incompleteDirPath := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) + + // We remove incomplete entries, as we expect the processes/clients who were writing to + // these entries to also have crashed, which would leak the files on disk. + err := os.RemoveAll(incompleteDirPath) + if err != nil { + return nil, fmt.Errorf("remove incomplete blobs left from a previous service run: %w", err) + } + keys, err := rebootKeys(completeDirPath) + if err != nil { + return nil, err + } + pather := newPather(rootDir) + + completeEntries := make([]*blobState, 0) + unevictableEntries := make([]*blobState, 0) + for _, key := range keys { + bState, ok, err := rebootBlobState(key, pather) + if err != nil { + return nil, err + } + if !ok { + continue + } + if bState.evictable { + completeEntries = append(completeEntries, bState) + } else { + unevictableEntries = append(unevictableEntries, bState) + } + } + + storeSize := uint64(0) + unevictableBlobs := make(map[string]uint64, len(unevictableEntries)) + for _, e := range unevictableEntries { + unevictableBlobs[e.key] = e.size + storeSize += e.size + } + + slices.SortFunc(completeEntries, func(left, right *blobState) int { + // left-most is oldest, i.e. next-to-evict. + return left.mTime.Compare(right.mTime) + }) + blobs := make(map[string]*list.Element, len(completeEntries)) + evictQueue := list.New() + for _, e := range completeEntries { + node := evictQueue.PushBack(el{key: e.key, size: e.size}) + blobs[e.key] = node + storeSize += e.size + } + + store := &DiskStore{ + blobs: blobs, + evictQueue: evictQueue, + incompleteBlobs: make(map[string]uint64), + unevictableBlobs: unevictableBlobs, + capacity: capacityBytes, + pather: pather, + size: storeSize, + log: log, + } + + if store.size > store.capacity { + prevSize := store.size + // evicts entries until size <= capacity. + err = store.reserveSpace(0) + if err != nil { + log.With("error", err).Error("DiskStore size exceeds its capacity after service reboot. Evicting blobs from disk did not work to reduce size within capacity.") + return nil, fmt.Errorf("remove blobs to reduce store size within configured capacity") + } + evictedBytes := prevSize - store.size + log.With("evicted_bytes", evictedBytes).Warn("DiskStore size exceeds its capacity after service reboot. Successfully evicted blobs to reduce size within capacity.") + } + return store, nil +} + +func rebootBlobState(key string, pather *pather) (res *blobState, ok bool, err error) { + complete := true + blobPath := pather.blobPath(key, complete) + fInfo, err := os.Stat(blobPath) + if errors.Is(err, os.ErrNotExist) { + // For some reason, the directory for the blob exists but not the blob itself. + return nil, true, nil + } + if err != nil { + return nil, false, fmt.Errorf("stat blob entry: %w", err) + } + flagBlobPath := pather.sidecarFilePath(key, complete, _evictionBannedFileName) + isUnevictable, err := exists(flagBlobPath) + if err != nil { + return nil, false, err + } + size := fInfo.Size() + mTime := fInfo.ModTime() + return &blobState{ + key: key, + size: uint64(size), + mTime: mTime, + evictable: !isUnevictable, + }, true, nil +} + +func rebootKeys(completeDirPath string) ([]string, error) { + keys := make([]string, 0) + err := filepath.WalkDir(completeDirPath, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + + if !entry.IsDir() { + return nil + } + relPath, err := filepath.Rel(completeDirPath, path) + if err != nil { + return err + } + nameParts := strings.Split(relPath, string(filepath.Separator)) + isBlobDir := len(nameParts) == _numShards+1 + if !isBlobDir { + return nil + } + key := nameParts[len(nameParts)-1] + keys = append(keys, key) + return fs.SkipDir + }) + if err != nil { + return nil, fmt.Errorf("walk through complete dir to collect keys of blobs: %w", err) + } + return keys, nil +} + +func existsPersistedState(rootDir string) (ok bool, err error) { + completeDir, incompleteDir := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) + completeExists, err := exists(completeDir) + if err != nil { + return false, fmt.Errorf("check if store has persisted state left on disk from previous service runs: %w", err) + } + incompleteExists, err := exists(incompleteDir) + if err != nil { + return false, fmt.Errorf("check if store has persisted state left on disk from previous service runs: %w", err) + } + + return completeExists || incompleteExists, nil +} diff --git a/lib/store/crash_recovery_test.go b/lib/store/crash_recovery_test.go new file mode 100644 index 000000000..5c5cec32e --- /dev/null +++ b/lib/store/crash_recovery_test.go @@ -0,0 +1,171 @@ +package store + +import ( + "bytes" + "crypto/rand" + "io" + "os" + "testing" + + "github.com/stretchr/testify/require" + "github.com/uber/kraken/core" + "github.com/uber/kraken/lib/store/metadata" + "github.com/uber/kraken/utils/memsize" +) + +func TestCrashRecovery(t *testing.T) { + t.Run("complete blobs, evictability, and completeness are recovered", func(t *testing.T) { + require := require.New(t) + store, rootDir := newTestStore(t, 10*memsize.KB) + + completeEvictableF, completeEvictableKey := newTestFile(t, store, 2*memsize.KB) + completeEvictableData := fillWithRandomData(t, completeEvictableF, 2*memsize.KB) + // We don't have to test case where file is not closed, as linux closes all FDs owned by a process upon process death. + require.NoError(completeEvictableF.Close()) + require.NoError(store.MarkComplete(completeEvictableKey)) + + completeUnevictableF, completeUnevictableKey := newTestFile(t, store, 2*memsize.KB) + completeUnevictableData := fillWithRandomData(t, completeUnevictableF, 2*memsize.KB) + require.NoError(completeUnevictableF.Close()) + require.NoError(store.MarkComplete(completeUnevictableKey)) + require.NoError(store.BanEviction(completeUnevictableKey)) + + incompleteEvictableF, incompleteEvictableKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, incompleteEvictableF, 2*memsize.KB) + require.NoError(incompleteEvictableF.Close()) + + incompleteUnevictableF, incompleteUnevictableKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, incompleteUnevictableF, 2*memsize.KB) + require.NoError(incompleteUnevictableF.Close()) + require.NoError(store.BanEviction(incompleteUnevictableKey)) + + // Assume that the application crashes here. The application would restart and call `NewDiskStore`. + rebootedStore, err := NewDiskStore(10*memsize.KB, rootDir) + require.NoError(err) + + // Incomplete files get dropped. + _, err = rebootedStore.Stat(incompleteEvictableKey, _dontIgnoreIncompleteFiles) + require.ErrorIs(err, os.ErrNotExist) + _, err = rebootedStore.Stat(incompleteUnevictableKey, _dontIgnoreIncompleteFiles) + require.ErrorIs(err, os.ErrNotExist) + + // Complete files are recovered. + f, err := rebootedStore.Open(completeEvictableKey, _ignoreIncompleteFiles) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err := io.ReadAll(f) + require.NoError(err) + require.Equal(completeEvictableData, data) + complete := true + unevictable, err := rebootedStore.checkDiskIfUnevictable(completeEvictableKey, complete) + require.NoError(err) + require.False(unevictable) + + f, err = rebootedStore.Open(completeUnevictableKey, _ignoreIncompleteFiles) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(completeUnevictableData, data) + unevictable, err = rebootedStore.checkDiskIfUnevictable(completeUnevictableKey, complete) + require.NoError(err) + require.True(unevictable) + }) + + t.Run("metadata is recovered", func(t *testing.T) { + require := require.New(t) + store, rootDir := newTestStore(t, 10*memsize.KB) + + f, key := newTestFile(t, store, 2*memsize.KB) + require.NoError(f.Close()) + require.NoError(store.MarkComplete(key)) + writtenMd := metadata.NewTorrentMeta(core.MetaInfoFixture()) + require.NoError(store.SetMetadata(key, writtenMd)) + + // Assume that the application crashes here. The application would restart and call `NewDiskStore`. + rebootedStore, err := NewDiskStore(10*memsize.KB, rootDir) + require.NoError(err) + + var readMd metadata.TorrentMeta + ok, err := rebootedStore.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + require.NoError(err) + require.True(ok) + require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) + }) + + t.Run("lru order is approximated and blobs are evicted if store size exceeds capacity", func(t *testing.T) { + require := require.New(t) + store, rootDir := newTestStore(t, 10*memsize.KB) + + aF, aKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, aF, 2*memsize.KB) + require.NoError(aF.Close()) + require.NoError(store.MarkComplete(aKey)) + require.NoError(store.BanEviction(aKey)) + + bF, _ := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, bF, 1*memsize.KB) + require.NoError(bF.Close()) + + cF, cKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, cF, 2*memsize.KB) + require.NoError(cF.Close()) + require.NoError(store.MarkComplete(cKey)) + + dF, dKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, dF, 2*memsize.KB) + require.NoError(dF.Close()) + require.NoError(store.MarkComplete(dKey)) + + eF, eKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, eF, 2*memsize.KB) + require.NoError(eF.Close()) + require.NoError(store.MarkComplete(eKey)) + require.Equal([]string{cKey, dKey, eKey}, store.evictionOrder()) // a is unevictable and b is incomplete + + // reset the access time for d + dF, err := store.Open(dKey, _ignoreIncompleteFiles) + require.NoError(err) + require.NoError(dF.Close()) + evictionOrderBeforeCrash := store.evictionOrder() + require.Equal([]string{cKey, eKey, dKey}, evictionOrderBeforeCrash) // a is unevictable and b is incomplete + + // Assume that the application restarts here. + rebootedStore, err := NewDiskStore(10*memsize.KB, rootDir) + require.NoError(err) + + // LRU order is approximated, but not exact. + rebootedEvictionOrder := rebootedStore.evictionOrder() + require.NotEqual(evictionOrderBeforeCrash, rebootedEvictionOrder) + wantEvictionOrder := []string{cKey, dKey, eKey} + require.Equal(wantEvictionOrder, rebootedEvictionOrder) + + // Assume we redeploy the service with a smaller capacity for the disk store: + rebootedSmallerStore, err := NewDiskStore(6*memsize.KB, rootDir) + require.NoError(err) + + // since 10KB of blobs are in store, `c` gets evicted to put the store back within its capacity. + _, err = rebootedSmallerStore.Stat(cKey, _dontIgnoreIncompleteFiles) + require.ErrorIs(err, os.ErrNotExist) + + require.Equal([]string{dKey, eKey}, rebootedSmallerStore.evictionOrder()) + _, err = rebootedSmallerStore.Stat(aKey, _dontIgnoreIncompleteFiles) + require.NoError(err) + }) +} + +func TestStoreWorksWhenFileSizeNotCorrect(t *testing.T) { + // Verify that the store works correctly when the reserved size for a file (the one passed by the client in Create) is different + // than its actual size. The store is expected to consistently use EITHER the client-given size OR the actual size of files, but not both. + // If we mix them, this could break the eviction logic - imagine the user uploads a 2GB size but reports it as 1.9GB. Eviction works correctly + // as long as we reserve 2GB upon upload to store and release 2GB upon deletion/eviction from store. BUT if we reserve 2GB and free 1.9GB + // or vice-versa, it could lead to over/under-reservation. + t.Skip("TODO") +} + +func fillWithRandomData(t *testing.T, f FileReadWriter, sizeBytes uint64) []byte { + data := make([]byte, sizeBytes) + _, err := rand.Read(data) + require.NoError(t, err) + _, err = io.Copy(f, bytes.NewReader(data)) + require.NoError(t, err) + return data +} diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index 7b79f535a..8a3d806e1 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -11,17 +11,17 @@ import ( "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/closers" + "github.com/uber/kraken/utils/log" + "go.uber.org/zap" ) const ( // _defaultShardIDLength is the number of bytes of file digest to be used for shard ID. // For every byte (2 HEX char), one more level of directories will be created. _defaultFilePerm = 0775 - _defaultShardIDLength = 2 - _incompleteSubDir = "incomplete" - _completeSubDir = "complete" - _blobFileName = "data" _evictionBannedFileName = "_eviction_banned" + // TODO - change this not to be hardcoded when configurable sharding is implemented + _numShards = 2 ) // DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. @@ -34,13 +34,13 @@ const ( // // - Supports (un-)marking blobs as non-evictable (may be needed when that data must be written back to remote storage). // -// - Crash-resistant - all state is restored upon restart (LRU order is approximated through file `ctime`). +// - Crash-resistant - all state is restored upon restart (LRU order is approximated through file `mtime`). // // - Uses directory sharding to speed up disk performance. type DiskStore struct { + *pather capacity uint64 size uint64 // includes both used and reserved space - dir string // synchronizes mem state access and syscalls to the fs in the APIs (opening, moving files, etc.) mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. // complete, evictable blobs. @@ -51,6 +51,7 @@ type DiskStore struct { incompleteBlobs map[string]uint64 // complete blobs that cannot be evicted. unevictableBlobs map[string]uint64 + log *zap.SugaredLogger } type el struct { @@ -59,23 +60,45 @@ type el struct { } // NewDiskStore creates a [*DiskStore]. +// incomplete blobs are deleted +// approximates lru eviction order by file `mtime` func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { // TODO - create a Config struct. - // TODO - recover persisted state in case of crash. might need to evict. emit log if so. // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. - return &DiskStore{ - dir: rootDir, - capacity: capacityBytes, - blobs: make(map[string]*list.Element), - evictQueue: list.New(), - incompleteBlobs: make(map[string]uint64), - unevictableBlobs: make(map[string]uint64), - size: 0, - }, nil + // TODO - move disk store files into their own directory and package. + + log := log.Default().With("module", "disk_store") + ok, err := existsPersistedState(rootDir) + if err != nil { + err = fmt.Errorf("could not check if previously-left persisted state exists on disk: %w", err) + log.With("error", err).Error("Failed to initialize disk store") + return nil, err + } + if !ok { + log.Info("Did not find any previously persisted state to reboot for DiskStore - initializing a new, empty DiskStore") + return &DiskStore{ + capacity: capacityBytes, + blobs: make(map[string]*list.Element), + evictQueue: list.New(), + incompleteBlobs: make(map[string]uint64), + unevictableBlobs: make(map[string]uint64), + size: 0, + log: log, + pather: newPather(rootDir), + }, nil + } + + store, err := rebootPersistedStateAfterCrash(capacityBytes, rootDir, log) + if err != nil { + err = fmt.Errorf("reboot persisted state into memory: %w", err) + log.With("error", err).Error("Failed to initialize disk store") + return nil, err + } + log.With("num_blobs", len(store.blobs)+len(store.unevictableBlobs)).Info("Successfully initialized disk store") + return store, nil } // Open returns an FD to a file in the store. [os.ErrNotExists] is returned on missing entry. -// The blob cannot be evicted before the client calls Close on it. func (s *DiskStore) Open(key string, ignoreIncomplete bool) (FileReadWriter, error) { s.mu.Lock() defer s.mu.Unlock() @@ -258,14 +281,11 @@ func (s *DiskStore) MarkComplete(key string) error { func (s *DiskStore) checkDiskIfUnevictable(key string, complete bool) (bool, error) { flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) - _, err := os.Stat(flagBlobPath) - if err == nil { - return true, nil - } - if errors.Is(err, os.ErrNotExist) { - return false, nil + unevictable, err := exists(flagBlobPath) + if err != nil { + return false, err } - return false, fmt.Errorf("stat: %w", err) + return unevictable, nil } // Delete removes a blob and its [metadata.Metadata] from the store. Works on any blob. @@ -513,28 +533,26 @@ func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { return nil } -func (s *DiskStore) blobPath(key string, complete bool) string { - dirName := s.dirPath(key, complete) - return filepath.Join(dirName, _blobFileName) -} +// used during testing +func (s *DiskStore) evictionOrder() []string { + s.mu.RLock() + defer s.mu.RUnlock() -func (s *DiskStore) dirPath(key string, complete bool) string { - // TODO - allow config to specify whether to shard or not. Allow no sharding, so we can replace [SimpleStore]. - subDirName := _incompleteSubDir - if complete { - subDirName = _completeSubDir + evictionOrder := make([]string, 0) + for curr := s.evictQueue.Front(); curr != nil; curr = curr.Next() { + currEl := curr.Value.(el) + evictionOrder = append(evictionOrder, currEl.key) } - dirPath := filepath.Join(s.dir, subDirName) - for i := 0; i < int(_defaultShardIDLength) && i < len(key)/2; i++ { - // (1 byte = 2 char of file name assumming file name is in HEX) - dirName := key[i*2 : i*2+2] - dirPath = filepath.Join(dirPath, dirName) - } - - return filepath.Join(dirPath, key) + return evictionOrder } -func (s *DiskStore) sidecarFilePath(key string, complete bool, sidecarFilePath string) string { - dirPath := s.dirPath(key, complete) - return filepath.Join(dirPath, sidecarFilePath) +func exists(path string) (ok bool, err error) { + _, err = os.Stat(path) + if err == nil { + return true, nil + } + if errors.Is(err, os.ErrNotExist) { + return false, nil + } + return false, fmt.Errorf("stat: %w", err) } diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index 88433c926..59463bd2f 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -20,6 +20,7 @@ import ( ) const ( + // TODO - consider exporting these constants for clients to consume to avoid naked params _dontIgnoreIncompleteFiles = false _ignoreIncompleteFiles = true ) @@ -46,7 +47,11 @@ func newTestFile(t *testing.T, store *DiskStore, size uint64) (f FileReadWriter, // does not count 1) the directories for sharding, 2) metadata files, and 3) the _eviction_banned flag file. func numBlobsOnDisk(t *testing.T, store *DiskStore) int { numBlobs := 0 - err := filepath.Walk(store.dir, func(path string, info fs.FileInfo, err error) error { + err := filepath.WalkDir(store.dir, func(path string, _ fs.DirEntry, err error) error { + if err != nil { + return err + } + if !strings.HasSuffix(path, _blobFileName) { return nil } @@ -217,13 +222,6 @@ func TestEviction(t *testing.T) { require.Equal(2, numBlobsOnDisk(t, store)) } -func TestCrashRecovery(t *testing.T) { - t.Skip("TODO - implement") - - // test each blob type is as expected in the store's internal state - // test LRU order is approximated correctly after recovery (ctime shd be different than access time) -} - func TestParallelAccessToSingleFile(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB) @@ -721,32 +719,3 @@ func TestMetadata(t *testing.T) { require.True(errors.Is(err, os.ErrNotExist)) }) } - -func TestPathing(t *testing.T) { - require := require.New(t) - store, rootDir := newTestStore(t, 10*memsize.KB) - md := metadata.NewTorrentMeta(core.MetaInfoFixture()) - - key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - complete := false - dirPath := store.dirPath(key, complete) - wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - require.Equal(wantDirPath, dirPath) - blobPath := store.blobPath(key, complete) - wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" - require.Equal(wantBlobPath, blobPath) - sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) - wantSidecarFilePath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" - require.Equal(wantSidecarFilePath, sidecarFilePath) - - complete = true - dirPath = store.dirPath(key, complete) - wantDirPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - require.Equal(wantDirPath, dirPath) - blobPath = store.blobPath(key, complete) - wantBlobPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" - require.Equal(wantBlobPath, blobPath) - sidecarFilePath = store.sidecarFilePath(key, complete, md.GetSuffix()) - wantSidecarFilePath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" - require.Equal(wantSidecarFilePath, sidecarFilePath) -} diff --git a/lib/store/pather.go b/lib/store/pather.go new file mode 100644 index 000000000..7dccb199f --- /dev/null +++ b/lib/store/pather.go @@ -0,0 +1,44 @@ +package store + +import "path/filepath" + +const ( + _defaultShardIDLength = 2 + _incompleteSubDir = "incomplete" + _completeSubDir = "complete" + _blobFileName = "data" +) + +type pather struct { + dir string +} + +func newPather(rootDir string) *pather { + return &pather{dir: rootDir} +} + +func (p *pather) blobPath(key string, complete bool) string { + dirName := p.dirPath(key, complete) + return filepath.Join(dirName, _blobFileName) +} + +func (p *pather) dirPath(key string, complete bool) string { + // TODO - allow config to specify whether to shard or not. Allow no sharding, so we can replace [SimpleStore]. + subDirName := _incompleteSubDir + if complete { + subDirName = _completeSubDir + } + dirPath := filepath.Join(p.dir, subDirName) + for i := 0; i < int(_defaultShardIDLength) && i < len(key)/2; i++ { + // (1 byte = 2 char of file name assumming file name is in HEX) + dirName := key[i*2 : i*2+2] + dirPath = filepath.Join(dirPath, dirName) + } + + return filepath.Join(dirPath, key) +} + +func (p *pather) sidecarFilePath(key string, complete bool, sidecarFilePath string) string { + dirPath := p.dirPath(key, complete) + return filepath.Join(dirPath, sidecarFilePath) +} diff --git a/lib/store/pather_test.go b/lib/store/pather_test.go new file mode 100644 index 000000000..1a1545e78 --- /dev/null +++ b/lib/store/pather_test.go @@ -0,0 +1,39 @@ +package store + +import ( + "testing" + + "github.com/stretchr/testify/require" + "github.com/uber/kraken/core" + "github.com/uber/kraken/lib/store/metadata" + "github.com/uber/kraken/utils/memsize" +) + +func TestPather(t *testing.T) { + require := require.New(t) + store, rootDir := newTestStore(t, 10*memsize.KB) + md := metadata.NewTorrentMeta(core.MetaInfoFixture()) + + key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + complete := false + dirPath := store.dirPath(key, complete) + wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + require.Equal(wantDirPath, dirPath) + blobPath := store.blobPath(key, complete) + wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" + require.Equal(wantBlobPath, blobPath) + sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) + wantSidecarFilePath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" + require.Equal(wantSidecarFilePath, sidecarFilePath) + + complete = true + dirPath = store.dirPath(key, complete) + wantDirPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + require.Equal(wantDirPath, dirPath) + blobPath = store.blobPath(key, complete) + wantBlobPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" + require.Equal(wantBlobPath, blobPath) + sidecarFilePath = store.sidecarFilePath(key, complete, md.GetSuffix()) + wantSidecarFilePath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" + require.Equal(wantSidecarFilePath, sidecarFilePath) +} From ea4e5f69ff8608dd03f5dd750e37088fd34caf53 Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Thu, 11 Jun 2026 22:10:40 +0200 Subject: [PATCH 05/10] feat(DiskStore): simplify implementation 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. --- lib/store/crash_recovery.go | 49 +++-- lib/store/crash_recovery_test.go | 18 +- lib/store/disk_store.go | 335 ++++++++++++------------------- lib/store/disk_store_test.go | 217 +++++++++++--------- lib/store/pather.go | 2 + lib/store/pather_test.go | 46 +++-- 6 files changed, 312 insertions(+), 355 deletions(-) diff --git a/lib/store/crash_recovery.go b/lib/store/crash_recovery.go index 3f19227e3..15b1625e9 100644 --- a/lib/store/crash_recovery.go +++ b/lib/store/crash_recovery.go @@ -14,6 +14,8 @@ import ( "go.uber.org/zap" ) +const _numShards = 2 // TODO - change this not to be hardcoded when configurable sharding is implemented + type blobState struct { key string size uint64 @@ -36,7 +38,7 @@ func rebootPersistedStateAfterCrash(capacityBytes uint64, rootDir string, log *z } pather := newPather(rootDir) - completeEntries := make([]*blobState, 0) + evictableEntries := make([]*blobState, 0) unevictableEntries := make([]*blobState, 0) for _, key := range keys { bState, ok, err := rebootBlobState(key, pather) @@ -47,40 +49,47 @@ func rebootPersistedStateAfterCrash(capacityBytes uint64, rootDir string, log *z continue } if bState.evictable { - completeEntries = append(completeEntries, bState) + evictableEntries = append(evictableEntries, bState) } else { unevictableEntries = append(unevictableEntries, bState) } } storeSize := uint64(0) - unevictableBlobs := make(map[string]uint64, len(unevictableEntries)) - for _, e := range unevictableEntries { - unevictableBlobs[e.key] = e.size - storeSize += e.size + blobs := make(map[string]*blob, 0) + for _, bState := range unevictableEntries { + blobs[bState.key] = &blob{ + size: bState.size, + complete: true, + node: nil, + evictionBanned: true, + } + storeSize += bState.size } - slices.SortFunc(completeEntries, func(left, right *blobState) int { + slices.SortFunc(evictableEntries, func(left, right *blobState) int { // left-most is oldest, i.e. next-to-evict. return left.mTime.Compare(right.mTime) }) - blobs := make(map[string]*list.Element, len(completeEntries)) evictQueue := list.New() - for _, e := range completeEntries { - node := evictQueue.PushBack(el{key: e.key, size: e.size}) - blobs[e.key] = node - storeSize += e.size + for _, bState := range evictableEntries { + node := evictQueue.PushBack(bState.key) + blobs[bState.key] = &blob{ + size: bState.size, + complete: true, + node: node, + evictionBanned: false, + } + storeSize += bState.size } store := &DiskStore{ - blobs: blobs, - evictQueue: evictQueue, - incompleteBlobs: make(map[string]uint64), - unevictableBlobs: unevictableBlobs, - capacity: capacityBytes, - pather: pather, - size: storeSize, - log: log, + blobs: blobs, + evictQueue: evictQueue, + capacity: capacityBytes, + pather: pather, + size: storeSize, + log: log, } if store.size > store.capacity { diff --git a/lib/store/crash_recovery_test.go b/lib/store/crash_recovery_test.go index 5c5cec32e..9aee117d3 100644 --- a/lib/store/crash_recovery_test.go +++ b/lib/store/crash_recovery_test.go @@ -44,13 +44,14 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) // Incomplete files get dropped. - _, err = rebootedStore.Stat(incompleteEvictableKey, _dontIgnoreIncompleteFiles) + _, err = rebootedStore.Stat(incompleteEvictableKey, CheckIncompleteBlobs) require.ErrorIs(err, os.ErrNotExist) - _, err = rebootedStore.Stat(incompleteUnevictableKey, _dontIgnoreIncompleteFiles) + _, err = rebootedStore.Stat(incompleteUnevictableKey, CheckIncompleteBlobs) require.ErrorIs(err, os.ErrNotExist) // Complete files are recovered. - f, err := rebootedStore.Open(completeEvictableKey, _ignoreIncompleteFiles) + f, err := rebootedStore.Open(completeEvictableKey, IgnoreIncompleteBlobs) + require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err := io.ReadAll(f) require.NoError(err) @@ -60,7 +61,8 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) require.False(unevictable) - f, err = rebootedStore.Open(completeUnevictableKey, _ignoreIncompleteFiles) + f, err = rebootedStore.Open(completeUnevictableKey, IgnoreIncompleteBlobs) + require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) require.NoError(err) @@ -85,7 +87,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) var readMd metadata.TorrentMeta - ok, err := rebootedStore.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + ok, err := rebootedStore.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) @@ -122,7 +124,7 @@ func TestCrashRecovery(t *testing.T) { require.Equal([]string{cKey, dKey, eKey}, store.evictionOrder()) // a is unevictable and b is incomplete // reset the access time for d - dF, err := store.Open(dKey, _ignoreIncompleteFiles) + dF, err := store.Open(dKey, IgnoreIncompleteBlobs) require.NoError(err) require.NoError(dF.Close()) evictionOrderBeforeCrash := store.evictionOrder() @@ -143,11 +145,11 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) // since 10KB of blobs are in store, `c` gets evicted to put the store back within its capacity. - _, err = rebootedSmallerStore.Stat(cKey, _dontIgnoreIncompleteFiles) + _, err = rebootedSmallerStore.Stat(cKey, CheckIncompleteBlobs) require.ErrorIs(err, os.ErrNotExist) require.Equal([]string{dKey, eKey}, rebootedSmallerStore.evictionOrder()) - _, err = rebootedSmallerStore.Stat(aKey, _dontIgnoreIncompleteFiles) + _, err = rebootedSmallerStore.Stat(aKey, CheckIncompleteBlobs) require.NoError(err) }) } diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index 8a3d806e1..849025752 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -15,13 +15,17 @@ import ( "go.uber.org/zap" ) +// Whether the store's read APIs ignore incomplete blobs. const ( - // _defaultShardIDLength is the number of bytes of file digest to be used for shard ID. - // For every byte (2 HEX char), one more level of directories will be created. + IgnoreIncompleteBlobs = true + CheckIncompleteBlobs = false +) + +const ( + _completeBlob = true + _incompleteBlob = false _defaultFilePerm = 0775 _evictionBannedFileName = "_eviction_banned" - // TODO - change this not to be hardcoded when configurable sharding is implemented - _numShards = 2 ) // DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. @@ -34,34 +38,36 @@ const ( // // - Supports (un-)marking blobs as non-evictable (may be needed when that data must be written back to remote storage). // -// - Crash-resistant - all state is restored upon restart (LRU order is approximated through file `mtime`). +// - Crash-resistant - all state is restored upon restart (check [NewDiskStore] for details). // // - Uses directory sharding to speed up disk performance. type DiskStore struct { - *pather - capacity uint64 - size uint64 // includes both used and reserved space + capacity uint64 + size uint64 // includes both used and reserved space. + blobs map[string]*blob // TODO - consider whether it's better to use struct instead of pointer to reduce GC stress. + evictQueue *list.List // Back is most recently used, front is the next to evict. // synchronizes mem state access and syscalls to the fs in the APIs (opening, moving files, etc.) - mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. - // complete, evictable blobs. - blobs map[string]*list.Element // value of [list.Element] is [el]. - // Back is most recently used, front is the next to evict. - evictQueue *list.List - // incomplete blobs that may or may not be evictable. - incompleteBlobs map[string]uint64 - // complete blobs that cannot be evicted. - unevictableBlobs map[string]uint64 - log *zap.SugaredLogger + mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. + log *zap.SugaredLogger + *pather } -type el struct { - key string - size uint64 +type blob struct { + node *list.Element // value of [list.Element] is [string]. + size uint64 + complete bool + evictionBanned bool } -// NewDiskStore creates a [*DiskStore]. -// incomplete blobs are deleted -// approximates lru eviction order by file `mtime` +// NewDiskStore initializes a new [*DiskStore]. If the store has been initialized in the same +// directory before, its state is recovered from disk with the following caveats: +// +// - incomplete blobs are deleted to prevent leaks, in case the clients have crashed. +// +// - the blobs' sizes are recovered from disk, which may differ from the client-provided sizes in Create. This is not a problem. +// +// - if the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), +// it evicts blobs until size is within capacity. func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { // TODO - create a Config struct. // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. @@ -77,14 +83,12 @@ func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { if !ok { log.Info("Did not find any previously persisted state to reboot for DiskStore - initializing a new, empty DiskStore") return &DiskStore{ - capacity: capacityBytes, - blobs: make(map[string]*list.Element), - evictQueue: list.New(), - incompleteBlobs: make(map[string]uint64), - unevictableBlobs: make(map[string]uint64), - size: 0, - log: log, - pather: newPather(rootDir), + capacity: capacityBytes, + size: 0, + blobs: make(map[string]*blob), + evictQueue: list.New(), + log: log, + pather: newPather(rootDir), }, nil } @@ -94,7 +98,7 @@ func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { log.With("error", err).Error("Failed to initialize disk store") return nil, err } - log.With("num_blobs", len(store.blobs)+len(store.unevictableBlobs)).Info("Successfully initialized disk store") + log.With("num_blobs", len(store.blobs)).Info("Successfully initialized disk store") return store, nil } @@ -103,36 +107,20 @@ func (s *DiskStore) Open(key string, ignoreIncomplete bool) (FileReadWriter, err s.mu.Lock() defer s.mu.Unlock() - unevictableBlobSize, okUnevictable := s.unevictableBlobs[key] - if okUnevictable { - complete := true - return s.open(key, unevictableBlobSize, complete) - } - - incompleteBlobSize, okIncomplete := s.incompleteBlobs[key] - if okIncomplete && !ignoreIncomplete { - complete := false - return s.open(key, incompleteBlobSize, complete) - } - - node, okBlob := s.blobs[key] - if !okBlob { + b, ok := s.blobs[key] + if !ok || (ignoreIncomplete && !b.complete) { return nil, os.ErrNotExist } - size := node.Value.(el).size - s.evictQueue.MoveToBack(node) - complete := true - return s.open(key, size, complete) -} - -func (s *DiskStore) open(key string, size uint64, complete bool) (FileReadWriter, error) { - path := s.blobPath(key, complete) + if b.node != nil { + s.evictQueue.MoveToBack(b.node) + } + path := s.blobPath(key, b.complete) f, err := os.OpenFile(path, os.O_RDWR, _defaultFilePerm) if err != nil { return nil, fmt.Errorf("open: %w", err) } - return newReadWriter(f, size), nil + return newReadWriter(f, b.size), nil } // Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExists] if the blob is not found. @@ -142,15 +130,11 @@ func (s *DiskStore) Stat(key string, ignoreIncomplete bool) (os.FileInfo, error) s.mu.Lock() defer s.mu.Unlock() - _, okBlob := s.blobs[key] - _, okUnevictable := s.unevictableBlobs[key] - _, okIncomplete := s.incompleteBlobs[key] - - if !okBlob && !okUnevictable && (ignoreIncomplete || !okIncomplete) { + b, ok := s.blobs[key] + if !ok || (ignoreIncomplete && !b.complete) { return nil, os.ErrNotExist } - complete := okBlob || okUnevictable - blobPath := s.blobPath(key, complete) + blobPath := s.blobPath(key, b.complete) return os.Stat(blobPath) } @@ -162,29 +146,26 @@ func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) s.mu.Lock() defer s.mu.Unlock() - if _, ok := s.blobs[key]; ok { + if b, ok := s.blobs[key]; ok { // TODO - consider whether we need public errors for these cases. - return nil, errors.New("blob is already in store") - } - if _, ok := s.unevictableBlobs[key]; ok { - return nil, errors.New("blob is already in store") - } - if _, ok := s.incompleteBlobs[key]; ok { - return nil, errors.New("blob is already in store (it is incomplete)") + if b.complete { + return nil, errors.New("blob is already in store") + } else { + return nil, errors.New("blob is already in store (it is incomplete)") + } } if err := s.reserveSpace(sizeBytes); err != nil { return nil, fmt.Errorf("reserve space: %w", err) } - complete := false - dirName := s.dirPath(key, complete) + dirName := s.dirPath(key, _incompleteBlob) err := os.MkdirAll(dirName, _defaultFilePerm) if err != nil { s.releaseSpace(sizeBytes) return nil, fmt.Errorf("ensure dir: %w", err) } - blobPath := s.blobPath(key, complete) + blobPath := s.blobPath(key, _incompleteBlob) flag := os.O_RDWR | os.O_CREATE | os.O_EXCL f, err := os.OpenFile(blobPath, flag, _defaultFilePerm) if err != nil { @@ -192,7 +173,12 @@ func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) return nil, fmt.Errorf("open file: %w", err) } - s.incompleteBlobs[key] = sizeBytes + s.blobs[key] = &blob{ + size: sizeBytes, + node: nil, + complete: false, + evictionBanned: false, + } return newReadWriter(f, sizeBytes), nil } @@ -205,17 +191,17 @@ func (s *DiskStore) reserveSpace(space uint64) error { } toEvictNode := s.evictQueue.Front() - toEvictEl := toEvictNode.Value.(el) + toEvictKey := toEvictNode.Value.(string) - complete := true - err := s.deleteFromDisk(toEvictEl.key, complete) + err := s.deleteFromDisk(toEvictKey, _completeBlob) if err != nil { // TODO - consider whether we want to fail-open by doing `continue` here. return fmt.Errorf("delete from disk: %w", err) } s.evictQueue.Remove(toEvictNode) - delete(s.blobs, toEvictEl.key) - s.releaseSpace(toEvictEl.size) + size := s.blobs[toEvictKey].size + s.releaseSpace(size) + delete(s.blobs, toEvictKey) } s.size += space @@ -241,41 +227,31 @@ func (s *DiskStore) MarkComplete(key string) error { s.mu.Lock() defer s.mu.Unlock() - size, okIncomplete := s.incompleteBlobs[key] - _, okBlob := s.blobs[key] - _, okUnevictable := s.unevictableBlobs[key] - if !okIncomplete && !okBlob && !okUnevictable { + b, ok := s.blobs[key] + if !ok { return os.ErrNotExist } - - if !okIncomplete { + if b.complete { // no-op return nil } - oldPathDir := s.dirPath(key, false) - newPathDir := s.dirPath(key, true) - complete := false - isUnevictable, err := s.checkDiskIfUnevictable(key, complete) - if err != nil { - return fmt.Errorf("check if blob is evictable or not: %w", err) - } - err = os.MkdirAll(filepath.Dir(newPathDir), _defaultFilePerm) + oldPathDir := s.dirPath(key, _incompleteBlob) + newPathDir := s.dirPath(key, _completeBlob) + err := os.MkdirAll(filepath.Dir(newPathDir), _defaultFilePerm) if err != nil { return fmt.Errorf("mkdirall: %w", err) } - err = os.Rename(oldPathDir, newPathDir) // atomic + // TODO - make sure that un-movable metadata is deleted after move + err = os.Rename(oldPathDir, newPathDir) if err != nil { return fmt.Errorf("move dir: %w", err) } - delete(s.incompleteBlobs, key) - if isUnevictable { - s.unevictableBlobs[key] = size - return nil + b.complete = true + if !b.evictionBanned { + node := s.evictQueue.PushBack(key) + b.node = node } - - node := s.evictQueue.PushBack(el{key: key, size: size}) - s.blobs[key] = node return nil } @@ -293,39 +269,22 @@ func (s *DiskStore) Delete(key string) error { s.mu.Lock() defer s.mu.Unlock() - if size, ok := s.unevictableBlobs[key]; ok { - complete := true - err := s.deleteFromDisk(key, complete) - if err != nil { - return fmt.Errorf("delete from disk: %w", err) - } - delete(s.unevictableBlobs, key) - s.releaseSpace(size) - return nil + b, ok := s.blobs[key] + if !ok { + return os.ErrNotExist } - if size, ok := s.incompleteBlobs[key]; ok { - complete := false - err := s.deleteFromDisk(key, complete) - if err != nil { - return fmt.Errorf("delete from disk: %w", err) - } - delete(s.incompleteBlobs, key) - s.releaseSpace(size) - return nil + err := s.deleteFromDisk(key, b.complete) + if err != nil { + return fmt.Errorf("delete from disk: %w", err) } - if node, ok := s.blobs[key]; ok { - complete := true - err := s.deleteFromDisk(key, complete) - if err != nil { - return fmt.Errorf("delete from disk: %w", err) - } - delete(s.blobs, key) - s.evictQueue.Remove(node) - s.releaseSpace(node.Value.(el).size) - return nil + if b.node != nil { + s.evictQueue.Remove(b.node) + b.node = nil } + delete(s.blobs, key) + s.releaseSpace(b.size) - return os.ErrNotExist + return nil } // List returns the blobs' keys. @@ -333,25 +292,12 @@ func (s *DiskStore) List(ignoreIncomplete bool) []string { s.mu.RLock() defer s.mu.RUnlock() - l := len(s.blobs) + len(s.unevictableBlobs) - if !ignoreIncomplete { - l += len(s.incompleteBlobs) - } - res := make([]string, l) - i := 0 - for key := range s.blobs { - res[i] = key - i++ - } - for key := range s.unevictableBlobs { - res[i] = key - i++ - } - if !ignoreIncomplete { - for key := range s.incompleteBlobs { - res[i] = key - i++ + res := make([]string, 0, len(s.blobs)) + for key, b := range s.blobs { + if ignoreIncomplete && !b.complete { + continue } + res = append(res, key) } return res } @@ -362,20 +308,16 @@ func (s *DiskStore) BanEviction(key string) error { s.mu.Lock() defer s.mu.Unlock() - node, okBlob := s.blobs[key] - _, okUnevictable := s.unevictableBlobs[key] - _, okIncomplete := s.incompleteBlobs[key] - - if !okBlob && !okUnevictable && !okIncomplete { + b, ok := s.blobs[key] + if !ok { return os.ErrNotExist } - if okUnevictable { + if b.evictionBanned { // no-op return nil } - complete := okBlob - flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) + flagBlobPath := s.sidecarFilePath(key, b.complete, _evictionBannedFileName) // We persist the ban as a flag file on disk, such that after // a system crash, we can recover the ban. f, err := os.OpenFile(flagBlobPath, os.O_RDONLY|os.O_CREATE, _defaultFilePerm) @@ -384,14 +326,11 @@ func (s *DiskStore) BanEviction(key string) error { } closers.Close(f) - if okIncomplete { - return nil + b.evictionBanned = true + if b.complete { + s.evictQueue.Remove(b.node) + b.node = nil } - - nodeEl := node.Value.(el) - delete(s.blobs, key) - s.evictQueue.Remove(node) - s.unevictableBlobs[key] = nodeEl.size return nil } @@ -400,43 +339,26 @@ func (s *DiskStore) UnbanEviction(key string) error { s.mu.Lock() defer s.mu.Unlock() - _, okBlob := s.blobs[key] - sizeUnevictable, okUnevictable := s.unevictableBlobs[key] - _, okIncomplete := s.incompleteBlobs[key] - - if !okBlob && !okUnevictable && !okIncomplete { + b, ok := s.blobs[key] + if !ok { return os.ErrNotExist } - - if okBlob { + if !b.evictionBanned { // no-op return nil } - complete := okUnevictable - isUnevictable, err := s.checkDiskIfUnevictable(key, complete) - if err != nil { - return fmt.Errorf("check on disk if file is evictable: %w", err) - } - if !isUnevictable { - // no-op - if okUnevictable { - // TODO - log invariant violation - } - return nil - } - flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) - err = os.Remove(flagBlobPath) + flagBlobPath := s.sidecarFilePath(key, b.complete, _evictionBannedFileName) + err := os.Remove(flagBlobPath) if err != nil { return fmt.Errorf("remove file that flags eviction as banned: %w", err) } - if okIncomplete { - return nil - } - delete(s.unevictableBlobs, key) - node := s.evictQueue.PushBack(el{key: key, size: sizeUnevictable}) - s.blobs[key] = node + b.evictionBanned = false + if b.complete { + node := s.evictQueue.PushBack(key) + b.node = node + } return nil } @@ -445,10 +367,8 @@ func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { s.mu.Lock() defer s.mu.Unlock() - _, okBlob := s.blobs[key] - _, okUnevictable := s.unevictableBlobs[key] - _, okIncomplete := s.incompleteBlobs[key] - if !okBlob && !okUnevictable && !okIncomplete { + b, ok := s.blobs[key] + if !ok { return os.ErrNotExist } @@ -456,8 +376,7 @@ func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { if err != nil { return fmt.Errorf("serialize metadata: %w", err) } - complete := !okIncomplete - mdFilePath := s.sidecarFilePath(key, complete, md.GetSuffix()) + mdFilePath := s.sidecarFilePath(key, b.complete, md.GetSuffix()) // We use a tmp file to ensure atomicity. tmpFilePath := mdFilePath + "-tmp" tmpFile, err := os.OpenFile(tmpFilePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, _defaultFilePerm) @@ -484,15 +403,12 @@ func (s *DiskStore) GetMetadata(key string, md metadata.Metadata, ignoreIncomple s.mu.RLock() defer s.mu.RUnlock() - _, okBlob := s.blobs[key] - _, okUnevictable := s.unevictableBlobs[key] - _, okIncomplete := s.incompleteBlobs[key] - if !okBlob && !okUnevictable && (ignoreIncomplete || !okIncomplete) { + b, ok := s.blobs[key] + if !ok || (ignoreIncomplete && !b.complete) { return false, os.ErrNotExist } - complete := !okIncomplete - mdFilePath := s.sidecarFilePath(key, complete, md.GetSuffix()) + mdFilePath := s.sidecarFilePath(key, b.complete, md.GetSuffix()) mdFile, err := os.OpenFile(mdFilePath, os.O_RDONLY, _defaultFilePerm) if errors.Is(err, os.ErrNotExist) { return false, nil @@ -514,14 +430,11 @@ func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { s.mu.Lock() defer s.mu.Unlock() - _, okBlob := s.blobs[key] - _, okUnevictable := s.unevictableBlobs[key] - _, okIncomplete := s.incompleteBlobs[key] - if !okBlob && !okUnevictable && !okIncomplete { + b, ok := s.blobs[key] + if !ok { return os.ErrNotExist } - complete := !okIncomplete - mdFilePath := s.sidecarFilePath(key, complete, md.GetSuffix()) + mdFilePath := s.sidecarFilePath(key, b.complete, md.GetSuffix()) err := os.Remove(mdFilePath) if errors.Is(err, os.ErrNotExist) { // no-op @@ -540,8 +453,8 @@ func (s *DiskStore) evictionOrder() []string { evictionOrder := make([]string, 0) for curr := s.evictQueue.Front(); curr != nil; curr = curr.Next() { - currEl := curr.Value.(el) - evictionOrder = append(evictionOrder, currEl.key) + currKey := curr.Value.(string) + evictionOrder = append(evictionOrder, currKey) } return evictionOrder } diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index 59463bd2f..bf2252ece 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -19,12 +19,6 @@ import ( "github.com/uber/kraken/utils/memsize" ) -const ( - // TODO - consider exporting these constants for clients to consume to avoid naked params - _dontIgnoreIncompleteFiles = false - _ignoreIncompleteFiles = true -) - func newTestStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) @@ -38,8 +32,7 @@ func newTestStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string func newTestFile(t *testing.T, store *DiskStore, size uint64) (f FileReadWriter, key string) { require := require.New(t) key = core.DigestFixture().Hex() - var err error - f, err = store.Create(key, size) + f, err := store.Create(key, size) require.NoError(err) return f, key } @@ -67,48 +60,49 @@ func TestDiskStore(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB) - digests := []core.Digest{} + keys := []string{} for i := range 10 { - digest := core.DigestFixture() - digests = append(digests, digest) - f, err := store.Create(digest.Hex(), memsize.KB) + key := core.DigestFixture().Hex() + keys = append(keys, key) + f, err := store.Create(key, memsize.KB) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) require.NoError(err) data := make([]byte, memsize.KB) for k := range data { data[k] = byte(i + 1) } - n, err := io.Copy(f, bytes.NewReader(make([]byte, memsize.KB))) - require.Equal(n, int64(memsize.KB)) + n, err := io.Copy(f, bytes.NewReader(data)) + require.Equal(int64(memsize.KB), n) require.NoError(err) } - require.Equal(store.size, 10*memsize.KB) + require.Equal(10*memsize.KB, store.size) - digest := core.DigestFixture() - writer, err := store.Create(digest.Hex(), memsize.B) + f, err := store.Create(core.DigestFixture().Hex(), memsize.B) require.EqualError(err, "reserve space: cannot evict enough, the unevictable/incomplete blobs are using up all the space") - require.Nil(writer) + require.Nil(f) - reader, err := store.Open(digests[0].Hex(), true) - require.Equal(err, os.ErrNotExist) - require.Nil(reader) + f, err = store.Open(keys[0], IgnoreIncompleteBlobs) + require.ErrorIs(err, os.ErrNotExist) + require.Nil(f) - require.NoError(store.MarkComplete(digests[0].Hex())) - reader, err = store.Open(digests[0].Hex(), true) + require.NoError(store.MarkComplete(keys[0])) + f, err = store.Open(keys[0], IgnoreIncompleteBlobs) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) require.NoError(err) wantData := make([]byte, memsize.KB) for k := range wantData { wantData[k] = byte(1) } - require.NoError(iotest.TestReader(reader, make([]byte, memsize.KB))) + require.NoError(iotest.TestReader(f, wantData)) // now test that LRU logic works - make sure that the 1 that is complete gets evicted. - digest = core.DigestFixture() - writer, err = store.Create(digest.Hex(), memsize.KB) + f, err = store.Create(core.DigestFixture().Hex(), memsize.KB) require.NoError(err) - reader, err = store.Open(digests[0].Hex(), true) - require.Equal(err, os.ErrNotExist) - require.Nil(reader) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + f, err = store.Open(keys[0], IgnoreIncompleteBlobs) + require.ErrorIs(err, os.ErrNotExist) + require.Nil(f) } func TestEviction(t *testing.T) { @@ -145,7 +139,7 @@ func TestEviction(t *testing.T) { f, fKey := newTestFile(t, store, 4*memsize.KB) require.NoError(f.Close()) require.NoError(store.MarkComplete(fKey)) - keys := store.List(_dontIgnoreIncompleteFiles) + keys := store.List(CheckIncompleteBlobs) require.NotContains(keys, cKey) // new size == 23KB == 24KB - 5KB (c) + 4KB (f) require.Equal(23*memsize.KB, store.size) @@ -165,7 +159,7 @@ func TestEviction(t *testing.T) { // size == 24KB == 24KB + 15KB (h) - 5KB (b) - 10KB (a) require.Equal(24*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) - keys = store.List(_dontIgnoreIncompleteFiles) + keys = store.List(CheckIncompleteBlobs) require.NotContains(keys, bKey) require.NotContains(keys, aKey) @@ -174,14 +168,14 @@ func TestEviction(t *testing.T) { require.NoError(store.UnbanEviction(eKey)) // eviction order (left-most is next to evict): f(4KB), g(1KB), h(15KB), e(1KB); d(3KB) is unevictable // we open g to change the order to f, h, e, g - g, err = store.Open(gKey, _dontIgnoreIncompleteFiles) + g, err = store.Open(gKey, CheckIncompleteBlobs) require.NoError(err) require.NoError(g.Close()) i, iKey := newTestFile(t, store, 5*memsize.KB) require.NoError(store.MarkComplete(iKey)) require.NoError(i.Close()) - keys = store.List(_dontIgnoreIncompleteFiles) + keys = store.List(CheckIncompleteBlobs) require.NotContains(keys, fKey) require.Equal(25*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -190,7 +184,7 @@ func TestEviction(t *testing.T) { j, jKey := newTestFile(t, store, 14*memsize.KB) require.NoError(j.Close()) require.NoError(store.MarkComplete(jKey)) - keys = store.List(_dontIgnoreIncompleteFiles) + keys = store.List(CheckIncompleteBlobs) require.NotContains(keys, hKey) require.Equal(24*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -199,7 +193,7 @@ func TestEviction(t *testing.T) { k, kKey := newTestFile(t, store, 2*memsize.KB) require.NoError(k.Close()) require.NoError(store.MarkComplete(kKey)) - keys = store.List(_dontIgnoreIncompleteFiles) + keys = store.List(CheckIncompleteBlobs) require.NotContains(keys, eKey) require.Equal(25*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -208,7 +202,7 @@ func TestEviction(t *testing.T) { l, lKey := newTestFile(t, store, 1*memsize.KB) require.NoError(store.MarkComplete(lKey)) require.NoError(l.Close()) - keys = store.List(_dontIgnoreIncompleteFiles) + keys = store.List(CheckIncompleteBlobs) require.NotContains(keys, gKey) require.Equal(25*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -234,36 +228,59 @@ func TestParallelAccessToSingleFile(t *testing.T) { var wg sync.WaitGroup wg.Add(5) + type res struct { + written, read []byte + err error + } + + results := make([]res, 5) + for idx := range 5 { go func(idx int64) { defer wg.Done() - ignoreIncomplete := false - f, err := store.Open(key, ignoreIncomplete) - require.NoError(err) + f, err := store.Open(key, CheckIncompleteBlobs) + if err != nil { + results[idx].err = err + return + } pos := idx * 10 writtenData := make([]byte, 10) for k := range writtenData { writtenData[k] = byte(idx) } - n, err := f.WriteAt(writtenData, pos) - require.NoError(err) - require.Equal(10, n) + _, err = f.WriteAt(writtenData, pos) + if err != nil { + results[idx].err = err + return + } - defer func() { require.NoError(f.Close()) }() readData := make([]byte, 10) - n, err = f.ReadAt(readData, pos) - require.NoError(err) - require.Equal(10, n) - require.Equal(writtenData, readData) + _, err = f.ReadAt(readData, pos) + if err != nil { + results[idx].err = err + return + } + err = f.Close() + if err != nil { + results[idx].err = err + return + } + + results[idx].read = readData + results[idx].written = writtenData }(int64(idx)) } wg.Wait() + for idx := range 5 { + require.NoError(results[idx].err) + require.Equal(results[idx].written, results[idx].written) + } + require.NoError(store.MarkComplete(key)) - ignoreIncomplete := true - f, err = store.Open(key, ignoreIncomplete) + f, err = store.Open(key, IgnoreIncompleteBlobs) require.NoError(err) defer func() { require.NoError(f.Close()) }() @@ -287,15 +304,13 @@ func TestOpenedFileAccessibleAfterMarkedComplete(t *testing.T) { require.NoError(err) require.NoError(f.Close()) - ignoreIncomplete := false - incompleteFile, err := store.Open(key, ignoreIncomplete) + incompleteFile, err := store.Open(key, CheckIncompleteBlobs) require.NoError(err) defer func() { require.NoError(incompleteFile.Close()) }() require.NoError(store.MarkComplete(key)) - ignoreIncomplete = true - completeFile, err := store.Open(key, ignoreIncomplete) + completeFile, err := store.Open(key, IgnoreIncompleteBlobs) require.NoError(err) defer func() { require.NoError(completeFile.Close()) }() @@ -327,8 +342,9 @@ func TestDelete(t *testing.T) { require.NoError(f.Close()) require.NoError(store.Delete(key)) - require.Empty(store.List(false)) + require.Empty(store.List(CheckIncompleteBlobs)) require.Equal(uint64(0), store.size) + require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("incomplete, unevictable blob", func(t *testing.T) { require := require.New(t) @@ -343,8 +359,9 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) - require.Empty(store.List(false)) + require.Empty(store.List(CheckIncompleteBlobs)) require.Equal(uint64(0), store.size) + require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("complete blob", func(t *testing.T) { require := require.New(t) @@ -359,8 +376,9 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) - require.Empty(store.List(false)) + require.Empty(store.List(CheckIncompleteBlobs)) require.Equal(uint64(0), store.size) + require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("complete, unevictable blob", func(t *testing.T) { require := require.New(t) @@ -376,8 +394,9 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) - require.Empty(store.List(false)) + require.Empty(store.List(CheckIncompleteBlobs)) require.Equal(uint64(0), store.size) + require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("not found", func(t *testing.T) { require := require.New(t) @@ -402,9 +421,9 @@ func TestMarkComplete(t *testing.T) { require.NoError(store.MarkComplete(key)) - require.Equal([]string{key}, store.List(true)) + require.Equal([]string{key}, store.List(IgnoreIncompleteBlobs)) require.Equal(uint64(100), store.size) - _, err = store.Open(key, true) + _, err = store.Open(key, IgnoreIncompleteBlobs) require.NoError(err) }) t.Run("incomplete blob with forbidden eviction", func(t *testing.T) { @@ -420,9 +439,9 @@ func TestMarkComplete(t *testing.T) { require.NoError(store.MarkComplete(key)) - require.Equal([]string{key}, store.List(true)) + require.Equal([]string{key}, store.List(IgnoreIncompleteBlobs)) require.Equal(uint64(100), store.size) - _, err = store.Open(key, true) + _, err = store.Open(key, IgnoreIncompleteBlobs) require.NoError(err) }) t.Run("already complete blob", func(t *testing.T) { @@ -474,15 +493,16 @@ func TestStat(t *testing.T) { require.NoError(err) require.NoError(store.MarkComplete(key)) - fInfo, err := store.Stat(key, true) + fInfo, err := store.Stat(key, IgnoreIncompleteBlobs) require.NoError(err) - _, err = store.Stat(key, false) + _, err = store.Stat(key, CheckIncompleteBlobs) require.NoError(err) require.False(fInfo.IsDir()) require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) require.Equal(_blobFileName, fInfo.Name()) require.Equal(int64(10), fInfo.Size()) + require.Equal(fs.FileMode(0755), fInfo.Mode()) }) t.Run("complete, unevictable blob", func(t *testing.T) { require := require.New(t) @@ -496,15 +516,16 @@ func TestStat(t *testing.T) { require.NoError(store.MarkComplete(key)) require.NoError(store.BanEviction(key)) - fInfo, err := store.Stat(key, true) + fInfo, err := store.Stat(key, IgnoreIncompleteBlobs) require.NoError(err) - _, err = store.Stat(key, false) + _, err = store.Stat(key, CheckIncompleteBlobs) require.NoError(err) require.False(fInfo.IsDir()) require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) require.Equal(_blobFileName, fInfo.Name()) require.Equal(int64(10), fInfo.Size()) + require.Equal(fs.FileMode(0755), fInfo.Mode()) }) t.Run("incomplete blob", func(t *testing.T) { require := require.New(t) @@ -516,15 +537,16 @@ func TestStat(t *testing.T) { _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) require.NoError(err) - _, err = store.Stat(key, true) + _, err = store.Stat(key, IgnoreIncompleteBlobs) require.Equal(os.ErrNotExist, err) - fInfo, err := store.Stat(key, false) + fInfo, err := store.Stat(key, CheckIncompleteBlobs) require.NoError(err) require.False(fInfo.IsDir()) require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) require.Equal(_blobFileName, fInfo.Name()) require.Equal(int64(10), fInfo.Size()) + require.Equal(fs.FileMode(0755), fInfo.Mode()) }) t.Run("incomplete, unevictable blob", func(t *testing.T) { @@ -538,24 +560,25 @@ func TestStat(t *testing.T) { require.NoError(err) require.NoError(store.BanEviction(key)) - _, err = store.Stat(key, true) + _, err = store.Stat(key, IgnoreIncompleteBlobs) require.Equal(os.ErrNotExist, err) - fInfo, err := store.Stat(key, false) + fInfo, err := store.Stat(key, CheckIncompleteBlobs) require.NoError(err) require.False(fInfo.IsDir()) require.WithinDuration(time.Now(), fInfo.ModTime(), 500*time.Millisecond) require.Equal(_blobFileName, fInfo.Name()) require.Equal(int64(10), fInfo.Size()) + require.Equal(fs.FileMode(0755), fInfo.Mode()) }) t.Run("non-existent blob", func(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB) key := core.DigestFixture().Hex() - _, err := store.Stat(key, true) + _, err := store.Stat(key, IgnoreIncompleteBlobs) require.Equal(os.ErrNotExist, err) - _, err = store.Stat(key, false) + _, err = store.Stat(key, CheckIncompleteBlobs) require.Equal(os.ErrNotExist, err) }) } @@ -564,8 +587,8 @@ func TestList(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB) - require.Empty(store.List(false)) - require.Empty(store.List(true)) + require.Empty(store.List(CheckIncompleteBlobs)) + require.Empty(store.List(IgnoreIncompleteBlobs)) incompleteBlobKey := core.DigestFixture().Hex() f, err := store.Create(incompleteBlobKey, 10*memsize.B) @@ -594,8 +617,13 @@ func TestList(t *testing.T) { defer func() { require.NoError(f.Close()) }() require.NoError(store.MarkComplete(unevictableCompleteBlobKey)) - require.Equal([]string{completeBlobKey, unevictableCompleteBlobKey}, store.List(true)) - require.Equal([]string{completeBlobKey, unevictableCompleteBlobKey, incompleteBlobKey, unevictableIncompleteBlobKey}, store.List(false)) + wantRes := []string{completeBlobKey, unevictableCompleteBlobKey} + res := store.List(IgnoreIncompleteBlobs) + require.ElementsMatch(wantRes, res) + + wantRes = []string{incompleteBlobKey, completeBlobKey, unevictableCompleteBlobKey, unevictableIncompleteBlobKey} + res = store.List(CheckIncompleteBlobs) + require.ElementsMatch(wantRes, res) } func TestMetadata(t *testing.T) { @@ -610,20 +638,20 @@ func TestMetadata(t *testing.T) { mdStruct := core.MetaInfoFixture() writtenMd := metadata.NewTorrentMeta(mdStruct) err = store.SetMetadata(key, writtenMd) - // ensure metadata is not included in LRU eviction calculation. + // asserts metadata is not included in LRU eviction calculation. require.NoError(err) var readMd metadata.TorrentMeta - ok, err := store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + ok, err := store.GetMetadata(key, &readMd, CheckIncompleteBlobs) require.NoError(err) require.True(ok) - require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) require.NoError(store.DeleteMetadata(key, &readMd)) - ok, err = store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + ok, err = store.GetMetadata(key, &readMd, CheckIncompleteBlobs) require.NoError(err) require.False(ok) - mdFilePath := store.sidecarFilePath(key, false, readMd.GetSuffix()) + mdFilePath := store.sidecarFilePath(key, _incompleteBlob, readMd.GetSuffix()) // ensure the metadata file is deleted from disk _, err = os.Stat(mdFilePath) require.True(errors.Is(err, os.ErrNotExist)) @@ -631,7 +659,7 @@ func TestMetadata(t *testing.T) { require.NoError(store.DeleteMetadata(key, &readMd)) }) - t.Run("non-existant blob", func(t *testing.T) { + t.Run("non-existent blob", func(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB) nonExistentKey := core.DigestFixture().Hex() @@ -641,7 +669,7 @@ func TestMetadata(t *testing.T) { err := store.SetMetadata(nonExistentKey, md) require.Equal(os.ErrNotExist, err) - ok, err := store.GetMetadata(nonExistentKey, md, _dontIgnoreIncompleteFiles) + ok, err := store.GetMetadata(nonExistentKey, md, CheckIncompleteBlobs) require.Equal(os.ErrNotExist, err) require.False(ok) @@ -662,33 +690,33 @@ func TestMetadata(t *testing.T) { require.NoError(store.SetMetadata(key, writtenMd)) var readMd metadata.TorrentMeta - ok, err := store.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + ok, err := store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) require.Equal(os.ErrNotExist, err) // incomplete files are ignored require.False(ok) - ok, err = store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + ok, err = store.GetMetadata(key, &readMd, CheckIncompleteBlobs) require.NoError(err) require.True(ok) - require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) - // Repeat the tests above for an evictable file + // Repeat the tests above for an unevictable file require.NoError(store.BanEviction(key)) - ok, err = store.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + ok, err = store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) require.Equal(os.ErrNotExist, err) // incomplete files are ignored require.False(ok) - ok, err = store.GetMetadata(key, &readMd, _dontIgnoreIncompleteFiles) + ok, err = store.GetMetadata(key, &readMd, CheckIncompleteBlobs) require.NoError(err) require.True(ok) - require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) require.NoError(store.MarkComplete(key)) - ok, err = store.GetMetadata(key, &readMd, _ignoreIncompleteFiles) + ok, err = store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) require.NoError(err) require.True(ok) - require.Equal(readMd.MetaInfo, writtenMd.MetaInfo) + require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) }) t.Run("metadata fully gone after blob is evicted", func(t *testing.T) { require := require.New(t) @@ -701,9 +729,8 @@ func TestMetadata(t *testing.T) { md := metadata.NewTorrentMeta(core.MetaInfoFixture()) err = store.SetMetadata(keyA, md) - store.SetMetadata(keyA, md) - complete := true - mdFilePath := store.sidecarFilePath(keyA, complete, md.GetSuffix()) + require.NoError(err) + mdFilePath := store.sidecarFilePath(keyA, _completeBlob, md.GetSuffix()) _, err = os.Stat(mdFilePath) require.NoError(err) @@ -712,7 +739,7 @@ func TestMetadata(t *testing.T) { require.NoError(err) defer func() { require.NoError(fB.Close()) }() - ok, err := store.GetMetadata(keyA, md, _dontIgnoreIncompleteFiles) + ok, err := store.GetMetadata(keyA, md, CheckIncompleteBlobs) require.Equal(os.ErrNotExist, err) require.False(ok) _, err = os.Stat(mdFilePath) diff --git a/lib/store/pather.go b/lib/store/pather.go index 7dccb199f..fed4065fa 100644 --- a/lib/store/pather.go +++ b/lib/store/pather.go @@ -3,6 +3,8 @@ package store import "path/filepath" const ( + // _defaultShardIDLength is the number of bytes of file digest to be used for shard ID. + // For every byte (2 HEX char), one more level of directories will be created. _defaultShardIDLength = 2 _incompleteSubDir = "incomplete" _completeSubDir = "complete" diff --git a/lib/store/pather_test.go b/lib/store/pather_test.go index 1a1545e78..8ea7cd80e 100644 --- a/lib/store/pather_test.go +++ b/lib/store/pather_test.go @@ -13,27 +13,31 @@ func TestPather(t *testing.T) { require := require.New(t) store, rootDir := newTestStore(t, 10*memsize.KB) md := metadata.NewTorrentMeta(core.MetaInfoFixture()) - key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - complete := false - dirPath := store.dirPath(key, complete) - wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - require.Equal(wantDirPath, dirPath) - blobPath := store.blobPath(key, complete) - wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" - require.Equal(wantBlobPath, blobPath) - sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) - wantSidecarFilePath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" - require.Equal(wantSidecarFilePath, sidecarFilePath) - complete = true - dirPath = store.dirPath(key, complete) - wantDirPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - require.Equal(wantDirPath, dirPath) - blobPath = store.blobPath(key, complete) - wantBlobPath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" - require.Equal(wantBlobPath, blobPath) - sidecarFilePath = store.sidecarFilePath(key, complete, md.GetSuffix()) - wantSidecarFilePath = rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" - require.Equal(wantSidecarFilePath, sidecarFilePath) + t.Run("incomplete entry", func(t *testing.T) { + complete := false + dirPath := store.dirPath(key, complete) + wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + require.Equal(wantDirPath, dirPath) + blobPath := store.blobPath(key, complete) + wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" + require.Equal(wantBlobPath, blobPath) + sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) + wantSidecarFilePath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" + require.Equal(wantSidecarFilePath, sidecarFilePath) + }) + + t.Run("complete entry", func(t *testing.T) { + complete := true + dirPath := store.dirPath(key, complete) + wantDirPath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" + require.Equal(wantDirPath, dirPath) + blobPath := store.blobPath(key, complete) + wantBlobPath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" + require.Equal(wantBlobPath, blobPath) + sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) + wantSidecarFilePath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" + require.Equal(wantSidecarFilePath, sidecarFilePath) + }) } From c427171ca160c2b1054386d1d8875e94e6329c8d Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Mon, 20 Jul 2026 20:09:27 +0200 Subject: [PATCH 06/10] feat(DiskStore): support rebooting of incomplete blobs 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 --- lib/store/crash_recovery.go | 133 +++++++++++++++------ lib/store/crash_recovery_test.go | 197 +++++++++++++++++++++++++++---- lib/store/disk_store.go | 65 +++++++--- lib/store/disk_store_test.go | 52 ++++---- lib/store/file.go | 12 +- lib/store/pather.go | 4 +- lib/store/pather_test.go | 19 +-- 7 files changed, 359 insertions(+), 123 deletions(-) diff --git a/lib/store/crash_recovery.go b/lib/store/crash_recovery.go index 15b1625e9..e5e23230e 100644 --- a/lib/store/crash_recovery.go +++ b/lib/store/crash_recovery.go @@ -4,13 +4,16 @@ import ( "container/list" "errors" "fmt" + "io" "io/fs" "os" "path/filepath" "slices" + "strconv" "strings" "time" + "github.com/uber/kraken/utils/closers" "go.uber.org/zap" ) @@ -21,58 +24,68 @@ type blobState struct { size uint64 mTime time.Time evictable bool + complete bool } -func rebootPersistedStateAfterCrash(capacityBytes uint64, rootDir string, log *zap.SugaredLogger) (*DiskStore, error) { +func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, log *zap.SugaredLogger) (*DiskStore, error) { completeDirPath, incompleteDirPath := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) - - // We remove incomplete entries, as we expect the processes/clients who were writing to - // these entries to also have crashed, which would leak the files on disk. - err := os.RemoveAll(incompleteDirPath) - if err != nil { - return nil, fmt.Errorf("remove incomplete blobs left from a previous service run: %w", err) + if !rebootIncompleteBlobs { + err := os.RemoveAll(incompleteDirPath) + if err != nil { + return nil, fmt.Errorf("remove incomplete blobs left from a previous service run: %w", err) + } } + keys, err := rebootKeys(completeDirPath) if err != nil { return nil, err } - pather := newPather(rootDir) + numCompleteBlobs := len(keys) + if rebootIncompleteBlobs { + incompleteKeys, err := rebootKeys(incompleteDirPath) + if err != nil { + return nil, err + } + keys = append(keys, incompleteKeys...) + } - evictableEntries := make([]*blobState, 0) - unevictableEntries := make([]*blobState, 0) - for _, key := range keys { - bState, ok, err := rebootBlobState(key, pather) + pather := newPather(rootDir) + completeEvictableEntries := make([]*blobState, 0) + otherEntries := make([]*blobState, 0) + for i, key := range keys { + complete := i < numCompleteBlobs + bState, ok, err := rebootBlobState(key, complete, pather) if err != nil { return nil, err } if !ok { continue } - if bState.evictable { - evictableEntries = append(evictableEntries, bState) + if bState.complete && bState.evictable { + completeEvictableEntries = append(completeEvictableEntries, bState) } else { - unevictableEntries = append(unevictableEntries, bState) + otherEntries = append(otherEntries, bState) } } storeSize := uint64(0) blobs := make(map[string]*blob, 0) - for _, bState := range unevictableEntries { + for _, bState := range otherEntries { blobs[bState.key] = &blob{ size: bState.size, - complete: true, + complete: bState.complete, + evictionBanned: !bState.evictable, node: nil, - evictionBanned: true, } storeSize += bState.size } - slices.SortFunc(evictableEntries, func(left, right *blobState) int { + slices.SortFunc(completeEvictableEntries, func(left, right *blobState) int { // left-most is oldest, i.e. next-to-evict. return left.mTime.Compare(right.mTime) }) evictQueue := list.New() - for _, bState := range evictableEntries { + for _, bState := range completeEvictableEntries { node := evictQueue.PushBack(bState.key) blobs[bState.key] = &blob{ size: bState.size, @@ -84,12 +97,13 @@ func rebootPersistedStateAfterCrash(capacityBytes uint64, rootDir string, log *z } store := &DiskStore{ - blobs: blobs, - evictQueue: evictQueue, - capacity: capacityBytes, - pather: pather, - size: storeSize, - log: log, + blobs: blobs, + evictQueue: evictQueue, + capacity: capacityBytes, + pather: pather, + size: storeSize, + rebootIncompleteBlobs: rebootIncompleteBlobs, + log: log, } if store.size > store.capacity { @@ -98,43 +112,84 @@ func rebootPersistedStateAfterCrash(capacityBytes uint64, rootDir string, log *z err = store.reserveSpace(0) if err != nil { log.With("error", err).Error("DiskStore size exceeds its capacity after service reboot. Evicting blobs from disk did not work to reduce size within capacity.") - return nil, fmt.Errorf("remove blobs to reduce store size within configured capacity") + return nil, fmt.Errorf("remove blobs to reduce store size within configured capacity: %w", err) } evictedBytes := prevSize - store.size - log.With("evicted_bytes", evictedBytes).Warn("DiskStore size exceeds its capacity after service reboot. Successfully evicted blobs to reduce size within capacity.") + log.With("evicted_bytes", evictedBytes).Warn("DiskStore size exceeded its capacity after service reboot. Successfully evicted blobs to reduce size within capacity.") } return store, nil } -func rebootBlobState(key string, pather *pather) (res *blobState, ok bool, err error) { - complete := true +func rebootBlobState(key string, complete bool, pather *pather) (res *blobState, ok bool, err error) { blobPath := pather.blobPath(key, complete) fInfo, err := os.Stat(blobPath) if errors.Is(err, os.ErrNotExist) { - // For some reason, the directory for the blob exists but not the blob itself. - return nil, true, nil + // The directory for the blob exists but not the blob itself. + return nil, false, nil } if err != nil { return nil, false, fmt.Errorf("stat blob entry: %w", err) } + flagBlobPath := pather.sidecarFilePath(key, complete, _evictionBannedFileName) isUnevictable, err := exists(flagBlobPath) if err != nil { return nil, false, err } - size := fInfo.Size() + var size uint64 + if complete { + size = uint64(fInfo.Size()) + } else { + size, ok, err = rebootIncompleteBlobSize(key, pather) + if err != nil { + return nil, false, fmt.Errorf("get incomplete blob size from sidecar file: %w", err) + } + if !ok { + return nil, false, nil + } + } mTime := fInfo.ModTime() return &blobState{ key: key, - size: uint64(size), + size: size, mTime: mTime, evictable: !isUnevictable, + complete: complete, }, true, nil } -func rebootKeys(completeDirPath string) ([]string, error) { +func rebootIncompleteBlobSize(key string, pather *pather) (size uint64, ok bool, err error) { + blobSizeFilePath := pather.sidecarFilePath(key, _incompleteBlob, _blobSizeFileName) + blobSizeF, err := os.OpenFile(blobSizeFilePath, os.O_RDONLY, _defaultFilePerm) + if errors.Is(err, os.ErrNotExist) { + // The size metadata file is not present, we fail-open by evicting the blob. + return 0, false, nil + } + if err != nil { + return 0, false, fmt.Errorf("open blob size sidecar file: %w", err) + } + blobSizeData, err := io.ReadAll(blobSizeF) + if err != nil { + return 0, false, fmt.Errorf("read blob size sidecar file: %w", err) + } + blobSize, err := strconv.Atoi(string(blobSizeData)) + if err != nil { + return 0, false, fmt.Errorf("blob size sidecar file is in unexpected format: %w", err) + } + closers.Close(blobSizeF) + return uint64(blobSize), true, nil +} + +func rebootKeys(subDir string) ([]string, error) { keys := make([]string, 0) - err := filepath.WalkDir(completeDirPath, func(path string, entry fs.DirEntry, err error) error { + ok, err := exists(subDir) + if err != nil { + return nil, fmt.Errorf("exists: %w", err) + } + if !ok { + return []string{}, nil + } + err = filepath.WalkDir(subDir, func(path string, entry fs.DirEntry, err error) error { if err != nil { return err } @@ -142,7 +197,7 @@ func rebootKeys(completeDirPath string) ([]string, error) { if !entry.IsDir() { return nil } - relPath, err := filepath.Rel(completeDirPath, path) + relPath, err := filepath.Rel(subDir, path) if err != nil { return err } @@ -156,12 +211,12 @@ func rebootKeys(completeDirPath string) ([]string, error) { return fs.SkipDir }) if err != nil { - return nil, fmt.Errorf("walk through complete dir to collect keys of blobs: %w", err) + return nil, fmt.Errorf("walk through subdir '%v' to reboot blob keys: %w", subDir, err) } return keys, nil } -func existsPersistedState(rootDir string) (ok bool, err error) { +func existsPersistedStore(rootDir string) (ok bool, err error) { completeDir, incompleteDir := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) completeExists, err := exists(completeDir) if err != nil { diff --git a/lib/store/crash_recovery_test.go b/lib/store/crash_recovery_test.go index 9aee117d3..aa99759a0 100644 --- a/lib/store/crash_recovery_test.go +++ b/lib/store/crash_recovery_test.go @@ -14,9 +14,9 @@ import ( ) func TestCrashRecovery(t *testing.T) { - t.Run("complete blobs, evictability, and completeness are recovered", func(t *testing.T) { + t.Run("blobs, evictability, completeness, and size are recovered", func(t *testing.T) { require := require.New(t) - store, rootDir := newTestStore(t, 10*memsize.KB) + store, rootDir := newTestStore(t, 10*memsize.KB, true) completeEvictableF, completeEvictableKey := newTestFile(t, store, 2*memsize.KB) completeEvictableData := fillWithRandomData(t, completeEvictableF, 2*memsize.KB) @@ -31,50 +31,100 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.BanEviction(completeUnevictableKey)) incompleteEvictableF, incompleteEvictableKey := newTestFile(t, store, 2*memsize.KB) - _ = fillWithRandomData(t, incompleteEvictableF, 2*memsize.KB) + incompleteEvictableData := fillWithRandomData(t, incompleteEvictableF, 2*memsize.KB) require.NoError(incompleteEvictableF.Close()) incompleteUnevictableF, incompleteUnevictableKey := newTestFile(t, store, 2*memsize.KB) - _ = fillWithRandomData(t, incompleteUnevictableF, 2*memsize.KB) + incompleteUnevictableData := fillWithRandomData(t, incompleteUnevictableF, 2*memsize.KB) require.NoError(incompleteUnevictableF.Close()) require.NoError(store.BanEviction(incompleteUnevictableKey)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - rebootedStore, err := NewDiskStore(10*memsize.KB, rootDir) + store, err := NewDiskStore(10*memsize.KB, rootDir, true) require.NoError(err) - // Incomplete files get dropped. - _, err = rebootedStore.Stat(incompleteEvictableKey, CheckIncompleteBlobs) + // Incomplete files are recovered (since `rebootIncompleteBlobs` is true). + f, err := store.Open(incompleteEvictableKey, CheckIncompleteBlobs) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err := io.ReadAll(f) + require.NoError(err) + require.Equal(incompleteEvictableData, data) + unevictable, err := store.checkDiskIfUnevictable(completeEvictableKey, _incompleteBlob) + require.NoError(err) + require.False(unevictable) + + f, err = store.Open(incompleteUnevictableKey, CheckIncompleteBlobs) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(incompleteUnevictableData, data) + unevictable, err = store.checkDiskIfUnevictable(completeEvictableKey, _incompleteBlob) + require.NoError(err) + require.False(unevictable) + + // Complete files are always recovered. + f, err = store.Open(completeEvictableKey, IgnoreIncompleteBlobs) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(completeEvictableData, data) + unevictable, err = store.checkDiskIfUnevictable(completeEvictableKey, _completeBlob) + require.NoError(err) + require.False(unevictable) + + f, err = store.Open(completeUnevictableKey, IgnoreIncompleteBlobs) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(completeUnevictableData, data) + unevictable, err = store.checkDiskIfUnevictable(completeUnevictableKey, _completeBlob) + require.NoError(err) + require.True(unevictable) + + // The sizes of both incomplete and complete files are recovered correctly. + require.Equal(8*memsize.KB, store.size) + + // Run the store with `rebootIncompleteBlobs` as false. + store, err = NewDiskStore(10*memsize.KB, rootDir, false) + require.NoError(err) + + // Incomplete files are dropped. + _, err = store.Stat(incompleteEvictableKey, CheckIncompleteBlobs) require.ErrorIs(err, os.ErrNotExist) - _, err = rebootedStore.Stat(incompleteUnevictableKey, CheckIncompleteBlobs) + _, err = store.Stat(incompleteUnevictableKey, CheckIncompleteBlobs) require.ErrorIs(err, os.ErrNotExist) - // Complete files are recovered. - f, err := rebootedStore.Open(completeEvictableKey, IgnoreIncompleteBlobs) + // Complete files are always recovered. + f, err = store.Open(completeEvictableKey, IgnoreIncompleteBlobs) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) - data, err := io.ReadAll(f) + data, err = io.ReadAll(f) require.NoError(err) require.Equal(completeEvictableData, data) - complete := true - unevictable, err := rebootedStore.checkDiskIfUnevictable(completeEvictableKey, complete) + unevictable, err = store.checkDiskIfUnevictable(completeEvictableKey, _completeBlob) require.NoError(err) require.False(unevictable) - f, err = rebootedStore.Open(completeUnevictableKey, IgnoreIncompleteBlobs) + f, err = store.Open(completeUnevictableKey, IgnoreIncompleteBlobs) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) require.NoError(err) require.Equal(completeUnevictableData, data) - unevictable, err = rebootedStore.checkDiskIfUnevictable(completeUnevictableKey, complete) + unevictable, err = store.checkDiskIfUnevictable(completeUnevictableKey, _completeBlob) require.NoError(err) require.True(unevictable) + + require.Equal(4*memsize.KB, store.size) }) - t.Run("metadata is recovered", func(t *testing.T) { + t.Run("metadata is recovered for complete blob", func(t *testing.T) { require := require.New(t) - store, rootDir := newTestStore(t, 10*memsize.KB) + store, rootDir := newTestStore(t, 10*memsize.KB, true) f, key := newTestFile(t, store, 2*memsize.KB) require.NoError(f.Close()) @@ -83,11 +133,31 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.SetMetadata(key, writtenMd)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - rebootedStore, err := NewDiskStore(10*memsize.KB, rootDir) + store, err := NewDiskStore(10*memsize.KB, rootDir, true) require.NoError(err) var readMd metadata.TorrentMeta - ok, err := rebootedStore.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) + ok, err := store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) + require.NoError(err) + require.True(ok) + require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) + }) + + t.Run("metadata is recovered for incomplete blob", func(t *testing.T) { + require := require.New(t) + store, rootDir := newTestStore(t, 10*memsize.KB, true) + + f, key := newTestFile(t, store, 2*memsize.KB) + require.NoError(f.Close()) + writtenMd := metadata.NewTorrentMeta(core.MetaInfoFixture()) + require.NoError(store.SetMetadata(key, writtenMd)) + + // Assume that the application crashes here. The application would restart and call `NewDiskStore`. + store, err := NewDiskStore(10*memsize.KB, rootDir, true) + require.NoError(err) + + var readMd metadata.TorrentMeta + ok, err := store.GetMetadata(key, &readMd, CheckIncompleteBlobs) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) @@ -95,7 +165,7 @@ func TestCrashRecovery(t *testing.T) { t.Run("lru order is approximated and blobs are evicted if store size exceeds capacity", func(t *testing.T) { require := require.New(t) - store, rootDir := newTestStore(t, 10*memsize.KB) + store, rootDir := newTestStore(t, 10*memsize.KB, false) aF, aKey := newTestFile(t, store, 2*memsize.KB) _ = fillWithRandomData(t, aF, 2*memsize.KB) @@ -131,17 +201,17 @@ func TestCrashRecovery(t *testing.T) { require.Equal([]string{cKey, eKey, dKey}, evictionOrderBeforeCrash) // a is unevictable and b is incomplete // Assume that the application restarts here. - rebootedStore, err := NewDiskStore(10*memsize.KB, rootDir) + store, err = NewDiskStore(10*memsize.KB, rootDir, false) require.NoError(err) // LRU order is approximated, but not exact. - rebootedEvictionOrder := rebootedStore.evictionOrder() + rebootedEvictionOrder := store.evictionOrder() require.NotEqual(evictionOrderBeforeCrash, rebootedEvictionOrder) wantEvictionOrder := []string{cKey, dKey, eKey} require.Equal(wantEvictionOrder, rebootedEvictionOrder) // Assume we redeploy the service with a smaller capacity for the disk store: - rebootedSmallerStore, err := NewDiskStore(6*memsize.KB, rootDir) + rebootedSmallerStore, err := NewDiskStore(6*memsize.KB, rootDir, false) require.NoError(err) // since 10KB of blobs are in store, `c` gets evicted to put the store back within its capacity. @@ -154,13 +224,92 @@ func TestCrashRecovery(t *testing.T) { }) } +func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { + require := require.New(t) + store, rootDir := newTestStore(t, 10*memsize.KB, true) + + f, key := newTestFile(t, store, 4*memsize.KB) + firstData := fillWithRandomData(t, f, 2*memsize.KB) + + // First crash. + store, err := NewDiskStore(10*memsize.KB, rootDir, true) + require.NoError(err) + + f, err = store.Open(key, CheckIncompleteBlobs) + require.NoError(err) + + // Write 5KB in total, while reporting only 4KB. + secondData := make([]byte, 3*memsize.KB) + _, err = rand.Read(secondData) + require.NoError(err) + _, err = f.WriteAt(secondData, int64(2*memsize.KB)) + require.NoError(err) + + // Second crash. + store, err = NewDiskStore(10*memsize.KB, rootDir, true) + require.NoError(err) + + wantData := make([]byte, 5*memsize.KB) + copy(wantData, firstData) + copy(wantData[2*memsize.KB:], secondData) + + f, err = store.Open(key, CheckIncompleteBlobs) + require.NoError(err) + data, err := io.ReadAll(f) + require.NoError(err) + require.Equal(wantData, data) + // The declared size is recovered correctly both times, without leaking or duplicating reservation. + require.Equal(4*memsize.KB, store.size) + require.NoError(store.MarkComplete(key)) + + // Third crash. + store, err = NewDiskStore(10*memsize.KB, rootDir, true) + // now that the blob is complete, its actual size should be rebooted through stat, instead of trusting the _size sidecar file. + require.Equal(5*memsize.KB, store.size) + f, err = store.Open(key, IgnoreIncompleteBlobs) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(wantData, data) +} + func TestStoreWorksWhenFileSizeNotCorrect(t *testing.T) { // Verify that the store works correctly when the reserved size for a file (the one passed by the client in Create) is different // than its actual size. The store is expected to consistently use EITHER the client-given size OR the actual size of files, but not both. // If we mix them, this could break the eviction logic - imagine the user uploads a 2GB size but reports it as 1.9GB. Eviction works correctly // as long as we reserve 2GB upon upload to store and release 2GB upon deletion/eviction from store. BUT if we reserve 2GB and free 1.9GB // or vice-versa, it could lead to over/under-reservation. - t.Skip("TODO") + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB, false) + + // Declares 8KB but only writes 1KB. + underreportedF, underreportedKey := newTestFile(t, store, 8*memsize.KB) + _ = fillWithRandomData(t, underreportedF, 1*memsize.KB) + require.NoError(underreportedF.Close()) + require.NoError(store.MarkComplete(underreportedKey)) + require.NoError(store.BanEviction(underreportedKey)) + require.Equal(8*memsize.KB, store.size) + + // Even though only 1KB is actually used on disk, the store enforces capacity based on the + // declared 8KB, so a 3KB blob doesn't fit alongside it (there's nothing evictable to make room). + _, err := store.Create(core.DigestFixture().Hex(), 3*memsize.KB) + require.EqualError(err, "reserve space: cannot evict enough, the unevictable/incomplete blobs are using up all the space") + + // Declares 2KB (fits exactly within the remaining capacity) but writes 5KB. + overreportedF, overreportedKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, overreportedF, 5*memsize.KB) + require.NoError(overreportedF.Close()) + require.NoError(store.MarkComplete(overreportedKey)) + // The store still only accounts for the declared 2KB, not the actual 5KB written. + require.Equal(10*memsize.KB, store.size) + + // Deleting releases exactly the declared size that was reserved, not the actual bytes found on + // disk, keeping reservation accounting self-consistent in both directions. + require.NoError(store.Delete(overreportedKey)) + require.Equal(8*memsize.KB, store.size) + require.NoError(store.Delete(underreportedKey)) + require.Equal(uint64(0), store.size) } func fillWithRandomData(t *testing.T, f FileReadWriter, sizeBytes uint64) []byte { diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index 849025752..1de21b2ad 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -7,6 +7,7 @@ import ( "io" "os" "path/filepath" + "strconv" "sync" "github.com/uber/kraken/lib/store/metadata" @@ -26,6 +27,7 @@ const ( _incompleteBlob = false _defaultFilePerm = 0775 _evictionBannedFileName = "_eviction_banned" + _blobSizeFileName = "_size" ) // DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. @@ -47,8 +49,11 @@ type DiskStore struct { blobs map[string]*blob // TODO - consider whether it's better to use struct instead of pointer to reduce GC stress. evictQueue *list.List // Back is most recently used, front is the next to evict. // synchronizes mem state access and syscalls to the fs in the APIs (opening, moving files, etc.) - mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. - log *zap.SugaredLogger + mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. + // If enabled, incomplete blobs are rebooted in the store upon restart (or after a crash), allowing users to continue using the blob. + // If disabled, incomplete blobs are discarded (usually done to prevent leaks). + rebootIncompleteBlobs bool + log *zap.SugaredLogger *pather } @@ -62,43 +67,42 @@ type blob struct { // NewDiskStore initializes a new [*DiskStore]. If the store has been initialized in the same // directory before, its state is recovered from disk with the following caveats: // -// - incomplete blobs are deleted to prevent leaks, in case the clients have crashed. +// - `rebootIncompleteBlobs` configures whether incomplete blobs are evicted or rebooted on restart. // -// - the blobs' sizes are recovered from disk, which may differ from the client-provided sizes in Create. This is not a problem. -// -// - if the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), +// - If the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), // it evicts blobs until size is within capacity. -func NewDiskStore(capacityBytes uint64, rootDir string) (*DiskStore, error) { +func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool) (*DiskStore, error) { // TODO - create a Config struct. // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. // TODO - move disk store files into their own directory and package. log := log.Default().With("module", "disk_store") - ok, err := existsPersistedState(rootDir) + ok, err := existsPersistedStore(rootDir) if err != nil { err = fmt.Errorf("could not check if previously-left persisted state exists on disk: %w", err) log.With("error", err).Error("Failed to initialize disk store") return nil, err } if !ok { - log.Info("Did not find any previously persisted state to reboot for DiskStore - initializing a new, empty DiskStore") + log.Info("Initialized a new, empty DiskStore (did not find any previously persisted state to reboot for DiskStore)") return &DiskStore{ - capacity: capacityBytes, - size: 0, - blobs: make(map[string]*blob), - evictQueue: list.New(), - log: log, - pather: newPather(rootDir), + capacity: capacityBytes, + size: 0, + blobs: make(map[string]*blob), + evictQueue: list.New(), + log: log, + pather: newPather(rootDir), + rebootIncompleteBlobs: rebootIncompleteBlobs, }, nil } - store, err := rebootPersistedStateAfterCrash(capacityBytes, rootDir, log) + store, err := rebootPersistedStore(capacityBytes, rootDir, rebootIncompleteBlobs, log) if err != nil { err = fmt.Errorf("reboot persisted state into memory: %w", err) log.With("error", err).Error("Failed to initialize disk store") return nil, err } - log.With("num_blobs", len(store.blobs)).Info("Successfully initialized disk store") + log.With("num_blobs", len(store.blobs)).Info("Successfully rebooted DiskStore's previously left state on disk") return store, nil } @@ -120,7 +124,7 @@ func (s *DiskStore) Open(key string, ignoreIncomplete bool) (FileReadWriter, err if err != nil { return nil, fmt.Errorf("open: %w", err) } - return newReadWriter(f, b.size), nil + return newReadWriter(f), nil } // Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExists] if the blob is not found. @@ -173,6 +177,14 @@ func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) return nil, fmt.Errorf("open file: %w", err) } + if s.rebootIncompleteBlobs { + err = s.persistBlobSize(key, sizeBytes) + if err != nil { + // Fail-open: the blob will be discarded upon reboot if incomplete. + s.log.With("error", err).Error("Could not persist client-provided blob size on disk") + } + } + s.blobs[key] = &blob{ size: sizeBytes, node: nil, @@ -180,10 +192,25 @@ func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) evictionBanned: false, } - return newReadWriter(f, sizeBytes), nil + return newReadWriter(f), nil +} + +func (s *DiskStore) persistBlobSize(key string, sizeBytes uint64) error { + blobSizeFilePath := s.sidecarFilePath(key, _incompleteBlob, _blobSizeFileName) + blobSizeF, err := os.OpenFile(blobSizeFilePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, _defaultFilePerm) + if err != nil { + return fmt.Errorf("create file: %w", err) + } + _, err = blobSizeF.Write([]byte(strconv.Itoa(int(sizeBytes)))) + if err != nil { + return fmt.Errorf("write to size file: %w", err) + } + closers.Close(blobSizeF) + return nil } func (s *DiskStore) reserveSpace(space uint64) error { + // TODO - benchmark and consider whether async eviction makes more sense. // TODO - emit latency to reserve space for a blob. for s.size+space > s.capacity { if s.evictQueue.Len() == 0 { diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index bf2252ece..42ca539e2 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -19,12 +19,12 @@ import ( "github.com/uber/kraken/utils/memsize" ) -func newTestStore(t *testing.T, capacity uint64) (res *DiskStore, rootDir string) { +func newTestStore(t *testing.T, capacity uint64, rebootIncompleteBlobs bool) (res *DiskStore, rootDir string) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) t.Cleanup(func() { os.RemoveAll(rootDir) }) - store, err := NewDiskStore(capacity, rootDir) + store, err := NewDiskStore(capacity, rootDir, rebootIncompleteBlobs) require.NoError(t, err) return store, rootDir } @@ -58,7 +58,7 @@ func numBlobsOnDisk(t *testing.T, store *DiskStore) int { func TestDiskStore(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) keys := []string{} for i := range 10 { @@ -107,7 +107,7 @@ func TestDiskStore(t *testing.T) { func TestEviction(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 25*memsize.KB) + store, _ := newTestStore(t, 25*memsize.KB, false) // create 5 blobs - a, b, c, d, e with different sizes. a, aKey := newTestFile(t, store, 10*memsize.KB) require.NoError(a.Close()) @@ -218,7 +218,7 @@ func TestEviction(t *testing.T) { func TestParallelAccessToSingleFile(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 1*memsize.KB) @@ -295,7 +295,7 @@ func TestParallelAccessToSingleFile(t *testing.T) { func TestOpenedFileAccessibleAfterMarkedComplete(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 1*memsize.KB) @@ -333,7 +333,7 @@ func TestOpenedFileAccessibleAfterEviction(t *testing.T) { func TestDelete(t *testing.T) { t.Run("incomplete blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -348,7 +348,7 @@ func TestDelete(t *testing.T) { }) t.Run("incomplete, unevictable blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -365,7 +365,7 @@ func TestDelete(t *testing.T) { }) t.Run("complete blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -382,7 +382,7 @@ func TestDelete(t *testing.T) { }) t.Run("complete, unevictable blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -400,7 +400,7 @@ func TestDelete(t *testing.T) { }) t.Run("not found", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() err := store.Delete(key) @@ -411,7 +411,7 @@ func TestDelete(t *testing.T) { func TestMarkComplete(t *testing.T) { t.Run("incomplete blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -428,7 +428,7 @@ func TestMarkComplete(t *testing.T) { }) t.Run("incomplete blob with forbidden eviction", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -446,7 +446,7 @@ func TestMarkComplete(t *testing.T) { }) t.Run("already complete blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -459,7 +459,7 @@ func TestMarkComplete(t *testing.T) { }) t.Run("already complete blob with forbidden eviction", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 100*memsize.B) require.NoError(err) @@ -473,7 +473,7 @@ func TestMarkComplete(t *testing.T) { }) t.Run("not found", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() err := store.MarkComplete(key) @@ -484,7 +484,7 @@ func TestMarkComplete(t *testing.T) { func TestStat(t *testing.T) { t.Run("complete blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 10*memsize.B) require.NoError(err) @@ -506,7 +506,7 @@ func TestStat(t *testing.T) { }) t.Run("complete, unevictable blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 10*memsize.B) require.NoError(err) @@ -529,7 +529,7 @@ func TestStat(t *testing.T) { }) t.Run("incomplete blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 10*memsize.B) require.NoError(err) @@ -551,7 +551,7 @@ func TestStat(t *testing.T) { t.Run("incomplete, unevictable blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 10*memsize.B) require.NoError(err) @@ -573,7 +573,7 @@ func TestStat(t *testing.T) { }) t.Run("non-existent blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() _, err := store.Stat(key, IgnoreIncompleteBlobs) @@ -585,7 +585,7 @@ func TestStat(t *testing.T) { func TestList(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) require.Empty(store.List(CheckIncompleteBlobs)) require.Empty(store.List(IgnoreIncompleteBlobs)) @@ -629,7 +629,7 @@ func TestList(t *testing.T) { func TestMetadata(t *testing.T) { t.Run("basic functionality", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 10*memsize.KB) require.NoError(err) @@ -661,7 +661,7 @@ func TestMetadata(t *testing.T) { t.Run("non-existent blob", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) nonExistentKey := core.DigestFixture().Hex() mdStruct := core.MetaInfoFixture() md := metadata.NewTorrentMeta(mdStruct) @@ -679,7 +679,7 @@ func TestMetadata(t *testing.T) { t.Run("metadata does not change after marking a file as complete and/or evictable/unevictable", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() f, err := store.Create(key, 1*memsize.KB) require.NoError(err) @@ -720,7 +720,7 @@ func TestMetadata(t *testing.T) { }) t.Run("metadata fully gone after blob is evicted", func(t *testing.T) { require := require.New(t) - store, _ := newTestStore(t, 10*memsize.KB) + store, _ := newTestStore(t, 10*memsize.KB, false) keyA := core.DigestFixture().Hex() fA, err := store.Create(keyA, 10*memsize.KB) require.NoError(err) diff --git a/lib/store/file.go b/lib/store/file.go index 852dae1d4..2f36a1d7b 100644 --- a/lib/store/file.go +++ b/lib/store/file.go @@ -25,10 +25,9 @@ type FileReadWriter = base.FileReadWriter // FileReader is a read-only file. type FileReader = base.FileReader -func newReadWriter(f *os.File, size uint64) FileReadWriter { +func newReadWriter(f *os.File) FileReadWriter { return &rwImpl{ File: f, - size: int64(size), } } @@ -36,12 +35,15 @@ var _ FileReadWriter = &rwImpl{} type rwImpl struct { *os.File - size int64 } -// Size returns the full size of the blob in bytes, even if the blob is not fully written yet. +// Size returns the number of bytes the file contains. func (i *rwImpl) Size() int64 { - return i.size + info, err := i.Stat() + if err != nil { + return 0 + } + return info.Size() } // Cancel is supposed to remove any written content. diff --git a/lib/store/pather.go b/lib/store/pather.go index fed4065fa..d4d505b4d 100644 --- a/lib/store/pather.go +++ b/lib/store/pather.go @@ -40,7 +40,7 @@ func (p *pather) dirPath(key string, complete bool) string { return filepath.Join(dirPath, key) } -func (p *pather) sidecarFilePath(key string, complete bool, sidecarFilePath string) string { +func (p *pather) sidecarFilePath(key string, complete bool, sidecarFileName string) string { dirPath := p.dirPath(key, complete) - return filepath.Join(dirPath, sidecarFilePath) + return filepath.Join(dirPath, sidecarFileName) } diff --git a/lib/store/pather_test.go b/lib/store/pather_test.go index 8ea7cd80e..1a17d2a85 100644 --- a/lib/store/pather_test.go +++ b/lib/store/pather_test.go @@ -1,42 +1,45 @@ package store import ( + "os" "testing" "github.com/stretchr/testify/require" "github.com/uber/kraken/core" "github.com/uber/kraken/lib/store/metadata" - "github.com/uber/kraken/utils/memsize" ) func TestPather(t *testing.T) { require := require.New(t) - store, rootDir := newTestStore(t, 10*memsize.KB) + rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") + require.NoError(err) + t.Cleanup(func() { os.RemoveAll(rootDir) }) + pather := newPather(rootDir) md := metadata.NewTorrentMeta(core.MetaInfoFixture()) key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" t.Run("incomplete entry", func(t *testing.T) { complete := false - dirPath := store.dirPath(key, complete) + dirPath := pather.dirPath(key, complete) wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" require.Equal(wantDirPath, dirPath) - blobPath := store.blobPath(key, complete) + blobPath := pather.blobPath(key, complete) wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" require.Equal(wantBlobPath, blobPath) - sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) + sidecarFilePath := pather.sidecarFilePath(key, complete, md.GetSuffix()) wantSidecarFilePath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" require.Equal(wantSidecarFilePath, sidecarFilePath) }) t.Run("complete entry", func(t *testing.T) { complete := true - dirPath := store.dirPath(key, complete) + dirPath := pather.dirPath(key, complete) wantDirPath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" require.Equal(wantDirPath, dirPath) - blobPath := store.blobPath(key, complete) + blobPath := pather.blobPath(key, complete) wantBlobPath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" require.Equal(wantBlobPath, blobPath) - sidecarFilePath := store.sidecarFilePath(key, complete, md.GetSuffix()) + sidecarFilePath := pather.sidecarFilePath(key, complete, md.GetSuffix()) wantSidecarFilePath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" require.Equal(wantSidecarFilePath, sidecarFilePath) }) From 8981dd13404f14648ee7e415aa79f1981ddf6589 Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Wed, 22 Jul 2026 19:26:21 +0200 Subject: [PATCH 07/10] feat(DiskStore): Support scoping store to only (in-)complete blobs 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. --- lib/store/crash_recovery.go | 4 +- lib/store/crash_recovery_test.go | 32 ++-- lib/store/disk_store.go | 141 +++++++++--------- lib/store/disk_store_test.go | 246 ++++++++++++++++++++++++------- lib/store/scoped_store.go | 101 +++++++++++++ 5 files changed, 387 insertions(+), 137 deletions(-) create mode 100644 lib/store/scoped_store.go diff --git a/lib/store/crash_recovery.go b/lib/store/crash_recovery.go index e5e23230e..e7b177844 100644 --- a/lib/store/crash_recovery.go +++ b/lib/store/crash_recovery.go @@ -27,7 +27,7 @@ type blobState struct { complete bool } -func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, log *zap.SugaredLogger) (*DiskStore, error) { +func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, log *zap.SugaredLogger) (*diskStore, error) { completeDirPath, incompleteDirPath := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) if !rebootIncompleteBlobs { err := os.RemoveAll(incompleteDirPath) @@ -96,7 +96,7 @@ func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncomplete storeSize += bState.size } - store := &DiskStore{ + store := &diskStore{ blobs: blobs, evictQueue: evictQueue, capacity: capacityBytes, diff --git a/lib/store/crash_recovery_test.go b/lib/store/crash_recovery_test.go index aa99759a0..9ede66a82 100644 --- a/lib/store/crash_recovery_test.go +++ b/lib/store/crash_recovery_test.go @@ -44,7 +44,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) // Incomplete files are recovered (since `rebootIncompleteBlobs` is true). - f, err := store.Open(incompleteEvictableKey, CheckIncompleteBlobs) + f, err := store.Open(incompleteEvictableKey) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err := io.ReadAll(f) @@ -54,7 +54,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) require.False(unevictable) - f, err = store.Open(incompleteUnevictableKey, CheckIncompleteBlobs) + f, err = store.Open(incompleteUnevictableKey) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) @@ -65,7 +65,7 @@ func TestCrashRecovery(t *testing.T) { require.False(unevictable) // Complete files are always recovered. - f, err = store.Open(completeEvictableKey, IgnoreIncompleteBlobs) + f, err = store.ScopeComplete().Open(completeEvictableKey) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) @@ -75,7 +75,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) require.False(unevictable) - f, err = store.Open(completeUnevictableKey, IgnoreIncompleteBlobs) + f, err = store.ScopeComplete().Open(completeUnevictableKey) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) @@ -93,13 +93,13 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) // Incomplete files are dropped. - _, err = store.Stat(incompleteEvictableKey, CheckIncompleteBlobs) + _, err = store.Stat(incompleteEvictableKey) require.ErrorIs(err, os.ErrNotExist) - _, err = store.Stat(incompleteUnevictableKey, CheckIncompleteBlobs) + _, err = store.Stat(incompleteUnevictableKey) require.ErrorIs(err, os.ErrNotExist) // Complete files are always recovered. - f, err = store.Open(completeEvictableKey, IgnoreIncompleteBlobs) + f, err = store.ScopeComplete().Open(completeEvictableKey) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) @@ -109,7 +109,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) require.False(unevictable) - f, err = store.Open(completeUnevictableKey, IgnoreIncompleteBlobs) + f, err = store.ScopeComplete().Open(completeUnevictableKey) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) @@ -137,7 +137,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) var readMd metadata.TorrentMeta - ok, err := store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) + ok, err := store.ScopeComplete().GetMetadata(key, &readMd) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) @@ -157,7 +157,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) var readMd metadata.TorrentMeta - ok, err := store.GetMetadata(key, &readMd, CheckIncompleteBlobs) + ok, err := store.GetMetadata(key, &readMd) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) @@ -194,7 +194,7 @@ func TestCrashRecovery(t *testing.T) { require.Equal([]string{cKey, dKey, eKey}, store.evictionOrder()) // a is unevictable and b is incomplete // reset the access time for d - dF, err := store.Open(dKey, IgnoreIncompleteBlobs) + dF, err := store.ScopeComplete().Open(dKey) require.NoError(err) require.NoError(dF.Close()) evictionOrderBeforeCrash := store.evictionOrder() @@ -215,11 +215,11 @@ func TestCrashRecovery(t *testing.T) { require.NoError(err) // since 10KB of blobs are in store, `c` gets evicted to put the store back within its capacity. - _, err = rebootedSmallerStore.Stat(cKey, CheckIncompleteBlobs) + _, err = rebootedSmallerStore.Stat(cKey) require.ErrorIs(err, os.ErrNotExist) require.Equal([]string{dKey, eKey}, rebootedSmallerStore.evictionOrder()) - _, err = rebootedSmallerStore.Stat(aKey, CheckIncompleteBlobs) + _, err = rebootedSmallerStore.Stat(aKey) require.NoError(err) }) } @@ -235,7 +235,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { store, err := NewDiskStore(10*memsize.KB, rootDir, true) require.NoError(err) - f, err = store.Open(key, CheckIncompleteBlobs) + f, err = store.Open(key) require.NoError(err) // Write 5KB in total, while reporting only 4KB. @@ -253,7 +253,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { copy(wantData, firstData) copy(wantData[2*memsize.KB:], secondData) - f, err = store.Open(key, CheckIncompleteBlobs) + f, err = store.Open(key) require.NoError(err) data, err := io.ReadAll(f) require.NoError(err) @@ -266,7 +266,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { store, err = NewDiskStore(10*memsize.KB, rootDir, true) // now that the blob is complete, its actual size should be rebooted through stat, instead of trusting the _size sidecar file. require.Equal(5*memsize.KB, store.size) - f, err = store.Open(key, IgnoreIncompleteBlobs) + f, err = store.ScopeComplete().Open(key) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) data, err = io.ReadAll(f) diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index 1de21b2ad..085208fd0 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -16,10 +16,14 @@ import ( "go.uber.org/zap" ) -// Whether the store's read APIs ignore incomplete blobs. +// the set of blobs that the store's APIs can operate on. +type blobScope int + +// flags to scope [diskStore]'s APIs to a subset of blobs. const ( - IgnoreIncompleteBlobs = true - CheckIncompleteBlobs = false + blobScopeAny blobScope = iota + blobScopeComplete + blobScopeIncomplete ) const ( @@ -30,20 +34,11 @@ const ( _blobSizeFileName = "_size" ) -// DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. -// -// - Supports pagination of blobs during reading/writing, such that blobs don't need to be fully loaded into memory. -// -// - New blobs are considered 'incomplete', which unlists them from LRU eviction. Read APIs may filter out incomplete blobs. -// -// - All APIs are thread-safe. Parallel access to a single file is allowed but clients must ensure they don't intervene with one another. +// diskStore implements the APIs of [DiskStore]. [diskStore]'s APIs expose the [blobScope] arg, +// while [DiskStore]'s APIs omit that arg (cleaner interface) and instead expose other APIs to scope the whole store. // -// - Supports (un-)marking blobs as non-evictable (may be needed when that data must be written back to remote storage). -// -// - Crash-resistant - all state is restored upon restart (check [NewDiskStore] for details). -// -// - Uses directory sharding to speed up disk performance. -type DiskStore struct { +// Check [DiskStore]'s comments for details on functionality. +type diskStore struct { capacity uint64 size uint64 // includes both used and reserved space. blobs map[string]*blob // TODO - consider whether it's better to use struct instead of pointer to reduce GC stress. @@ -64,14 +59,7 @@ type blob struct { evictionBanned bool } -// NewDiskStore initializes a new [*DiskStore]. If the store has been initialized in the same -// directory before, its state is recovered from disk with the following caveats: -// -// - `rebootIncompleteBlobs` configures whether incomplete blobs are evicted or rebooted on restart. -// -// - If the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), -// it evicts blobs until size is within capacity. -func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool) (*DiskStore, error) { +func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool) (*diskStore, error) { // TODO - create a Config struct. // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. // TODO - move disk store files into their own directory and package. @@ -85,7 +73,7 @@ func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bo } if !ok { log.Info("Initialized a new, empty DiskStore (did not find any previously persisted state to reboot for DiskStore)") - return &DiskStore{ + return &diskStore{ capacity: capacityBytes, size: 0, blobs: make(map[string]*blob), @@ -106,15 +94,17 @@ func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bo return store, nil } -// Open returns an FD to a file in the store. [os.ErrNotExists] is returned on missing entry. -func (s *DiskStore) Open(key string, ignoreIncomplete bool) (FileReadWriter, error) { +func (s *diskStore) open(key string, scope blobScope) (FileReadWriter, error) { s.mu.Lock() defer s.mu.Unlock() b, ok := s.blobs[key] - if !ok || (ignoreIncomplete && !b.complete) { + if !ok { return nil, os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return nil, err + } if b.node != nil { s.evictQueue.MoveToBack(b.node) @@ -127,25 +117,22 @@ func (s *DiskStore) Open(key string, ignoreIncomplete bool) (FileReadWriter, err return newReadWriter(f), nil } -// Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExists] if the blob is not found. -func (s *DiskStore) Stat(key string, ignoreIncomplete bool) (os.FileInfo, error) { - // We **could** avoid locking the mutex by just statting the file directly. However, the current implementation - // prefers mutex contention over extra disk usage, as origin is bottlenecked by disk IO. +func (s *diskStore) stat(key string, scope blobScope) (os.FileInfo, error) { s.mu.Lock() defer s.mu.Unlock() b, ok := s.blobs[key] - if !ok || (ignoreIncomplete && !b.complete) { + if !ok { return nil, os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return nil, err + } blobPath := s.blobPath(key, b.complete) return os.Stat(blobPath) } -// Create adds a new, incomplete blob to the store and reserves space for it. -// Incomplete entries cannot be automatically evicted. MarkComplete must be called once the blob is complete. -// DiskStore does not ever check/use the real size of the blob and only uses `sizeBytes` for its eviction logic. -func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) { +func (s *diskStore) create(key string, sizeBytes uint64) (FileReadWriter, error) { // TODO - we might want some TTI on uploads to the store, after which we cancel the upload, e.g. 1min without the client uploading more data. s.mu.Lock() defer s.mu.Unlock() @@ -195,7 +182,7 @@ func (s *DiskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) return newReadWriter(f), nil } -func (s *DiskStore) persistBlobSize(key string, sizeBytes uint64) error { +func (s *diskStore) persistBlobSize(key string, sizeBytes uint64) error { blobSizeFilePath := s.sidecarFilePath(key, _incompleteBlob, _blobSizeFileName) blobSizeF, err := os.OpenFile(blobSizeFilePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, _defaultFilePerm) if err != nil { @@ -209,7 +196,7 @@ func (s *DiskStore) persistBlobSize(key string, sizeBytes uint64) error { return nil } -func (s *DiskStore) reserveSpace(space uint64) error { +func (s *diskStore) reserveSpace(space uint64) error { // TODO - benchmark and consider whether async eviction makes more sense. // TODO - emit latency to reserve space for a blob. for s.size+space > s.capacity { @@ -235,22 +222,22 @@ func (s *DiskStore) reserveSpace(space uint64) error { return nil } -func (s *DiskStore) releaseSpace(space uint64) { - // TODO - if space > s.size, emit an error log for an invariant violation +func (s *diskStore) releaseSpace(space uint64) { + if space > s.size { + s.log.Error("Invariant violation - DiskStore wants to release more disk space than actually reserved. Failing open by releasing all reserved space.") + s.size = 0 + return + } s.size -= space } -// fully deletes the disk state of a blob, including metadata. Works on any blob. -func (s *DiskStore) deleteFromDisk(key string, complete bool) error { +// Fully deletes the disk state of a blob, including metadata. Works on any blob. +func (s *diskStore) deleteFromDisk(key string, complete bool) error { dir := s.dirPath(key, complete) return os.RemoveAll(dir) } -// MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless BanEviction has been called). -// Additionally, read APIs may optionally filter out incomplete blobs. -func (s *DiskStore) MarkComplete(key string) error { - // TODO - check if we can derive when a blob is considered complete (e.g. when client calls Close on file (although that depends on the - // assumption that Close means the file is complete which may not be true if the client that created the file expects another client to continue mutating it)). +func (s *diskStore) markComplete(key string) error { s.mu.Lock() defer s.mu.Unlock() @@ -282,7 +269,7 @@ func (s *DiskStore) MarkComplete(key string) error { return nil } -func (s *DiskStore) checkDiskIfUnevictable(key string, complete bool) (bool, error) { +func (s *diskStore) checkDiskIfUnevictable(key string, complete bool) (bool, error) { flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) unevictable, err := exists(flagBlobPath) if err != nil { @@ -291,8 +278,7 @@ func (s *DiskStore) checkDiskIfUnevictable(key string, complete bool) (bool, err return unevictable, nil } -// Delete removes a blob and its [metadata.Metadata] from the store. Works on any blob. -func (s *DiskStore) Delete(key string) error { +func (s *diskStore) delete(key string, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -300,6 +286,9 @@ func (s *DiskStore) Delete(key string) error { if !ok { return os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return err + } err := s.deleteFromDisk(key, b.complete) if err != nil { return fmt.Errorf("delete from disk: %w", err) @@ -314,14 +303,13 @@ func (s *DiskStore) Delete(key string) error { return nil } -// List returns the blobs' keys. -func (s *DiskStore) List(ignoreIncomplete bool) []string { +func (s *diskStore) list(scope blobScope) []string { s.mu.RLock() defer s.mu.RUnlock() res := make([]string, 0, len(s.blobs)) for key, b := range s.blobs { - if ignoreIncomplete && !b.complete { + if err := isOutOfScope(b, scope); err != nil { continue } res = append(res, key) @@ -329,9 +317,7 @@ func (s *DiskStore) List(ignoreIncomplete bool) []string { return res } -// BanEviction marks a blob as unevictable by LRU eviction. It is idempotent. -// Needed when e.g. blobs must be written back to GCS/S3 and eviction before that is unacceptable. -func (s *DiskStore) BanEviction(key string) error { +func (s *diskStore) banEviction(key string, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -339,14 +325,16 @@ func (s *DiskStore) BanEviction(key string) error { if !ok { return os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return err + } if b.evictionBanned { // no-op return nil } flagBlobPath := s.sidecarFilePath(key, b.complete, _evictionBannedFileName) - // We persist the ban as a flag file on disk, such that after - // a system crash, we can recover the ban. + // We persist the ban as a flag file on disk for crash-resilience. f, err := os.OpenFile(flagBlobPath, os.O_RDONLY|os.O_CREATE, _defaultFilePerm) if err != nil { return fmt.Errorf("create file that flags eviction as banned: %w", err) @@ -361,8 +349,7 @@ func (s *DiskStore) BanEviction(key string) error { return nil } -// UnbanDeletion removes the effect of BanDeletion for a blob. It is idempotent. -func (s *DiskStore) UnbanEviction(key string) error { +func (s *diskStore) unbanEviction(key string, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -370,6 +357,9 @@ func (s *DiskStore) UnbanEviction(key string) error { if !ok { return os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return err + } if !b.evictionBanned { // no-op return nil @@ -389,8 +379,7 @@ func (s *DiskStore) UnbanEviction(key string) error { return nil } -// SetMetadata atomically sets the respective metadata for a blob. Works on any blob. -func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { +func (s *diskStore) setMetadata(key string, md metadata.Metadata, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -398,6 +387,9 @@ func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { if !ok { return os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return err + } mdData, err := md.Serialize() if err != nil { @@ -425,15 +417,17 @@ func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { return nil } -// GetMetadata populates `md` if the metadata is present. Returns [os.ErrNotExists] if key is not in store. -func (s *DiskStore) GetMetadata(key string, md metadata.Metadata, ignoreIncomplete bool) (ok bool, err error) { +func (s *diskStore) getMetadata(key string, md metadata.Metadata, scope blobScope) (ok bool, err error) { s.mu.RLock() defer s.mu.RUnlock() b, ok := s.blobs[key] - if !ok || (ignoreIncomplete && !b.complete) { + if !ok { return false, os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return false, err + } mdFilePath := s.sidecarFilePath(key, b.complete, md.GetSuffix()) mdFile, err := os.OpenFile(mdFilePath, os.O_RDONLY, _defaultFilePerm) @@ -452,8 +446,7 @@ func (s *DiskStore) GetMetadata(key string, md metadata.Metadata, ignoreIncomple return true, nil } -// DeleteMetadata removes any metadata of a blob with `md`'s suffix, if present. -func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { +func (s *diskStore) deleteMetadata(key string, md metadata.Metadata, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -461,6 +454,9 @@ func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { if !ok { return os.ErrNotExist } + if err := isOutOfScope(b, scope); err != nil { + return err + } mdFilePath := s.sidecarFilePath(key, b.complete, md.GetSuffix()) err := os.Remove(mdFilePath) if errors.Is(err, os.ErrNotExist) { @@ -474,7 +470,7 @@ func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { } // used during testing -func (s *DiskStore) evictionOrder() []string { +func (s *diskStore) evictionOrder() []string { s.mu.RLock() defer s.mu.RUnlock() @@ -496,3 +492,10 @@ func exists(path string) (ok bool, err error) { } return false, fmt.Errorf("stat: %w", err) } + +func isOutOfScope(b *blob, scope blobScope) error { + if (b.complete && scope == blobScopeIncomplete) || (!b.complete && scope == blobScopeComplete) { + return ErrOutOfScope + } + return nil +} diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index 42ca539e2..3c86b60c5 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -82,12 +82,12 @@ func TestDiskStore(t *testing.T) { require.EqualError(err, "reserve space: cannot evict enough, the unevictable/incomplete blobs are using up all the space") require.Nil(f) - f, err = store.Open(keys[0], IgnoreIncompleteBlobs) - require.ErrorIs(err, os.ErrNotExist) + f, err = store.ScopeComplete().Open(keys[0]) + require.ErrorIs(err, ErrOutOfScope) require.Nil(f) require.NoError(store.MarkComplete(keys[0])) - f, err = store.Open(keys[0], IgnoreIncompleteBlobs) + f, err = store.Open(keys[0]) defer func(f io.Closer) { require.NoError(f.Close()) }(f) require.NoError(err) wantData := make([]byte, memsize.KB) @@ -100,7 +100,7 @@ func TestDiskStore(t *testing.T) { f, err = store.Create(core.DigestFixture().Hex(), memsize.KB) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) - f, err = store.Open(keys[0], IgnoreIncompleteBlobs) + f, err = store.Open(keys[0]) require.ErrorIs(err, os.ErrNotExist) require.Nil(f) } @@ -139,7 +139,7 @@ func TestEviction(t *testing.T) { f, fKey := newTestFile(t, store, 4*memsize.KB) require.NoError(f.Close()) require.NoError(store.MarkComplete(fKey)) - keys := store.List(CheckIncompleteBlobs) + keys := store.List() require.NotContains(keys, cKey) // new size == 23KB == 24KB - 5KB (c) + 4KB (f) require.Equal(23*memsize.KB, store.size) @@ -159,7 +159,7 @@ func TestEviction(t *testing.T) { // size == 24KB == 24KB + 15KB (h) - 5KB (b) - 10KB (a) require.Equal(24*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) - keys = store.List(CheckIncompleteBlobs) + keys = store.List() require.NotContains(keys, bKey) require.NotContains(keys, aKey) @@ -168,14 +168,14 @@ func TestEviction(t *testing.T) { require.NoError(store.UnbanEviction(eKey)) // eviction order (left-most is next to evict): f(4KB), g(1KB), h(15KB), e(1KB); d(3KB) is unevictable // we open g to change the order to f, h, e, g - g, err = store.Open(gKey, CheckIncompleteBlobs) + g, err = store.Open(gKey) require.NoError(err) require.NoError(g.Close()) i, iKey := newTestFile(t, store, 5*memsize.KB) require.NoError(store.MarkComplete(iKey)) require.NoError(i.Close()) - keys = store.List(CheckIncompleteBlobs) + keys = store.List() require.NotContains(keys, fKey) require.Equal(25*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -184,7 +184,7 @@ func TestEviction(t *testing.T) { j, jKey := newTestFile(t, store, 14*memsize.KB) require.NoError(j.Close()) require.NoError(store.MarkComplete(jKey)) - keys = store.List(CheckIncompleteBlobs) + keys = store.List() require.NotContains(keys, hKey) require.Equal(24*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -193,7 +193,7 @@ func TestEviction(t *testing.T) { k, kKey := newTestFile(t, store, 2*memsize.KB) require.NoError(k.Close()) require.NoError(store.MarkComplete(kKey)) - keys = store.List(CheckIncompleteBlobs) + keys = store.List() require.NotContains(keys, eKey) require.Equal(25*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -202,7 +202,7 @@ func TestEviction(t *testing.T) { l, lKey := newTestFile(t, store, 1*memsize.KB) require.NoError(store.MarkComplete(lKey)) require.NoError(l.Close()) - keys = store.List(CheckIncompleteBlobs) + keys = store.List() require.NotContains(keys, gKey) require.Equal(25*memsize.KB, store.size) require.Equal(5, numBlobsOnDisk(t, store)) @@ -239,7 +239,7 @@ func TestParallelAccessToSingleFile(t *testing.T) { go func(idx int64) { defer wg.Done() - f, err := store.Open(key, CheckIncompleteBlobs) + f, err := store.Open(key) if err != nil { results[idx].err = err return @@ -280,7 +280,7 @@ func TestParallelAccessToSingleFile(t *testing.T) { require.NoError(store.MarkComplete(key)) - f, err = store.Open(key, IgnoreIncompleteBlobs) + f, err = store.ScopeComplete().Open(key) require.NoError(err) defer func() { require.NoError(f.Close()) }() @@ -304,13 +304,13 @@ func TestOpenedFileAccessibleAfterMarkedComplete(t *testing.T) { require.NoError(err) require.NoError(f.Close()) - incompleteFile, err := store.Open(key, CheckIncompleteBlobs) + incompleteFile, err := store.Open(key) require.NoError(err) defer func() { require.NoError(incompleteFile.Close()) }() require.NoError(store.MarkComplete(key)) - completeFile, err := store.Open(key, IgnoreIncompleteBlobs) + completeFile, err := store.ScopeComplete().Open(key) require.NoError(err) defer func() { require.NoError(completeFile.Close()) }() @@ -342,7 +342,7 @@ func TestDelete(t *testing.T) { require.NoError(f.Close()) require.NoError(store.Delete(key)) - require.Empty(store.List(CheckIncompleteBlobs)) + require.Empty(store.List()) require.Equal(uint64(0), store.size) require.Equal(0, numBlobsOnDisk(t, store)) }) @@ -359,7 +359,7 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) - require.Empty(store.List(CheckIncompleteBlobs)) + require.Empty(store.List()) require.Equal(uint64(0), store.size) require.Equal(0, numBlobsOnDisk(t, store)) }) @@ -376,7 +376,7 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) - require.Empty(store.List(CheckIncompleteBlobs)) + require.Empty(store.List()) require.Equal(uint64(0), store.size) require.Equal(0, numBlobsOnDisk(t, store)) }) @@ -394,7 +394,7 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) - require.Empty(store.List(CheckIncompleteBlobs)) + require.Empty(store.List()) require.Equal(uint64(0), store.size) require.Equal(0, numBlobsOnDisk(t, store)) }) @@ -421,9 +421,9 @@ func TestMarkComplete(t *testing.T) { require.NoError(store.MarkComplete(key)) - require.Equal([]string{key}, store.List(IgnoreIncompleteBlobs)) + require.Equal([]string{key}, store.ScopeComplete().List()) require.Equal(uint64(100), store.size) - _, err = store.Open(key, IgnoreIncompleteBlobs) + _, err = store.ScopeComplete().Open(key) require.NoError(err) }) t.Run("incomplete blob with forbidden eviction", func(t *testing.T) { @@ -439,9 +439,9 @@ func TestMarkComplete(t *testing.T) { require.NoError(store.MarkComplete(key)) - require.Equal([]string{key}, store.List(IgnoreIncompleteBlobs)) + require.Equal([]string{key}, store.ScopeComplete().List()) require.Equal(uint64(100), store.size) - _, err = store.Open(key, IgnoreIncompleteBlobs) + _, err = store.ScopeComplete().Open(key) require.NoError(err) }) t.Run("already complete blob", func(t *testing.T) { @@ -493,9 +493,9 @@ func TestStat(t *testing.T) { require.NoError(err) require.NoError(store.MarkComplete(key)) - fInfo, err := store.Stat(key, IgnoreIncompleteBlobs) + fInfo, err := store.ScopeComplete().Stat(key) require.NoError(err) - _, err = store.Stat(key, CheckIncompleteBlobs) + _, err = store.Stat(key) require.NoError(err) require.False(fInfo.IsDir()) @@ -516,9 +516,9 @@ func TestStat(t *testing.T) { require.NoError(store.MarkComplete(key)) require.NoError(store.BanEviction(key)) - fInfo, err := store.Stat(key, IgnoreIncompleteBlobs) + fInfo, err := store.ScopeComplete().Stat(key) require.NoError(err) - _, err = store.Stat(key, CheckIncompleteBlobs) + _, err = store.Stat(key) require.NoError(err) require.False(fInfo.IsDir()) @@ -537,9 +537,9 @@ func TestStat(t *testing.T) { _, err = io.Copy(f, bytes.NewReader(make([]byte, 10))) require.NoError(err) - _, err = store.Stat(key, IgnoreIncompleteBlobs) - require.Equal(os.ErrNotExist, err) - fInfo, err := store.Stat(key, CheckIncompleteBlobs) + _, err = store.ScopeComplete().Stat(key) + require.Equal(ErrOutOfScope, err) + fInfo, err := store.Stat(key) require.NoError(err) require.False(fInfo.IsDir()) @@ -560,9 +560,9 @@ func TestStat(t *testing.T) { require.NoError(err) require.NoError(store.BanEviction(key)) - _, err = store.Stat(key, IgnoreIncompleteBlobs) - require.Equal(os.ErrNotExist, err) - fInfo, err := store.Stat(key, CheckIncompleteBlobs) + _, err = store.ScopeComplete().Stat(key) + require.Equal(ErrOutOfScope, err) + fInfo, err := store.Stat(key) require.NoError(err) require.False(fInfo.IsDir()) @@ -576,9 +576,9 @@ func TestStat(t *testing.T) { store, _ := newTestStore(t, 10*memsize.KB, false) key := core.DigestFixture().Hex() - _, err := store.Stat(key, IgnoreIncompleteBlobs) + _, err := store.ScopeComplete().Stat(key) require.Equal(os.ErrNotExist, err) - _, err = store.Stat(key, CheckIncompleteBlobs) + _, err = store.Stat(key) require.Equal(os.ErrNotExist, err) }) } @@ -587,8 +587,8 @@ func TestList(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB, false) - require.Empty(store.List(CheckIncompleteBlobs)) - require.Empty(store.List(IgnoreIncompleteBlobs)) + require.Empty(store.ScopeComplete().List()) + require.Empty(store.ScopeComplete().List()) incompleteBlobKey := core.DigestFixture().Hex() f, err := store.Create(incompleteBlobKey, 10*memsize.B) @@ -618,11 +618,11 @@ func TestList(t *testing.T) { require.NoError(store.MarkComplete(unevictableCompleteBlobKey)) wantRes := []string{completeBlobKey, unevictableCompleteBlobKey} - res := store.List(IgnoreIncompleteBlobs) + res := store.ScopeComplete().List() require.ElementsMatch(wantRes, res) wantRes = []string{incompleteBlobKey, completeBlobKey, unevictableCompleteBlobKey, unevictableIncompleteBlobKey} - res = store.List(CheckIncompleteBlobs) + res = store.List() require.ElementsMatch(wantRes, res) } @@ -642,13 +642,13 @@ func TestMetadata(t *testing.T) { require.NoError(err) var readMd metadata.TorrentMeta - ok, err := store.GetMetadata(key, &readMd, CheckIncompleteBlobs) + ok, err := store.GetMetadata(key, &readMd) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) require.NoError(store.DeleteMetadata(key, &readMd)) - ok, err = store.GetMetadata(key, &readMd, CheckIncompleteBlobs) + ok, err = store.GetMetadata(key, &readMd) require.NoError(err) require.False(ok) mdFilePath := store.sidecarFilePath(key, _incompleteBlob, readMd.GetSuffix()) @@ -669,7 +669,7 @@ func TestMetadata(t *testing.T) { err := store.SetMetadata(nonExistentKey, md) require.Equal(os.ErrNotExist, err) - ok, err := store.GetMetadata(nonExistentKey, md, CheckIncompleteBlobs) + ok, err := store.GetMetadata(nonExistentKey, md) require.Equal(os.ErrNotExist, err) require.False(ok) @@ -690,30 +690,30 @@ func TestMetadata(t *testing.T) { require.NoError(store.SetMetadata(key, writtenMd)) var readMd metadata.TorrentMeta - ok, err := store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) - require.Equal(os.ErrNotExist, err) + ok, err := store.ScopeComplete().GetMetadata(key, &readMd) + require.Equal(ErrOutOfScope, err) // incomplete files are ignored require.False(ok) - ok, err = store.GetMetadata(key, &readMd, CheckIncompleteBlobs) + ok, err = store.GetMetadata(key, &readMd) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) // Repeat the tests above for an unevictable file require.NoError(store.BanEviction(key)) - ok, err = store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) - require.Equal(os.ErrNotExist, err) + ok, err = store.ScopeComplete().GetMetadata(key, &readMd) + require.Equal(ErrOutOfScope, err) // incomplete files are ignored require.False(ok) - ok, err = store.GetMetadata(key, &readMd, CheckIncompleteBlobs) + ok, err = store.GetMetadata(key, &readMd) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) require.NoError(store.MarkComplete(key)) - ok, err = store.GetMetadata(key, &readMd, IgnoreIncompleteBlobs) + ok, err = store.ScopeComplete().GetMetadata(key, &readMd) require.NoError(err) require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) @@ -739,10 +739,156 @@ func TestMetadata(t *testing.T) { require.NoError(err) defer func() { require.NoError(fB.Close()) }() - ok, err := store.GetMetadata(keyA, md, CheckIncompleteBlobs) + ok, err := store.GetMetadata(keyA, md) require.Equal(os.ErrNotExist, err) require.False(ok) _, err = os.Stat(mdFilePath) require.True(errors.Is(err, os.ErrNotExist)) }) } + +func TestScopes(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB, false) + + f, key := newTestFile(t, store, 2*memsize.KB) + data := fillWithRandomData(t, f, 2*memsize.KB) + require.NoError(f.Close()) + md := metadata.NewTorrentMeta(core.MetaInfoFixture()) + + // While incomplete, ScopeComplete's APIs reject the blob. + _, err := store.ScopeComplete().Open(key) + require.ErrorIs(err, ErrOutOfScope) + _, err = store.ScopeComplete().Stat(key) + require.ErrorIs(err, ErrOutOfScope) + require.ErrorIs(store.ScopeComplete().BanEviction(key), ErrOutOfScope) + require.ErrorIs(store.ScopeComplete().UnbanEviction(key), ErrOutOfScope) + require.ErrorIs(store.ScopeComplete().SetMetadata(key, md), ErrOutOfScope) + var readMd metadata.TorrentMeta + ok, err := store.ScopeComplete().GetMetadata(key, &readMd) + require.ErrorIs(err, ErrOutOfScope) + require.False(ok) + require.ErrorIs(store.ScopeComplete().DeleteMetadata(key, &readMd), ErrOutOfScope) + require.NotContains(store.ScopeComplete().List(), key) + + // The unscoped store's APIs work regardless of completeness. + f, err = store.Open(key) + require.NoError(err) + readData, err := io.ReadAll(f) + require.NoError(err) + require.Equal(data, readData) + require.NoError(f.Close()) + _, err = store.Stat(key) + require.NoError(err) + require.NoError(store.BanEviction(key)) + require.NoError(store.UnbanEviction(key)) + require.NoError(store.SetMetadata(key, md)) + ok, err = store.GetMetadata(key, &readMd) + require.NoError(err) + require.True(ok) + require.Equal(md.MetaInfo, readMd.MetaInfo) + require.NoError(store.DeleteMetadata(key, &readMd)) + require.Contains(store.List(), key) + + // ScopeIncomplete's APIs also work, since the blob is (still) incomplete. + f, err = store.ScopeIncomplete().Open(key) + require.NoError(err) + readData, err = io.ReadAll(f) + require.NoError(err) + require.Equal(data, readData) + require.NoError(f.Close()) + _, err = store.ScopeIncomplete().Stat(key) + require.NoError(err) + require.NoError(store.ScopeIncomplete().BanEviction(key)) + require.NoError(store.ScopeIncomplete().UnbanEviction(key)) + require.NoError(store.ScopeIncomplete().SetMetadata(key, md)) + ok, err = store.ScopeIncomplete().GetMetadata(key, &readMd) + require.NoError(err) + require.True(ok) + require.Equal(md.MetaInfo, readMd.MetaInfo) + require.NoError(store.ScopeIncomplete().DeleteMetadata(key, &readMd)) + require.Contains(store.ScopeIncomplete().List(), key) + + require.NoError(store.MarkComplete(key)) + + // Now that the blob is complete, the roles reverse: ScopeIncomplete rejects it. + _, err = store.ScopeIncomplete().Open(key) + require.ErrorIs(err, ErrOutOfScope) + _, err = store.ScopeIncomplete().Stat(key) + require.ErrorIs(err, ErrOutOfScope) + require.ErrorIs(store.ScopeIncomplete().BanEviction(key), ErrOutOfScope) + require.ErrorIs(store.ScopeIncomplete().UnbanEviction(key), ErrOutOfScope) + require.ErrorIs(store.ScopeIncomplete().SetMetadata(key, md), ErrOutOfScope) + ok, err = store.ScopeIncomplete().GetMetadata(key, &readMd) + require.ErrorIs(err, ErrOutOfScope) + require.False(ok) + require.ErrorIs(store.ScopeIncomplete().DeleteMetadata(key, &readMd), ErrOutOfScope) + require.NotContains(store.ScopeIncomplete().List(), key) + + // The unscoped store's APIs still work regardless of completeness. + f, err = store.Open(key) + require.NoError(err) + readData, err = io.ReadAll(f) + require.NoError(err) + require.Equal(data, readData) + require.NoError(f.Close()) + _, err = store.Stat(key) + require.NoError(err) + require.NoError(store.BanEviction(key)) + require.NoError(store.UnbanEviction(key)) + require.NoError(store.SetMetadata(key, md)) + ok, err = store.GetMetadata(key, &readMd) + require.NoError(err) + require.True(ok) + require.Equal(md.MetaInfo, readMd.MetaInfo) + require.NoError(store.DeleteMetadata(key, &readMd)) + require.Contains(store.List(), key) + + // ScopeComplete's APIs now work, since the blob is complete. + f, err = store.ScopeComplete().Open(key) + require.NoError(err) + readData, err = io.ReadAll(f) + require.NoError(err) + require.Equal(data, readData) + require.NoError(f.Close()) + _, err = store.ScopeComplete().Stat(key) + require.NoError(err) + require.NoError(store.ScopeComplete().BanEviction(key)) + require.NoError(store.ScopeComplete().UnbanEviction(key)) + require.NoError(store.ScopeComplete().SetMetadata(key, md)) + ok, err = store.ScopeComplete().GetMetadata(key, &readMd) + require.NoError(err) + require.True(ok) + require.Equal(md.MetaInfo, readMd.MetaInfo) + require.NoError(store.ScopeComplete().DeleteMetadata(key, &readMd)) + require.Contains(store.ScopeComplete().List(), key) +} + +func TestScopesDelete(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB, false) + + f, key := newTestFile(t, store, 1*memsize.KB) + require.NoError(f.Close()) + require.ErrorIs(store.ScopeComplete().Delete(key), ErrOutOfScope) + require.NoError(store.ScopeIncomplete().Delete(key)) + _, err := store.Stat(key) + require.ErrorIs(err, os.ErrNotExist) + + f, key = newTestFile(t, store, 1*memsize.KB) + require.NoError(f.Close()) + require.NoError(store.MarkComplete(key)) + require.ErrorIs(store.ScopeIncomplete().Delete(key), ErrOutOfScope) + require.NoError(store.ScopeComplete().Delete(key)) + _, err = store.Stat(key) + require.ErrorIs(err, os.ErrNotExist) + + f, key = newTestFile(t, store, 1*memsize.KB) + require.NoError(f.Close()) + require.NoError(store.Delete(key)) + + f, key = newTestFile(t, store, 1*memsize.KB) + require.NoError(f.Close()) + require.NoError(store.MarkComplete(key)) + require.NoError(store.Delete(key)) +} diff --git a/lib/store/scoped_store.go b/lib/store/scoped_store.go new file mode 100644 index 000000000..5cad968ca --- /dev/null +++ b/lib/store/scoped_store.go @@ -0,0 +1,101 @@ +package store + +import ( + "errors" + "os" + + "github.com/uber/kraken/lib/store/metadata" +) + +// ErrOutOfScope is returned when the provided key is in the store, but not in the store's [blobScope], +// e.g. if the blob is incomplete, but ScopeComplete was called. +var ErrOutOfScope = errors.New("the blob is in DiskStore but filtered by the selected scope") + +// DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. +// +// - Supports pagination of blobs during reading/writing, such that blobs don't need to be fully loaded into memory. +// +// - New blobs are considered 'incomplete', which unlists them from LRU eviction. The store can be scoped to work on only (in-)complete blobs. +// +// - All APIs are thread-safe. Parallel access to a single file is allowed but clients must ensure they don't intervene with one another. +// +// - Supports (un-)marking blobs as non-evictable (may be needed when that data must be written back to remote storage). +// +// - Crash-resistant - all state is restored upon restart (check [newDiskStore] for details). +// +// - Uses directory sharding to speed up disk performance. +type DiskStore struct { + *diskStore + scope blobScope +} + +// NewDiskStore initializes a new [*DiskStore]. If the store has been initialized in the same +// directory before, its state is recovered from disk with the following caveats: +// +// - `rebootIncompleteBlobs` configures whether incomplete blobs are evicted or rebooted on restart. +// +// - If the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), +// it evicts blobs until size is within capacity. +func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool) (*DiskStore, error) { + diskStore, err := newDiskStore(capacityBytes, rootDir, rebootIncompleteBlobs) + if err != nil { + return nil, err + } + return &DiskStore{ + diskStore: diskStore, + scope: blobScopeAny, + }, nil +} + +// Create adds a new, incomplete blob to the store and reserves space for it. +// Incomplete entries cannot be automatically evicted. MarkComplete must be called once the blob is complete. +// DiskStore uses `sizeBytes` for its eviction logic even if the blob's real size differs. +func (s *diskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) { + return s.create(key, sizeBytes) +} + +// Open returns an FD to a file in the store. [os.ErrNotExists] is returned on missing entry. +func (s *DiskStore) Open(key string) (FileReadWriter, error) { return s.open(key, s.scope) } + +// Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExists] if the blob is not found. +func (s *DiskStore) Stat(key string) (os.FileInfo, error) { return s.stat(key, s.scope) } + +// MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless BanEviction has been called). +// Additionally, other store APIs may filter blobs based on completeness. +func (s *diskStore) MarkComplete(key string) error { return s.markComplete(key) } + +// Delete removes a blob and its [metadata.Metadata] from the store. +func (s *DiskStore) Delete(key string) error { return s.delete(key, s.scope) } + +// List returns the keys of all blobs (except those out of scope). +func (s *DiskStore) List() []string { return s.list(s.scope) } + +// BanEviction marks a blob as unevictable by LRU eviction. It is idempotent. +// Needed when e.g. blobs must be written back to GCS/S3 and eviction before that is unacceptable. +func (s *DiskStore) BanEviction(key string) error { return s.banEviction(key, s.scope) } + +// UnbanEviction removes the effect of BanEviction for a blob. It is idempotent. +func (s *DiskStore) UnbanEviction(key string) error { return s.unbanEviction(key, s.scope) } + +// SetMetadata atomically sets the respective metadata for a blob. +func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { + return s.setMetadata(key, md, s.scope) +} + +// GetMetadata populates `md` if the metadata is present. Returns [os.ErrNotExists] if key is not in store. +func (s *DiskStore) GetMetadata(key string, md metadata.Metadata) (ok bool, err error) { + return s.getMetadata(key, md, s.scope) +} + +// DeleteMetadata removes any metadata of a blob with `md`'s suffix, if present. +func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { + return s.deleteMetadata(key, md, s.scope) +} + +// ScopeComplete scopes [DiskStore]'s APIs such that they can only operate on complete blobs. +// [ErrOutOfScope] is returned if the user tries to operate on an incomplete blob. +func (s *DiskStore) ScopeComplete() *DiskStore { return &DiskStore{s.diskStore, blobScopeComplete} } + +// ScopeIncomplete scopes [DiskStore]'s APIs such that they can only operate on incomplete blobs. +// [ErrOutOfScope] is returned if the user tries to operate on a complete blob. +func (s *DiskStore) ScopeIncomplete() *DiskStore { return &DiskStore{s.diskStore, blobScopeIncomplete} } From 223422730a3bf4a867a02df649de82de198df6e0 Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Thu, 23 Jul 2026 20:44:21 +0200 Subject: [PATCH 08/10] feat(DiskStore): fully implement Metadata ops - 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 --- lib/store/crash_recovery.go | 6 +- lib/store/crash_recovery_test.go | 59 ++++++---- lib/store/disk_store.go | 131 +++++++++++++++++++--- lib/store/disk_store_test.go | 186 +++++++++++++++++++++++++++---- lib/store/metadata/metadata.go | 1 - lib/store/pather.go | 2 +- lib/store/pather_test.go | 7 +- lib/store/scoped_store.go | 21 +++- 8 files changed, 347 insertions(+), 66 deletions(-) diff --git a/lib/store/crash_recovery.go b/lib/store/crash_recovery.go index e7b177844..47d617661 100644 --- a/lib/store/crash_recovery.go +++ b/lib/store/crash_recovery.go @@ -13,6 +13,7 @@ import ( "strings" "time" + "github.com/uber-go/tally" "github.com/uber/kraken/utils/closers" "go.uber.org/zap" ) @@ -27,7 +28,7 @@ type blobState struct { complete bool } -func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, log *zap.SugaredLogger) (*diskStore, error) { +func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, log *zap.SugaredLogger, metrics tally.Scope) (*diskStore, error) { completeDirPath, incompleteDirPath := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) if !rebootIncompleteBlobs { err := os.RemoveAll(incompleteDirPath) @@ -59,6 +60,7 @@ func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncomplete return nil, err } if !ok { + log.Warn("Could not reboot blob from disk - its parent directory is there but the blob is missing") continue } if bState.complete && bState.evictable { @@ -104,6 +106,7 @@ func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncomplete size: storeSize, rebootIncompleteBlobs: rebootIncompleteBlobs, log: log, + metrics: metrics, } if store.size > store.capacity { @@ -181,6 +184,7 @@ func rebootIncompleteBlobSize(key string, pather *pather) (size uint64, ok bool, } func rebootKeys(subDir string) ([]string, error) { + // TODO - consider using glob matching to reboot the keys keys := make([]string, 0) ok, err := exists(subDir) if err != nil { diff --git a/lib/store/crash_recovery_test.go b/lib/store/crash_recovery_test.go index 9ede66a82..c84cb9487 100644 --- a/lib/store/crash_recovery_test.go +++ b/lib/store/crash_recovery_test.go @@ -4,10 +4,10 @@ import ( "bytes" "crypto/rand" "io" - "os" "testing" "github.com/stretchr/testify/require" + "github.com/uber-go/tally" "github.com/uber/kraken/core" "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/memsize" @@ -40,7 +40,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.BanEviction(incompleteUnevictableKey)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) require.NoError(err) // Incomplete files are recovered (since `rebootIncompleteBlobs` is true). @@ -89,14 +89,13 @@ func TestCrashRecovery(t *testing.T) { require.Equal(8*memsize.KB, store.size) // Run the store with `rebootIncompleteBlobs` as false. - store, err = NewDiskStore(10*memsize.KB, rootDir, false) + store, err = NewDiskStore(10*memsize.KB, rootDir, false, tally.NoopScope) require.NoError(err) // Incomplete files are dropped. - _, err = store.Stat(incompleteEvictableKey) - require.ErrorIs(err, os.ErrNotExist) - _, err = store.Stat(incompleteUnevictableKey) - require.ErrorIs(err, os.ErrNotExist) + rebootedKeys := store.List() + require.NotContains(rebootedKeys, incompleteEvictableKey) + require.NotContains(rebootedKeys, incompleteUnevictableKey) // Complete files are always recovered. f, err = store.ScopeComplete().Open(completeEvictableKey) @@ -133,7 +132,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.SetMetadata(key, writtenMd)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) require.NoError(err) var readMd metadata.TorrentMeta @@ -153,7 +152,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.SetMetadata(key, writtenMd)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) require.NoError(err) var readMd metadata.TorrentMeta @@ -201,7 +200,7 @@ func TestCrashRecovery(t *testing.T) { require.Equal([]string{cKey, eKey, dKey}, evictionOrderBeforeCrash) // a is unevictable and b is incomplete // Assume that the application restarts here. - store, err = NewDiskStore(10*memsize.KB, rootDir, false) + store, err = NewDiskStore(10*memsize.KB, rootDir, false, tally.NoopScope) require.NoError(err) // LRU order is approximated, but not exact. @@ -211,16 +210,37 @@ func TestCrashRecovery(t *testing.T) { require.Equal(wantEvictionOrder, rebootedEvictionOrder) // Assume we redeploy the service with a smaller capacity for the disk store: - rebootedSmallerStore, err := NewDiskStore(6*memsize.KB, rootDir, false) + store, err = NewDiskStore(6*memsize.KB, rootDir, false, tally.NoopScope) require.NoError(err) // since 10KB of blobs are in store, `c` gets evicted to put the store back within its capacity. - _, err = rebootedSmallerStore.Stat(cKey) - require.ErrorIs(err, os.ErrNotExist) + require.NotContains(store.List(), cKey) - require.Equal([]string{dKey, eKey}, rebootedSmallerStore.evictionOrder()) - _, err = rebootedSmallerStore.Stat(aKey) - require.NoError(err) + require.Equal([]string{dKey, eKey}, store.evictionOrder()) + require.Contains(store.List(), aKey) + }) + + t.Run("reboot fails if unevictable data is more than store capacity", func(t *testing.T) { + require := require.New(t) + store, rootDir := newTestStore(t, 10*memsize.KB, true) + + aF, aKey := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, aF, 2*memsize.KB) + require.NoError(aF.Close()) + require.NoError(store.MarkComplete(aKey)) + require.NoError(store.BanEviction(aKey)) + + bF, _ := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, bF, 1*memsize.KB) + require.NoError(bF.Close()) + + cF, _ := newTestFile(t, store, 2*memsize.KB) + _ = fillWithRandomData(t, cF, 1*memsize.KB) + require.NoError(cF.Close()) + + // Assume that the application restarts here. + store, err := NewDiskStore(5*memsize.KB, rootDir, true, tally.NoopScope) + require.ErrorContains(err, "cannot evict enough, the unevictable/incomplete blobs are using up all the space") }) } @@ -232,7 +252,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { firstData := fillWithRandomData(t, f, 2*memsize.KB) // First crash. - store, err := NewDiskStore(10*memsize.KB, rootDir, true) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) require.NoError(err) f, err = store.Open(key) @@ -246,7 +266,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { require.NoError(err) // Second crash. - store, err = NewDiskStore(10*memsize.KB, rootDir, true) + store, err = NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) require.NoError(err) wantData := make([]byte, 5*memsize.KB) @@ -263,7 +283,8 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { require.NoError(store.MarkComplete(key)) // Third crash. - store, err = NewDiskStore(10*memsize.KB, rootDir, true) + store, err = NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) + require.NoError(err) // now that the blob is complete, its actual size should be rebooted through stat, instead of trusting the _size sidecar file. require.Equal(5*memsize.KB, store.size) f, err = store.ScopeComplete().Open(key) diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index 085208fd0..40b6f69a7 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -9,7 +9,9 @@ import ( "path/filepath" "strconv" "sync" + "time" + "github.com/uber-go/tally" "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/closers" "github.com/uber/kraken/utils/log" @@ -34,6 +36,8 @@ const ( _blobSizeFileName = "_size" ) +var _syncEvictionLatencyBuckets = tally.MustMakeExponentialDurationBuckets(100*time.Millisecond, 1.4, 15) + // diskStore implements the APIs of [DiskStore]. [diskStore]'s APIs expose the [blobScope] arg, // while [DiskStore]'s APIs omit that arg (cleaner interface) and instead expose other APIs to scope the whole store. // @@ -49,6 +53,7 @@ type diskStore struct { // If disabled, incomplete blobs are discarded (usually done to prevent leaks). rebootIncompleteBlobs bool log *zap.SugaredLogger + metrics tally.Scope *pather } @@ -59,7 +64,7 @@ type blob struct { evictionBanned bool } -func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool) (*diskStore, error) { +func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, metrics tally.Scope) (*diskStore, error) { // TODO - create a Config struct. // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. // TODO - move disk store files into their own directory and package. @@ -79,12 +84,13 @@ func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bo blobs: make(map[string]*blob), evictQueue: list.New(), log: log, + metrics: metrics, pather: newPather(rootDir), rebootIncompleteBlobs: rebootIncompleteBlobs, }, nil } - store, err := rebootPersistedStore(capacityBytes, rootDir, rebootIncompleteBlobs, log) + store, err := rebootPersistedStore(capacityBytes, rootDir, rebootIncompleteBlobs, log, metrics) if err != nil { err = fmt.Errorf("reboot persisted state into memory: %w", err) log.With("error", err).Error("Failed to initialize disk store") @@ -157,8 +163,7 @@ func (s *diskStore) create(key string, sizeBytes uint64) (FileReadWriter, error) return nil, fmt.Errorf("ensure dir: %w", err) } blobPath := s.blobPath(key, _incompleteBlob) - flag := os.O_RDWR | os.O_CREATE | os.O_EXCL - f, err := os.OpenFile(blobPath, flag, _defaultFilePerm) + f, err := os.OpenFile(blobPath, os.O_RDWR|os.O_CREATE|os.O_EXCL, _defaultFilePerm) if err != nil { s.releaseSpace(sizeBytes) return nil, fmt.Errorf("open file: %w", err) @@ -188,19 +193,24 @@ func (s *diskStore) persistBlobSize(key string, sizeBytes uint64) error { if err != nil { return fmt.Errorf("create file: %w", err) } + defer closers.Close(blobSizeF) _, err = blobSizeF.Write([]byte(strconv.Itoa(int(sizeBytes)))) if err != nil { return fmt.Errorf("write to size file: %w", err) } - closers.Close(blobSizeF) return nil } func (s *diskStore) reserveSpace(space uint64) error { // TODO - benchmark and consider whether async eviction makes more sense. - // TODO - emit latency to reserve space for a blob. + startTime := time.Now() for s.size+space > s.capacity { if s.evictQueue.Len() == 0 { + s.log.With( + "unevictable_bytes", s.size, + "required_space", space, + "capacity", s.capacity, + ).Error("Cannot evict enough data to free space for new entry to DiskStore. The unevictable/incomplete blobs are using up all the space") return errors.New("cannot evict enough, the unevictable/incomplete blobs are using up all the space") } @@ -209,7 +219,6 @@ func (s *diskStore) reserveSpace(space uint64) error { err := s.deleteFromDisk(toEvictKey, _completeBlob) if err != nil { - // TODO - consider whether we want to fail-open by doing `continue` here. return fmt.Errorf("delete from disk: %w", err) } s.evictQueue.Remove(toEvictNode) @@ -217,8 +226,10 @@ func (s *diskStore) reserveSpace(space uint64) error { s.releaseSpace(size) delete(s.blobs, toEvictKey) } - s.size += space + + latency := time.Since(startTime) + s.metrics.Histogram("sync_eviction_latency", _syncEvictionLatencyBuckets).RecordDuration(latency) return nil } @@ -256,7 +267,6 @@ func (s *diskStore) markComplete(key string) error { if err != nil { return fmt.Errorf("mkdirall: %w", err) } - // TODO - make sure that un-movable metadata is deleted after move err = os.Rename(oldPathDir, newPathDir) if err != nil { return fmt.Errorf("move dir: %w", err) @@ -266,9 +276,35 @@ func (s *diskStore) markComplete(key string) error { node := s.evictQueue.PushBack(key) b.node = node } + + s.tryDeleteImmovableMetadata(key) return nil } +// Best-effort attempt to remove immovable metadata. Fail-open on failure +// to avoid an inconsistent state, as failure only costs a negligible amount +// of disk until the blob is evicted. +func (s *diskStore) tryDeleteImmovableMetadata(key string) { + mdList, err := s.listMetadataNoLock(key, blobScopeAny) + if err != nil { + err = fmt.Errorf("list metadata: %w", err) + s.log.With("error", err).Error("Failed to delete un-movable metadata upon marking a blob as complete") + return + } + for _, md := range mdList { + if md.Movable() { + continue + } + mdFilePath := s.sidecarFilePath(key, _completeBlob, md.GetSuffix()) + err = os.Remove(mdFilePath) + if err != nil && !errors.Is(err, os.ErrNotExist) { + err = fmt.Errorf("remove metadata file: %w", err) + s.log.With("error", err).Error("Failed to delete un-movable metadata upon marking a blob as complete") + continue + } + } +} + func (s *diskStore) checkDiskIfUnevictable(key string, complete bool) (bool, error) { flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) unevictable, err := exists(flagBlobPath) @@ -396,20 +432,17 @@ func (s *diskStore) setMetadata(key string, md metadata.Metadata, scope blobScop return fmt.Errorf("serialize metadata: %w", err) } mdFilePath := s.sidecarFilePath(key, b.complete, md.GetSuffix()) - // We use a tmp file to ensure atomicity. + // Use a tmp file to ensure atomicity. tmpFilePath := mdFilePath + "-tmp" tmpFile, err := os.OpenFile(tmpFilePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, _defaultFilePerm) if err != nil { return fmt.Errorf("create tmp file for md: %w", err) } + defer closers.Close(tmpFile) _, err = tmpFile.Write(mdData) if err != nil { return fmt.Errorf("write to tmp file: %w", err) } - err = tmpFile.Close() - if err != nil { - return fmt.Errorf("close tmp file: %w", err) - } err = os.Rename(tmpFile.Name(), mdFilePath) if err != nil { return fmt.Errorf("rename tmp file: %w", err) @@ -434,6 +467,9 @@ func (s *diskStore) getMetadata(key string, md metadata.Metadata, scope blobScop if errors.Is(err, os.ErrNotExist) { return false, nil } + if err != nil { + return false, fmt.Errorf("open metadata file: %w", err) + } defer closers.Close(mdFile) data, err := io.ReadAll(mdFile) if err != nil { @@ -469,6 +505,73 @@ func (s *diskStore) deleteMetadata(key string, md metadata.Metadata, scope blobS return nil } +func (s *diskStore) listMetadata(key string, scope blobScope) ([]metadata.Metadata, error) { + s.mu.RLock() + defer s.mu.RUnlock() + + return s.listMetadataNoLock(key, scope) +} + +func (s *diskStore) listMetadataNoLock(key string, scope blobScope) ([]metadata.Metadata, error) { + b, ok := s.blobs[key] + if !ok { + return nil, os.ErrNotExist + } + if err := isOutOfScope(b, scope); err != nil { + return nil, err + } + + res := make([]metadata.Metadata, 0) + entries, err := os.ReadDir(s.dirPath(key, b.complete)) + if err != nil { + return nil, fmt.Errorf("read metadata files from dir: %w", err) + } + for _, entry := range entries { + name := entry.Name() + if name == _blobFileName || name == _evictionBannedFileName || name == _blobSizeFileName { + continue + } + + md := metadata.CreateFromSuffix(name) + if md == nil { + s.log.With("key", key, "file_name", name).Warn("Found file in blob dir that does not successfully parse as a metadata file") + continue + } + res = append(res, md) + } + return res, nil +} + +func (s *diskStore) writeAtMetadata(key string, md metadata.Metadata, p []byte, off int64, scope blobScope) error { + s.mu.Lock() + defer s.mu.Unlock() + + b, ok := s.blobs[key] + if !ok { + return os.ErrNotExist + } + if err := isOutOfScope(b, scope); err != nil { + return err + } + + mdFilePath := s.sidecarFilePath(key, b.complete, md.GetSuffix()) + mdFile, err := os.OpenFile(mdFilePath, os.O_WRONLY, _defaultFilePerm) + if errors.Is(err, os.ErrNotExist) { + return errors.New("metadata does not exist") + } + if err != nil { + return fmt.Errorf("open: %w", err) + } + defer closers.Close(mdFile) + + _, err = mdFile.WriteAt(p, off) + if err != nil { + return fmt.Errorf("write at: %w", err) + } + + return nil +} + // used during testing func (s *diskStore) evictionOrder() []string { s.mu.RLock() diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index 3c86b60c5..d39168d10 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -2,11 +2,11 @@ package store import ( "bytes" - "errors" "io" "io/fs" "os" "path/filepath" + "regexp" "strings" "sync" "testing" @@ -14,6 +14,7 @@ import ( "time" "github.com/stretchr/testify/require" + "github.com/uber-go/tally" "github.com/uber/kraken/core" "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/memsize" @@ -22,9 +23,9 @@ import ( func newTestStore(t *testing.T, capacity uint64, rebootIncompleteBlobs bool) (res *DiskStore, rootDir string) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) - t.Cleanup(func() { os.RemoveAll(rootDir) }) + t.Cleanup(func() { _ = os.RemoveAll(rootDir) }) - store, err := NewDiskStore(capacity, rootDir, rebootIncompleteBlobs) + store, err := NewDiskStore(capacity, rootDir, rebootIncompleteBlobs, tally.NoopScope) require.NoError(t, err) return store, rootDir } @@ -275,7 +276,7 @@ func TestParallelAccessToSingleFile(t *testing.T) { wg.Wait() for idx := range 5 { require.NoError(results[idx].err) - require.Equal(results[idx].written, results[idx].written) + require.Equal(results[idx].written, results[idx].read) } require.NoError(store.MarkComplete(key)) @@ -323,13 +324,6 @@ func TestOpenedFileAccessibleAfterMarkedComplete(t *testing.T) { require.Equal([]byte("Hello World"), completeFileData) } -func TestOpenedFileAccessibleAfterEviction(t *testing.T) { - // when a file is created, then received with a Get, then deleted, check that the store doesn't count its size but it is still on disk and accessible. then check that when the client closes the file, it is not accessible anymore and not on disk anymore. - // test after both "Delete" and natural LRU eviction. - // what happens when we try to close the fd? i assume it works? - t.Skip("TODO") -} - func TestDelete(t *testing.T) { t.Run("incomplete blob", func(t *testing.T) { require := require.New(t) @@ -404,7 +398,7 @@ func TestDelete(t *testing.T) { key := core.DigestFixture().Hex() err := store.Delete(key) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) }) } @@ -477,7 +471,7 @@ func TestMarkComplete(t *testing.T) { key := core.DigestFixture().Hex() err := store.MarkComplete(key) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) }) } @@ -577,9 +571,9 @@ func TestStat(t *testing.T) { key := core.DigestFixture().Hex() _, err := store.ScopeComplete().Stat(key) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) _, err = store.Stat(key) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) }) } @@ -626,6 +620,19 @@ func TestList(t *testing.T) { require.ElementsMatch(wantRes, res) } +func init() { + metadata.Register(regexp.MustCompile("immovableMd"), &immovableMdFactory{}) +} + +// used for testing +type immovableMd struct{} +type immovableMdFactory struct{} + +func (f *immovableMdFactory) Create(suffix string) metadata.Metadata { return &immovableMd{} } +func (mda *immovableMd) GetSuffix() string { return "immovableMd" } +func (mda *immovableMd) Movable() bool { return false } +func (mda *immovableMd) Serialize() ([]byte, error) { return nil, nil } +func (mda *immovableMd) Deserialize(b []byte) error { return nil } func TestMetadata(t *testing.T) { t.Run("basic functionality", func(t *testing.T) { require := require.New(t) @@ -647,6 +654,27 @@ func TestMetadata(t *testing.T) { require.True(ok) require.Equal(writtenMd.MetaInfo, readMd.MetaInfo) + mdList, err := store.ListMetadata(key) + require.NoError(err) + require.Len(mdList, 1) + require.Equal(writtenMd.GetSuffix(), mdList[0].GetSuffix()) + + persistMd := metadata.NewPersist(true) + require.NoError(store.SetMetadata(key, persistMd)) + require.NoError(store.WriteAtMetadata(key, persistMd, []byte("false"), 0)) + var readPersistMd metadata.Persist + ok, err = store.GetMetadata(key, &readPersistMd) + require.NoError(err) + require.True(ok) + require.False(readPersistMd.Value) + + mdList, err = store.ListMetadata(key) + require.NoError(err) + require.ElementsMatch( + []string{writtenMd.GetSuffix(), persistMd.GetSuffix()}, + []string{mdList[0].GetSuffix(), mdList[1].GetSuffix()}, + ) + require.NoError(store.DeleteMetadata(key, &readMd)) ok, err = store.GetMetadata(key, &readMd) require.NoError(err) @@ -654,9 +682,19 @@ func TestMetadata(t *testing.T) { mdFilePath := store.sidecarFilePath(key, _incompleteBlob, readMd.GetSuffix()) // ensure the metadata file is deleted from disk _, err = os.Stat(mdFilePath) - require.True(errors.Is(err, os.ErrNotExist)) + require.ErrorIs(err, os.ErrNotExist) // deleting a second time should be a no-op. require.NoError(store.DeleteMetadata(key, &readMd)) + + mdList, err = store.ListMetadata(key) + require.NoError(err) + require.Len(mdList, 1) + require.Equal(persistMd.GetSuffix(), mdList[0].GetSuffix()) + + require.NoError(store.DeleteMetadata(key, persistMd)) + mdList, err = store.ListMetadata(key) + require.NoError(err) + require.Empty(mdList) }) t.Run("non-existent blob", func(t *testing.T) { @@ -667,14 +705,20 @@ func TestMetadata(t *testing.T) { md := metadata.NewTorrentMeta(mdStruct) err := store.SetMetadata(nonExistentKey, md) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) ok, err := store.GetMetadata(nonExistentKey, md) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) require.False(ok) + _, err = store.ListMetadata(nonExistentKey) + require.ErrorIs(os.ErrNotExist, err) + + err = store.WriteAtMetadata(nonExistentKey, md, []byte("data"), 0) + require.ErrorIs(os.ErrNotExist, err) + err = store.DeleteMetadata(nonExistentKey, md) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) }) t.Run("metadata does not change after marking a file as complete and/or evictable/unevictable", func(t *testing.T) { @@ -740,10 +784,80 @@ func TestMetadata(t *testing.T) { defer func() { require.NoError(fB.Close()) }() ok, err := store.GetMetadata(keyA, md) - require.Equal(os.ErrNotExist, err) + require.ErrorIs(os.ErrNotExist, err) require.False(ok) _, err = os.Stat(mdFilePath) - require.True(errors.Is(err, os.ErrNotExist)) + require.ErrorIs(err, os.ErrNotExist) + }) + + t.Run("immovable metadata is deleted when calling MarkComplete", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB, false) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 10*memsize.KB) + require.NoError(err) + require.NoError(f.Close()) + + md := &immovableMd{} + require.NoError(store.SetMetadata(key, md)) + movableMd := metadata.NewTorrentMeta(core.MetaInfoFixture()) + require.NoError(store.SetMetadata(key, movableMd)) + + require.NoError(store.MarkComplete(key)) + readMd := immovableMd{} + ok, err := store.GetMetadata(key, &readMd) + require.NoError(err) + require.False(ok) + + readMovableMd := metadata.TorrentMeta{} + ok, err = store.GetMetadata(key, &readMovableMd) + require.NoError(err) + require.True(ok) + }) + + t.Run("writeAt on metadata that does not exist", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB, false) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 10*memsize.KB) + require.NoError(err) + require.NoError(f.Close()) + + md := metadata.NewPersist(true) + err = store.WriteAtMetadata(key, md, []byte("true"), 0) + require.EqualError(err, "metadata does not exist") + }) + + t.Run("list metadata excludes non-metadata sidecar files", func(t *testing.T) { + require := require.New(t) + store, _ := newTestStore(t, 10*memsize.KB, true) + key := core.DigestFixture().Hex() + f, err := store.Create(key, 10*memsize.KB) + require.NoError(err) + require.NoError(f.Close()) + require.NoError(store.BanEviction(key)) + + mdList, err := store.ListMetadata(key) + require.NoError(err) + require.Empty(mdList) + + torrentMd := metadata.NewTorrentMeta(core.MetaInfoFixture()) + persistMd := metadata.NewPersist(true) + lastAccessMd := metadata.NewLastAccessTime(time.Now()) + require.NoError(store.SetMetadata(key, torrentMd)) + require.NoError(store.SetMetadata(key, persistMd)) + require.NoError(store.SetMetadata(key, lastAccessMd)) + + mdList, err = store.ListMetadata(key) + require.NoError(err) + gotSuffixes := make([]string, len(mdList)) + for i, md := range mdList { + gotSuffixes[i] = md.GetSuffix() + } + require.ElementsMatch( + []string{torrentMd.GetSuffix(), persistMd.GetSuffix(), lastAccessMd.GetSuffix()}, + gotSuffixes, + ) }) } @@ -755,9 +869,11 @@ func TestScopes(t *testing.T) { data := fillWithRandomData(t, f, 2*memsize.KB) require.NoError(f.Close()) md := metadata.NewTorrentMeta(core.MetaInfoFixture()) + mdData, err := md.Serialize() + require.NoError(err) // While incomplete, ScopeComplete's APIs reject the blob. - _, err := store.ScopeComplete().Open(key) + _, err = store.ScopeComplete().Open(key) require.ErrorIs(err, ErrOutOfScope) _, err = store.ScopeComplete().Stat(key) require.ErrorIs(err, ErrOutOfScope) @@ -768,6 +884,9 @@ func TestScopes(t *testing.T) { ok, err := store.ScopeComplete().GetMetadata(key, &readMd) require.ErrorIs(err, ErrOutOfScope) require.False(ok) + _, err = store.ScopeComplete().ListMetadata(key) + require.ErrorIs(err, ErrOutOfScope) + require.ErrorIs(store.ScopeComplete().WriteAtMetadata(key, md, mdData, 0), ErrOutOfScope) require.ErrorIs(store.ScopeComplete().DeleteMetadata(key, &readMd), ErrOutOfScope) require.NotContains(store.ScopeComplete().List(), key) @@ -787,6 +906,11 @@ func TestScopes(t *testing.T) { require.NoError(err) require.True(ok) require.Equal(md.MetaInfo, readMd.MetaInfo) + mdList, err := store.ListMetadata(key) + require.NoError(err) + require.Len(mdList, 1) + require.Equal(md.GetSuffix(), mdList[0].GetSuffix()) + require.NoError(store.WriteAtMetadata(key, md, mdData, 0)) require.NoError(store.DeleteMetadata(key, &readMd)) require.Contains(store.List(), key) @@ -806,6 +930,11 @@ func TestScopes(t *testing.T) { require.NoError(err) require.True(ok) require.Equal(md.MetaInfo, readMd.MetaInfo) + mdList, err = store.ScopeIncomplete().ListMetadata(key) + require.NoError(err) + require.Len(mdList, 1) + require.Equal(md.GetSuffix(), mdList[0].GetSuffix()) + require.NoError(store.ScopeIncomplete().WriteAtMetadata(key, md, mdData, 0)) require.NoError(store.ScopeIncomplete().DeleteMetadata(key, &readMd)) require.Contains(store.ScopeIncomplete().List(), key) @@ -822,6 +951,9 @@ func TestScopes(t *testing.T) { ok, err = store.ScopeIncomplete().GetMetadata(key, &readMd) require.ErrorIs(err, ErrOutOfScope) require.False(ok) + _, err = store.ScopeIncomplete().ListMetadata(key) + require.ErrorIs(err, ErrOutOfScope) + require.ErrorIs(store.ScopeIncomplete().WriteAtMetadata(key, md, mdData, 0), ErrOutOfScope) require.ErrorIs(store.ScopeIncomplete().DeleteMetadata(key, &readMd), ErrOutOfScope) require.NotContains(store.ScopeIncomplete().List(), key) @@ -841,6 +973,11 @@ func TestScopes(t *testing.T) { require.NoError(err) require.True(ok) require.Equal(md.MetaInfo, readMd.MetaInfo) + mdList, err = store.ListMetadata(key) + require.NoError(err) + require.Len(mdList, 1) + require.Equal(md.GetSuffix(), mdList[0].GetSuffix()) + require.NoError(store.WriteAtMetadata(key, md, mdData, 0)) require.NoError(store.DeleteMetadata(key, &readMd)) require.Contains(store.List(), key) @@ -860,6 +997,11 @@ func TestScopes(t *testing.T) { require.NoError(err) require.True(ok) require.Equal(md.MetaInfo, readMd.MetaInfo) + mdList, err = store.ScopeComplete().ListMetadata(key) + require.NoError(err) + require.Len(mdList, 1) + require.Equal(md.GetSuffix(), mdList[0].GetSuffix()) + require.NoError(store.ScopeComplete().WriteAtMetadata(key, md, mdData, 0)) require.NoError(store.ScopeComplete().DeleteMetadata(key, &readMd)) require.Contains(store.ScopeComplete().List(), key) } diff --git a/lib/store/metadata/metadata.go b/lib/store/metadata/metadata.go index 1c59ca07f..c8861a125 100644 --- a/lib/store/metadata/metadata.go +++ b/lib/store/metadata/metadata.go @@ -37,7 +37,6 @@ func Register(suffix *regexp.Regexp, factory Factory) { } // CreateFromSuffix creates a Metadata obj based on suffix. -// This is not a very efficient method; It's mostly used during reload. func CreateFromSuffix(suffix string) Metadata { for re, factory := range _factories { if re.MatchString(suffix) { diff --git a/lib/store/pather.go b/lib/store/pather.go index d4d505b4d..5ba18be0e 100644 --- a/lib/store/pather.go +++ b/lib/store/pather.go @@ -32,7 +32,7 @@ func (p *pather) dirPath(key string, complete bool) string { } dirPath := filepath.Join(p.dir, subDirName) for i := 0; i < int(_defaultShardIDLength) && i < len(key)/2; i++ { - // (1 byte = 2 char of file name assumming file name is in HEX) + // (1 byte = 2 char of file name assuming file name is in HEX) dirName := key[i*2 : i*2+2] dirPath = filepath.Join(dirPath, dirName) } diff --git a/lib/store/pather_test.go b/lib/store/pather_test.go index 1a17d2a85..2c239a227 100644 --- a/lib/store/pather_test.go +++ b/lib/store/pather_test.go @@ -10,15 +10,15 @@ import ( ) func TestPather(t *testing.T) { - require := require.New(t) rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") - require.NoError(err) - t.Cleanup(func() { os.RemoveAll(rootDir) }) + require.NoError(t, err) + t.Cleanup(func() { _ = os.RemoveAll(rootDir) }) pather := newPather(rootDir) md := metadata.NewTorrentMeta(core.MetaInfoFixture()) key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" t.Run("incomplete entry", func(t *testing.T) { + require := require.New(t) complete := false dirPath := pather.dirPath(key, complete) wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" @@ -32,6 +32,7 @@ func TestPather(t *testing.T) { }) t.Run("complete entry", func(t *testing.T) { + require := require.New(t) complete := true dirPath := pather.dirPath(key, complete) wantDirPath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" diff --git a/lib/store/scoped_store.go b/lib/store/scoped_store.go index 5cad968ca..fe2514141 100644 --- a/lib/store/scoped_store.go +++ b/lib/store/scoped_store.go @@ -4,6 +4,7 @@ import ( "errors" "os" + "github.com/uber-go/tally" "github.com/uber/kraken/lib/store/metadata" ) @@ -36,8 +37,8 @@ type DiskStore struct { // // - If the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), // it evicts blobs until size is within capacity. -func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool) (*DiskStore, error) { - diskStore, err := newDiskStore(capacityBytes, rootDir, rebootIncompleteBlobs) +func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, metrics tally.Scope) (*DiskStore, error) { + diskStore, err := newDiskStore(capacityBytes, rootDir, rebootIncompleteBlobs, metrics) if err != nil { return nil, err } @@ -54,10 +55,10 @@ func (s *diskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) return s.create(key, sizeBytes) } -// Open returns an FD to a file in the store. [os.ErrNotExists] is returned on missing entry. +// Open returns an FD to a file in the store. [os.ErrNotExist] is returned on missing entry. func (s *DiskStore) Open(key string) (FileReadWriter, error) { return s.open(key, s.scope) } -// Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExists] if the blob is not found. +// Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExist] if the blob is not found. func (s *DiskStore) Stat(key string) (os.FileInfo, error) { return s.stat(key, s.scope) } // MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless BanEviction has been called). @@ -82,7 +83,7 @@ func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { return s.setMetadata(key, md, s.scope) } -// GetMetadata populates `md` if the metadata is present. Returns [os.ErrNotExists] if key is not in store. +// GetMetadata populates `md` if the metadata is present. Returns [os.ErrNotExist] if key is not in store. func (s *DiskStore) GetMetadata(key string, md metadata.Metadata) (ok bool, err error) { return s.getMetadata(key, md, s.scope) } @@ -92,6 +93,16 @@ func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { return s.deleteMetadata(key, md, s.scope) } +// ListMetadata returns all [metadata.Metadata] of key. +func (s *DiskStore) ListMetadata(key string) ([]metadata.Metadata, error) { + return s.listMetadata(key, s.scope) +} + +// WriteAtMetadata implements [io.WriterAt] for the metadata file on disk. +func (s *DiskStore) WriteAtMetadata(key string, md metadata.Metadata, p []byte, off int64) error { + return s.writeAtMetadata(key, md, p, off, s.scope) +} + // ScopeComplete scopes [DiskStore]'s APIs such that they can only operate on complete blobs. // [ErrOutOfScope] is returned if the user tries to operate on an incomplete blob. func (s *DiskStore) ScopeComplete() *DiskStore { return &DiskStore{s.diskStore, blobScopeComplete} } From 3474cab20e60daeea97608b9efa4c672ae32ee9b Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Fri, 24 Jul 2026 16:28:07 +0200 Subject: [PATCH 09/10] feat(DiskStore): Make dir-sharding configurable 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. --- lib/store/crash_recovery.go | 49 ++---------------------- lib/store/crash_recovery_test.go | 20 +++++----- lib/store/disk_store.go | 6 +-- lib/store/disk_store_test.go | 6 ++- lib/store/pather.go | 66 ++++++++++++++++++++++++++------ lib/store/pather_test.go | 54 +++++++++++++++----------- lib/store/scoped_store.go | 4 +- 7 files changed, 109 insertions(+), 96 deletions(-) diff --git a/lib/store/crash_recovery.go b/lib/store/crash_recovery.go index 47d617661..5d9f3851a 100644 --- a/lib/store/crash_recovery.go +++ b/lib/store/crash_recovery.go @@ -5,12 +5,10 @@ import ( "errors" "fmt" "io" - "io/fs" "os" "path/filepath" "slices" "strconv" - "strings" "time" "github.com/uber-go/tally" @@ -18,8 +16,6 @@ import ( "go.uber.org/zap" ) -const _numShards = 2 // TODO - change this not to be hardcoded when configurable sharding is implemented - type blobState struct { key string size uint64 @@ -28,7 +24,7 @@ type blobState struct { complete bool } -func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, log *zap.SugaredLogger, metrics tally.Scope) (*diskStore, error) { +func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, shardLength int, log *zap.SugaredLogger, metrics tally.Scope) (*diskStore, error) { completeDirPath, incompleteDirPath := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) if !rebootIncompleteBlobs { err := os.RemoveAll(incompleteDirPath) @@ -37,20 +33,20 @@ func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncomplete } } - keys, err := rebootKeys(completeDirPath) + pather := newPather(rootDir, shardLength) + keys, err := pather.rebootKeys(completeDirPath) if err != nil { return nil, err } numCompleteBlobs := len(keys) if rebootIncompleteBlobs { - incompleteKeys, err := rebootKeys(incompleteDirPath) + incompleteKeys, err := pather.rebootKeys(incompleteDirPath) if err != nil { return nil, err } keys = append(keys, incompleteKeys...) } - pather := newPather(rootDir) completeEvictableEntries := make([]*blobState, 0) otherEntries := make([]*blobState, 0) for i, key := range keys { @@ -183,43 +179,6 @@ func rebootIncompleteBlobSize(key string, pather *pather) (size uint64, ok bool, return uint64(blobSize), true, nil } -func rebootKeys(subDir string) ([]string, error) { - // TODO - consider using glob matching to reboot the keys - keys := make([]string, 0) - ok, err := exists(subDir) - if err != nil { - return nil, fmt.Errorf("exists: %w", err) - } - if !ok { - return []string{}, nil - } - err = filepath.WalkDir(subDir, func(path string, entry fs.DirEntry, err error) error { - if err != nil { - return err - } - - if !entry.IsDir() { - return nil - } - relPath, err := filepath.Rel(subDir, path) - if err != nil { - return err - } - nameParts := strings.Split(relPath, string(filepath.Separator)) - isBlobDir := len(nameParts) == _numShards+1 - if !isBlobDir { - return nil - } - key := nameParts[len(nameParts)-1] - keys = append(keys, key) - return fs.SkipDir - }) - if err != nil { - return nil, fmt.Errorf("walk through subdir '%v' to reboot blob keys: %w", subDir, err) - } - return keys, nil -} - func existsPersistedStore(rootDir string) (ok bool, err error) { completeDir, incompleteDir := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) completeExists, err := exists(completeDir) diff --git a/lib/store/crash_recovery_test.go b/lib/store/crash_recovery_test.go index c84cb9487..b125aea31 100644 --- a/lib/store/crash_recovery_test.go +++ b/lib/store/crash_recovery_test.go @@ -40,7 +40,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.BanEviction(incompleteUnevictableKey)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) require.NoError(err) // Incomplete files are recovered (since `rebootIncompleteBlobs` is true). @@ -89,7 +89,7 @@ func TestCrashRecovery(t *testing.T) { require.Equal(8*memsize.KB, store.size) // Run the store with `rebootIncompleteBlobs` as false. - store, err = NewDiskStore(10*memsize.KB, rootDir, false, tally.NoopScope) + store, err = NewDiskStore(10*memsize.KB, rootDir, false, _defaultShardLength, tally.NoopScope) require.NoError(err) // Incomplete files are dropped. @@ -132,7 +132,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.SetMetadata(key, writtenMd)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) require.NoError(err) var readMd metadata.TorrentMeta @@ -152,7 +152,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(store.SetMetadata(key, writtenMd)) // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) require.NoError(err) var readMd metadata.TorrentMeta @@ -200,7 +200,7 @@ func TestCrashRecovery(t *testing.T) { require.Equal([]string{cKey, eKey, dKey}, evictionOrderBeforeCrash) // a is unevictable and b is incomplete // Assume that the application restarts here. - store, err = NewDiskStore(10*memsize.KB, rootDir, false, tally.NoopScope) + store, err = NewDiskStore(10*memsize.KB, rootDir, false, _defaultShardLength, tally.NoopScope) require.NoError(err) // LRU order is approximated, but not exact. @@ -210,7 +210,7 @@ func TestCrashRecovery(t *testing.T) { require.Equal(wantEvictionOrder, rebootedEvictionOrder) // Assume we redeploy the service with a smaller capacity for the disk store: - store, err = NewDiskStore(6*memsize.KB, rootDir, false, tally.NoopScope) + store, err = NewDiskStore(6*memsize.KB, rootDir, false, _defaultShardLength, tally.NoopScope) require.NoError(err) // since 10KB of blobs are in store, `c` gets evicted to put the store back within its capacity. @@ -239,7 +239,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(cF.Close()) // Assume that the application restarts here. - store, err := NewDiskStore(5*memsize.KB, rootDir, true, tally.NoopScope) + _, err := NewDiskStore(5*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) require.ErrorContains(err, "cannot evict enough, the unevictable/incomplete blobs are using up all the space") }) } @@ -252,7 +252,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { firstData := fillWithRandomData(t, f, 2*memsize.KB) // First crash. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) + store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) require.NoError(err) f, err = store.Open(key) @@ -266,7 +266,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { require.NoError(err) // Second crash. - store, err = NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) + store, err = NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) require.NoError(err) wantData := make([]byte, 5*memsize.KB) @@ -283,7 +283,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { require.NoError(store.MarkComplete(key)) // Third crash. - store, err = NewDiskStore(10*memsize.KB, rootDir, true, tally.NoopScope) + store, err = NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) require.NoError(err) // now that the blob is complete, its actual size should be rebooted through stat, instead of trusting the _size sidecar file. require.Equal(5*memsize.KB, store.size) diff --git a/lib/store/disk_store.go b/lib/store/disk_store.go index 40b6f69a7..ce5495652 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk_store.go @@ -64,7 +64,7 @@ type blob struct { evictionBanned bool } -func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, metrics tally.Scope) (*diskStore, error) { +func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, shardLength int, metrics tally.Scope) (*diskStore, error) { // TODO - create a Config struct. // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. // TODO - move disk store files into their own directory and package. @@ -85,12 +85,12 @@ func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bo evictQueue: list.New(), log: log, metrics: metrics, - pather: newPather(rootDir), + pather: newPather(rootDir, shardLength), rebootIncompleteBlobs: rebootIncompleteBlobs, }, nil } - store, err := rebootPersistedStore(capacityBytes, rootDir, rebootIncompleteBlobs, log, metrics) + store, err := rebootPersistedStore(capacityBytes, rootDir, rebootIncompleteBlobs, shardLength, log, metrics) if err != nil { err = fmt.Errorf("reboot persisted state into memory: %w", err) log.With("error", err).Error("Failed to initialize disk store") diff --git a/lib/store/disk_store_test.go b/lib/store/disk_store_test.go index d39168d10..5bc74129a 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk_store_test.go @@ -20,12 +20,16 @@ import ( "github.com/uber/kraken/utils/memsize" ) +const ( + _defaultShardLength = 2 +) + func newTestStore(t *testing.T, capacity uint64, rebootIncompleteBlobs bool) (res *DiskStore, rootDir string) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) t.Cleanup(func() { _ = os.RemoveAll(rootDir) }) - store, err := NewDiskStore(capacity, rootDir, rebootIncompleteBlobs, tally.NoopScope) + store, err := NewDiskStore(capacity, rootDir, rebootIncompleteBlobs, _defaultShardLength, tally.NoopScope) require.NoError(t, err) return store, rootDir } diff --git a/lib/store/pather.go b/lib/store/pather.go index 5ba18be0e..1f870b516 100644 --- a/lib/store/pather.go +++ b/lib/store/pather.go @@ -1,22 +1,28 @@ package store -import "path/filepath" +import ( + "fmt" + "io/fs" + "path/filepath" + "strings" +) const ( - // _defaultShardIDLength is the number of bytes of file digest to be used for shard ID. - // For every byte (2 HEX char), one more level of directories will be created. - _defaultShardIDLength = 2 - _incompleteSubDir = "incomplete" - _completeSubDir = "complete" - _blobFileName = "data" + _incompleteSubDir = "incomplete" + _completeSubDir = "complete" + _blobFileName = "data" ) type pather struct { - dir string + dir string + shardLength int } -func newPather(rootDir string) *pather { - return &pather{dir: rootDir} +// shardLength is the number of bytes of blob key to be used for shard ID. +// For every byte (2 HEX char), one more level of directories will be created. +// A shardLength of 0 denotes no sharding. +func newPather(rootDir string, shardLength int) *pather { + return &pather{dir: rootDir, shardLength: shardLength} } func (p *pather) blobPath(key string, complete bool) string { @@ -25,13 +31,12 @@ func (p *pather) blobPath(key string, complete bool) string { } func (p *pather) dirPath(key string, complete bool) string { - // TODO - allow config to specify whether to shard or not. Allow no sharding, so we can replace [SimpleStore]. subDirName := _incompleteSubDir if complete { subDirName = _completeSubDir } dirPath := filepath.Join(p.dir, subDirName) - for i := 0; i < int(_defaultShardIDLength) && i < len(key)/2; i++ { + for i := 0; i < int(p.shardLength) && i < len(key)/2; i++ { // (1 byte = 2 char of file name assuming file name is in HEX) dirName := key[i*2 : i*2+2] dirPath = filepath.Join(dirPath, dirName) @@ -44,3 +49,40 @@ func (p *pather) sidecarFilePath(key string, complete bool, sidecarFileName stri dirPath := p.dirPath(key, complete) return filepath.Join(dirPath, sidecarFileName) } + +func (p *pather) rebootKeys(subDir string) ([]string, error) { + keys := make([]string, 0) + ok, err := exists(subDir) + if err != nil { + return nil, fmt.Errorf("exists: %w", err) + } + if !ok { + return []string{}, nil + } + err = filepath.WalkDir(subDir, func(path string, entry fs.DirEntry, err error) error { + if err != nil { + return err + } + + if !entry.IsDir() { + return nil + } + relPath, err := filepath.Rel(subDir, path) + if err != nil { + return err + } + nameParts := strings.Split(relPath, string(filepath.Separator)) + numShards := p.shardLength + isBlobDir := len(nameParts) == numShards+1 + if !isBlobDir { + return nil + } + key := nameParts[len(nameParts)-1] + keys = append(keys, key) + return fs.SkipDir + }) + if err != nil { + return nil, fmt.Errorf("walk through subdir '%v' to reboot blob keys: %w", subDir, err) + } + return keys, nil +} diff --git a/lib/store/pather_test.go b/lib/store/pather_test.go index 2c239a227..3d456dacc 100644 --- a/lib/store/pather_test.go +++ b/lib/store/pather_test.go @@ -13,35 +13,43 @@ func TestPather(t *testing.T) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) t.Cleanup(func() { _ = os.RemoveAll(rootDir) }) - pather := newPather(rootDir) + shardLength := 2 + shardedPather := newPather(rootDir, shardLength) + unshardedPather := newPather(rootDir, 0) md := metadata.NewTorrentMeta(core.MetaInfoFixture()) key := "8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - t.Run("incomplete entry", func(t *testing.T) { + t.Run("incomplete entry, sharded", func(t *testing.T) { require := require.New(t) - complete := false - dirPath := pather.dirPath(key, complete) - wantDirPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - require.Equal(wantDirPath, dirPath) - blobPath := pather.blobPath(key, complete) - wantBlobPath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" - require.Equal(wantBlobPath, blobPath) - sidecarFilePath := pather.sidecarFilePath(key, complete, md.GetSuffix()) - wantSidecarFilePath := rootDir + "/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" - require.Equal(wantSidecarFilePath, sidecarFilePath) + require.Equal(rootDir+"/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78", shardedPather.dirPath(key, _incompleteBlob)) + require.Equal(rootDir+"/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data", shardedPather.blobPath(key, _incompleteBlob)) + require.Equal(rootDir+"/incomplete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta", shardedPather.sidecarFilePath(key, _incompleteBlob, md.GetSuffix())) + }) + t.Run("incomplete entry, unsharded", func(t *testing.T) { + require := require.New(t) + require.Equal(rootDir+"/incomplete/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78", unshardedPather.dirPath(key, _incompleteBlob)) + require.Equal(rootDir+"/incomplete/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data", unshardedPather.blobPath(key, _incompleteBlob)) + require.Equal(rootDir+"/incomplete/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta", unshardedPather.sidecarFilePath(key, _incompleteBlob, md.GetSuffix())) + }) + + t.Run("complete entry, sharded", func(t *testing.T) { + require := require.New(t) + require.Equal(rootDir+"/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78", shardedPather.dirPath(key, _completeBlob)) + require.Equal(rootDir+"/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data", shardedPather.blobPath(key, _completeBlob)) + require.Equal(rootDir+"/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta", shardedPather.sidecarFilePath(key, _completeBlob, md.GetSuffix())) + }) + t.Run("complete entry, unsharded", func(t *testing.T) { + require := require.New(t) + require.Equal(rootDir+"/complete/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78", unshardedPather.dirPath(key, _completeBlob)) + require.Equal(rootDir+"/complete/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data", unshardedPather.blobPath(key, _completeBlob)) + require.Equal(rootDir+"/complete/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta", unshardedPather.sidecarFilePath(key, _completeBlob, md.GetSuffix())) }) - t.Run("complete entry", func(t *testing.T) { + t.Run("unsharded entries with non-digest keys", func(t *testing.T) { + key := "foo" require := require.New(t) - complete := true - dirPath := pather.dirPath(key, complete) - wantDirPath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78" - require.Equal(wantDirPath, dirPath) - blobPath := pather.blobPath(key, complete) - wantBlobPath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/data" - require.Equal(wantBlobPath, blobPath) - sidecarFilePath := pather.sidecarFilePath(key, complete, md.GetSuffix()) - wantSidecarFilePath := rootDir + "/complete/8c/6a/8c6af6ca6458353bfa8cb3d756ca54a4fe7b1de04196bf1b37e0863c3f806a78/_torrentmeta" - require.Equal(wantSidecarFilePath, sidecarFilePath) + require.Equal(rootDir+"/complete/foo", unshardedPather.dirPath(key, _completeBlob)) + require.Equal(rootDir+"/complete/foo/data", unshardedPather.blobPath(key, _completeBlob)) + require.Equal(rootDir+"/complete/foo/_torrentmeta", unshardedPather.sidecarFilePath(key, _completeBlob, md.GetSuffix())) }) } diff --git a/lib/store/scoped_store.go b/lib/store/scoped_store.go index fe2514141..bf66a2777 100644 --- a/lib/store/scoped_store.go +++ b/lib/store/scoped_store.go @@ -37,8 +37,8 @@ type DiskStore struct { // // - If the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), // it evicts blobs until size is within capacity. -func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, metrics tally.Scope) (*DiskStore, error) { - diskStore, err := newDiskStore(capacityBytes, rootDir, rebootIncompleteBlobs, metrics) +func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, shardLength int, metrics tally.Scope) (*DiskStore, error) { + diskStore, err := newDiskStore(capacityBytes, rootDir, rebootIncompleteBlobs, shardLength, metrics) if err != nil { return nil, err } From 85e4790c2ce75bf95e99473628a305d5f6a1927c Mon Sep 17 00:00:00 2001 From: Anton Kalpakchiev Date: Fri, 24 Jul 2026 17:32:09 +0200 Subject: [PATCH 10/10] style(DiskStore): Move DiskStore into its own package, add Config, and 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 --- lib/store/disk/config.go | 16 ++ lib/store/{ => disk}/crash_recovery.go | 88 +++--- lib/store/{ => disk}/crash_recovery_test.go | 272 ++++++++++-------- lib/store/{ => disk}/pather.go | 20 +- lib/store/{ => disk}/pather_test.go | 4 +- lib/store/{ => disk}/scoped_store.go | 75 ++--- lib/store/{disk_store.go => disk/store.go} | 124 ++++---- .../store_test.go} | 72 +++-- lib/store/file.go | 2 +- 9 files changed, 363 insertions(+), 310 deletions(-) create mode 100644 lib/store/disk/config.go rename lib/store/{ => disk}/crash_recovery.go (60%) rename lib/store/{ => disk}/crash_recovery_test.go (54%) rename lib/store/{ => disk}/pather.go (79%) rename lib/store/{ => disk}/pather_test.go (97%) rename lib/store/{ => disk}/scoped_store.go (53%) rename lib/store/{disk_store.go => disk/store.go} (74%) rename lib/store/{disk_store_test.go => disk/store_test.go} (94%) diff --git a/lib/store/disk/config.go b/lib/store/disk/config.go new file mode 100644 index 000000000..3f3a0e75a --- /dev/null +++ b/lib/store/disk/config.go @@ -0,0 +1,16 @@ +package disk + +// Config configures [Store]. +type Config struct { + // The capacity of the store in bytes. When breached, LRU eviction is used. + CapacityBytes uint64 + // The root directory under which [Store]'s blobs and state are stored. + RootDir string + // Whether after crash/restart, the Store removes incomplete files from disk (usually to prevent leaks) OR + // reboots any incomplete files from disk (allowing users to continue the blob download/upload, where it was left off before the crash). + RebootIncompleteBlobs bool + // If > 0, directory sharding is used to speed up performance, where ShardLength denotes + // 1) the length of each directory shard's name and 2) the number of shards. + // A value of 0 denotes no sharding. + ShardLength int +} diff --git a/lib/store/crash_recovery.go b/lib/store/disk/crash_recovery.go similarity index 60% rename from lib/store/crash_recovery.go rename to lib/store/disk/crash_recovery.go index 5d9f3851a..04e908a62 100644 --- a/lib/store/crash_recovery.go +++ b/lib/store/disk/crash_recovery.go @@ -1,4 +1,4 @@ -package store +package disk import ( "container/list" @@ -16,7 +16,7 @@ import ( "go.uber.org/zap" ) -type blobState struct { +type rebootedBlob struct { key string size uint64 mTime time.Time @@ -24,102 +24,102 @@ type blobState struct { complete bool } -func rebootPersistedStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, shardLength int, log *zap.SugaredLogger, metrics tally.Scope) (*diskStore, error) { - completeDirPath, incompleteDirPath := filepath.Join(rootDir, _completeSubDir), filepath.Join(rootDir, _incompleteSubDir) - if !rebootIncompleteBlobs { +func rebootPersistedStore(config *Config, log *zap.SugaredLogger, metrics tally.Scope) (*store, error) { + incompleteDirPath := filepath.Join(config.RootDir, _incompleteSubDir) + if !config.RebootIncompleteBlobs { err := os.RemoveAll(incompleteDirPath) if err != nil { return nil, fmt.Errorf("remove incomplete blobs left from a previous service run: %w", err) } } - pather := newPather(rootDir, shardLength) - keys, err := pather.rebootKeys(completeDirPath) + pather := newPather(config.RootDir, config.ShardLength) + keys, err := pather.rebootKeys(_completeBlob) if err != nil { return nil, err } numCompleteBlobs := len(keys) - if rebootIncompleteBlobs { - incompleteKeys, err := pather.rebootKeys(incompleteDirPath) + if config.RebootIncompleteBlobs { + incompleteKeys, err := pather.rebootKeys(_incompleteBlob) if err != nil { return nil, err } keys = append(keys, incompleteKeys...) } - completeEvictableEntries := make([]*blobState, 0) - otherEntries := make([]*blobState, 0) + completeEvictableBlobs := make([]*rebootedBlob, 0) + otherBlobs := make([]*rebootedBlob, 0) for i, key := range keys { complete := i < numCompleteBlobs - bState, ok, err := rebootBlobState(key, complete, pather) + b, ok, err := rebootBlob(key, complete, pather) if err != nil { return nil, err } if !ok { - log.Warn("Could not reboot blob from disk - its parent directory is there but the blob is missing") + log.With("key", key).Warn("Could not reboot blob from disk - its parent directory is there but the blob is missing") continue } - if bState.complete && bState.evictable { - completeEvictableEntries = append(completeEvictableEntries, bState) + if b.complete && b.evictable { + completeEvictableBlobs = append(completeEvictableBlobs, b) } else { - otherEntries = append(otherEntries, bState) + otherBlobs = append(otherBlobs, b) } } storeSize := uint64(0) blobs := make(map[string]*blob, 0) - for _, bState := range otherEntries { - blobs[bState.key] = &blob{ - size: bState.size, - complete: bState.complete, - evictionBanned: !bState.evictable, + for _, b := range otherBlobs { + blobs[b.key] = &blob{ + size: b.size, + complete: b.complete, + evictionBanned: !b.evictable, node: nil, } - storeSize += bState.size + storeSize += b.size } - slices.SortFunc(completeEvictableEntries, func(left, right *blobState) int { + slices.SortFunc(completeEvictableBlobs, func(left, right *rebootedBlob) int { // left-most is oldest, i.e. next-to-evict. return left.mTime.Compare(right.mTime) }) evictQueue := list.New() - for _, bState := range completeEvictableEntries { - node := evictQueue.PushBack(bState.key) - blobs[bState.key] = &blob{ - size: bState.size, + for _, b := range completeEvictableBlobs { + node := evictQueue.PushBack(b.key) + blobs[b.key] = &blob{ + size: b.size, complete: true, node: node, evictionBanned: false, } - storeSize += bState.size + storeSize += b.size } - store := &diskStore{ - blobs: blobs, - evictQueue: evictQueue, - capacity: capacityBytes, - pather: pather, - size: storeSize, - rebootIncompleteBlobs: rebootIncompleteBlobs, - log: log, - metrics: metrics, + store := &store{ + blobs: blobs, + evictQueue: evictQueue, + capacity: config.CapacityBytes, + size: storeSize, + pather: pather, + config: config, + log: log, + metrics: metrics, } if store.size > store.capacity { prevSize := store.size - // evicts entries until size <= capacity. + // evicts blobs until size <= capacity. err = store.reserveSpace(0) if err != nil { - log.With("error", err).Error("DiskStore size exceeds its capacity after service reboot. Evicting blobs from disk did not work to reduce size within capacity.") + log.With("error", err).Error("Store size exceeds its capacity after service reboot. Evicting blobs from disk did not work to reduce size within capacity.") return nil, fmt.Errorf("remove blobs to reduce store size within configured capacity: %w", err) } evictedBytes := prevSize - store.size - log.With("evicted_bytes", evictedBytes).Warn("DiskStore size exceeded its capacity after service reboot. Successfully evicted blobs to reduce size within capacity.") + log.With("evicted_bytes", evictedBytes).Warn("Store size exceeded its capacity after service reboot. Successfully evicted blobs to reduce size within capacity.") } return store, nil } -func rebootBlobState(key string, complete bool, pather *pather) (res *blobState, ok bool, err error) { +func rebootBlob(key string, complete bool, pather *pather) (res *rebootedBlob, ok bool, err error) { blobPath := pather.blobPath(key, complete) fInfo, err := os.Stat(blobPath) if errors.Is(err, os.ErrNotExist) { @@ -127,7 +127,7 @@ func rebootBlobState(key string, complete bool, pather *pather) (res *blobState, return nil, false, nil } if err != nil { - return nil, false, fmt.Errorf("stat blob entry: %w", err) + return nil, false, fmt.Errorf("stat blob file: %w", err) } flagBlobPath := pather.sidecarFilePath(key, complete, _evictionBannedFileName) @@ -148,7 +148,7 @@ func rebootBlobState(key string, complete bool, pather *pather) (res *blobState, } } mTime := fInfo.ModTime() - return &blobState{ + return &rebootedBlob{ key: key, size: size, mTime: mTime, @@ -167,6 +167,7 @@ func rebootIncompleteBlobSize(key string, pather *pather) (size uint64, ok bool, if err != nil { return 0, false, fmt.Errorf("open blob size sidecar file: %w", err) } + defer closers.Close(blobSizeF) blobSizeData, err := io.ReadAll(blobSizeF) if err != nil { return 0, false, fmt.Errorf("read blob size sidecar file: %w", err) @@ -175,7 +176,6 @@ func rebootIncompleteBlobSize(key string, pather *pather) (size uint64, ok bool, if err != nil { return 0, false, fmt.Errorf("blob size sidecar file is in unexpected format: %w", err) } - closers.Close(blobSizeF) return uint64(blobSize), true, nil } diff --git a/lib/store/crash_recovery_test.go b/lib/store/disk/crash_recovery_test.go similarity index 54% rename from lib/store/crash_recovery_test.go rename to lib/store/disk/crash_recovery_test.go index b125aea31..4f5fd288b 100644 --- a/lib/store/crash_recovery_test.go +++ b/lib/store/disk/crash_recovery_test.go @@ -1,124 +1,142 @@ -package store +package disk import ( "bytes" "crypto/rand" "io" + "os" "testing" "github.com/stretchr/testify/require" "github.com/uber-go/tally" "github.com/uber/kraken/core" + storelib "github.com/uber/kraken/lib/store" "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/memsize" ) func TestCrashRecovery(t *testing.T) { - t.Run("blobs, evictability, completeness, and size are recovered", func(t *testing.T) { - require := require.New(t) - store, rootDir := newTestStore(t, 10*memsize.KB, true) - - completeEvictableF, completeEvictableKey := newTestFile(t, store, 2*memsize.KB) - completeEvictableData := fillWithRandomData(t, completeEvictableF, 2*memsize.KB) - // We don't have to test case where file is not closed, as linux closes all FDs owned by a process upon process death. - require.NoError(completeEvictableF.Close()) - require.NoError(store.MarkComplete(completeEvictableKey)) - - completeUnevictableF, completeUnevictableKey := newTestFile(t, store, 2*memsize.KB) - completeUnevictableData := fillWithRandomData(t, completeUnevictableF, 2*memsize.KB) - require.NoError(completeUnevictableF.Close()) - require.NoError(store.MarkComplete(completeUnevictableKey)) - require.NoError(store.BanEviction(completeUnevictableKey)) - - incompleteEvictableF, incompleteEvictableKey := newTestFile(t, store, 2*memsize.KB) - incompleteEvictableData := fillWithRandomData(t, incompleteEvictableF, 2*memsize.KB) - require.NoError(incompleteEvictableF.Close()) - - incompleteUnevictableF, incompleteUnevictableKey := newTestFile(t, store, 2*memsize.KB) - incompleteUnevictableData := fillWithRandomData(t, incompleteUnevictableF, 2*memsize.KB) - require.NoError(incompleteUnevictableF.Close()) - require.NoError(store.BanEviction(incompleteUnevictableKey)) - - // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) - require.NoError(err) - - // Incomplete files are recovered (since `rebootIncompleteBlobs` is true). - f, err := store.Open(incompleteEvictableKey) - require.NoError(err) - defer func(f io.Closer) { require.NoError(f.Close()) }(f) - data, err := io.ReadAll(f) - require.NoError(err) - require.Equal(incompleteEvictableData, data) - unevictable, err := store.checkDiskIfUnevictable(completeEvictableKey, _incompleteBlob) - require.NoError(err) - require.False(unevictable) - - f, err = store.Open(incompleteUnevictableKey) - require.NoError(err) - defer func(f io.Closer) { require.NoError(f.Close()) }(f) - data, err = io.ReadAll(f) - require.NoError(err) - require.Equal(incompleteUnevictableData, data) - unevictable, err = store.checkDiskIfUnevictable(completeEvictableKey, _incompleteBlob) - require.NoError(err) - require.False(unevictable) - - // Complete files are always recovered. - f, err = store.ScopeComplete().Open(completeEvictableKey) - require.NoError(err) - defer func(f io.Closer) { require.NoError(f.Close()) }(f) - data, err = io.ReadAll(f) - require.NoError(err) - require.Equal(completeEvictableData, data) - unevictable, err = store.checkDiskIfUnevictable(completeEvictableKey, _completeBlob) - require.NoError(err) - require.False(unevictable) - - f, err = store.ScopeComplete().Open(completeUnevictableKey) - require.NoError(err) - defer func(f io.Closer) { require.NoError(f.Close()) }(f) - data, err = io.ReadAll(f) - require.NoError(err) - require.Equal(completeUnevictableData, data) - unevictable, err = store.checkDiskIfUnevictable(completeUnevictableKey, _completeBlob) - require.NoError(err) - require.True(unevictable) - - // The sizes of both incomplete and complete files are recovered correctly. - require.Equal(8*memsize.KB, store.size) - - // Run the store with `rebootIncompleteBlobs` as false. - store, err = NewDiskStore(10*memsize.KB, rootDir, false, _defaultShardLength, tally.NoopScope) - require.NoError(err) - - // Incomplete files are dropped. - rebootedKeys := store.List() - require.NotContains(rebootedKeys, incompleteEvictableKey) - require.NotContains(rebootedKeys, incompleteUnevictableKey) - - // Complete files are always recovered. - f, err = store.ScopeComplete().Open(completeEvictableKey) - require.NoError(err) - defer func(f io.Closer) { require.NoError(f.Close()) }(f) - data, err = io.ReadAll(f) - require.NoError(err) - require.Equal(completeEvictableData, data) - unevictable, err = store.checkDiskIfUnevictable(completeEvictableKey, _completeBlob) - require.NoError(err) - require.False(unevictable) - - f, err = store.ScopeComplete().Open(completeUnevictableKey) - require.NoError(err) - defer func(f io.Closer) { require.NoError(f.Close()) }(f) - data, err = io.ReadAll(f) - require.NoError(err) - require.Equal(completeUnevictableData, data) - unevictable, err = store.checkDiskIfUnevictable(completeUnevictableKey, _completeBlob) - require.NoError(err) - require.True(unevictable) - - require.Equal(4*memsize.KB, store.size) + t.Run("blobs, evictability, completeness, and size are recovered regardless of sharding", func(t *testing.T) { + for _, shardLength := range []int{0, _defaultShardLength, 4} { + require := require.New(t) + + rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") + require.NoError(err) + t.Cleanup(func() { require.NoError(os.RemoveAll(rootDir)) }) + config := &Config{ + CapacityBytes: 10 * memsize.KB, + RootDir: rootDir, + RebootIncompleteBlobs: true, + ShardLength: shardLength, + } + + store, err := NewStore(config, tally.NoopScope) + require.NoError(err) + + completeEvictableF, completeEvictableKey := newTestFile(t, store, 2*memsize.KB) + completeEvictableData := fillWithRandomData(t, completeEvictableF, 2*memsize.KB) + // We don't have to test case where file is not closed, as linux closes all FDs owned by a process upon process death. + require.NoError(completeEvictableF.Close()) + require.NoError(store.MarkComplete(completeEvictableKey)) + + completeUnevictableF, completeUnevictableKey := newTestFile(t, store, 2*memsize.KB) + completeUnevictableData := fillWithRandomData(t, completeUnevictableF, 2*memsize.KB) + require.NoError(completeUnevictableF.Close()) + require.NoError(store.MarkComplete(completeUnevictableKey)) + require.NoError(store.BanEviction(completeUnevictableKey)) + + incompleteEvictableF, incompleteEvictableKey := newTestFile(t, store, 2*memsize.KB) + incompleteEvictableData := fillWithRandomData(t, incompleteEvictableF, 2*memsize.KB) + require.NoError(incompleteEvictableF.Close()) + + incompleteUnevictableF, incompleteUnevictableKey := newTestFile(t, store, 2*memsize.KB) + incompleteUnevictableData := fillWithRandomData(t, incompleteUnevictableF, 2*memsize.KB) + require.NoError(incompleteUnevictableF.Close()) + require.NoError(store.BanEviction(incompleteUnevictableKey)) + + // Assume that the application crashes here. The application would restart and call `NewStore`. + store, err = NewStore(&Config{10 * memsize.KB, rootDir, true, shardLength}, tally.NoopScope) + require.NoError(err) + + require.Equal(8*memsize.KB, store.impl.size) + + // Incomplete files are recovered (since `rebootIncompleteBlobs` is true). + f, err := store.Open(incompleteEvictableKey) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err := io.ReadAll(f) + require.NoError(err) + require.Equal(incompleteEvictableData, data) + unevictable, err := store.impl.checkDiskIfUnevictable(incompleteEvictableKey, _incompleteBlob) + require.NoError(err) + require.False(unevictable) + + f, err = store.Open(incompleteUnevictableKey) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(incompleteUnevictableData, data) + unevictable, err = store.impl.checkDiskIfUnevictable(incompleteUnevictableKey, _incompleteBlob) + require.NoError(err) + require.True(unevictable) + + // Complete files are always recovered. + f, err = store.ScopeComplete().Open(completeEvictableKey) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(completeEvictableData, data) + unevictable, err = store.impl.checkDiskIfUnevictable(completeEvictableKey, _completeBlob) + require.NoError(err) + require.False(unevictable) + + f, err = store.ScopeComplete().Open(completeUnevictableKey) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(completeUnevictableData, data) + unevictable, err = store.impl.checkDiskIfUnevictable(completeUnevictableKey, _completeBlob) + require.NoError(err) + require.True(unevictable) + + // The sizes of both incomplete and complete files are recovered correctly. + require.Equal(8*memsize.KB, store.impl.size) + + // Run the store with `rebootIncompleteBlobs` as false. + store, err = NewStore(&Config{10 * memsize.KB, rootDir, false, shardLength}, tally.NoopScope) + require.NoError(err) + + // Incomplete files are dropped. + rebootedKeys := store.List() + require.NotContains(rebootedKeys, incompleteEvictableKey) + require.NotContains(rebootedKeys, incompleteUnevictableKey) + + // Complete files are always recovered. + f, err = store.ScopeComplete().Open(completeEvictableKey) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(completeEvictableData, data) + unevictable, err = store.impl.checkDiskIfUnevictable(completeEvictableKey, _completeBlob) + require.NoError(err) + require.False(unevictable) + + f, err = store.ScopeComplete().Open(completeUnevictableKey) + require.NoError(err) + defer func(f io.Closer) { require.NoError(f.Close()) }(f) + data, err = io.ReadAll(f) + require.NoError(err) + require.Equal(completeUnevictableData, data) + unevictable, err = store.impl.checkDiskIfUnevictable(completeUnevictableKey, _completeBlob) + require.NoError(err) + require.True(unevictable) + + require.Equal(4*memsize.KB, store.impl.size) + } }) t.Run("metadata is recovered for complete blob", func(t *testing.T) { @@ -131,8 +149,8 @@ func TestCrashRecovery(t *testing.T) { writtenMd := metadata.NewTorrentMeta(core.MetaInfoFixture()) require.NoError(store.SetMetadata(key, writtenMd)) - // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) + // Assume that the application crashes here. The application would restart and call `NewStore`. + store, err := NewStore(&Config{10 * memsize.KB, rootDir, true, _defaultShardLength}, tally.NoopScope) require.NoError(err) var readMd metadata.TorrentMeta @@ -151,8 +169,8 @@ func TestCrashRecovery(t *testing.T) { writtenMd := metadata.NewTorrentMeta(core.MetaInfoFixture()) require.NoError(store.SetMetadata(key, writtenMd)) - // Assume that the application crashes here. The application would restart and call `NewDiskStore`. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) + // Assume that the application crashes here. The application would restart and call `NewStore`. + store, err := NewStore(&Config{10 * memsize.KB, rootDir, true, _defaultShardLength}, tally.NoopScope) require.NoError(err) var readMd metadata.TorrentMeta @@ -190,33 +208,33 @@ func TestCrashRecovery(t *testing.T) { _ = fillWithRandomData(t, eF, 2*memsize.KB) require.NoError(eF.Close()) require.NoError(store.MarkComplete(eKey)) - require.Equal([]string{cKey, dKey, eKey}, store.evictionOrder()) // a is unevictable and b is incomplete + require.Equal([]string{cKey, dKey, eKey}, store.impl.evictionOrder()) // a is unevictable and b is incomplete // reset the access time for d dF, err := store.ScopeComplete().Open(dKey) require.NoError(err) require.NoError(dF.Close()) - evictionOrderBeforeCrash := store.evictionOrder() + evictionOrderBeforeCrash := store.impl.evictionOrder() require.Equal([]string{cKey, eKey, dKey}, evictionOrderBeforeCrash) // a is unevictable and b is incomplete // Assume that the application restarts here. - store, err = NewDiskStore(10*memsize.KB, rootDir, false, _defaultShardLength, tally.NoopScope) + store, err = NewStore(&Config{10 * memsize.KB, rootDir, false, _defaultShardLength}, tally.NoopScope) require.NoError(err) // LRU order is approximated, but not exact. - rebootedEvictionOrder := store.evictionOrder() + rebootedEvictionOrder := store.impl.evictionOrder() require.NotEqual(evictionOrderBeforeCrash, rebootedEvictionOrder) wantEvictionOrder := []string{cKey, dKey, eKey} require.Equal(wantEvictionOrder, rebootedEvictionOrder) // Assume we redeploy the service with a smaller capacity for the disk store: - store, err = NewDiskStore(6*memsize.KB, rootDir, false, _defaultShardLength, tally.NoopScope) + store, err = NewStore(&Config{6 * memsize.KB, rootDir, false, _defaultShardLength}, tally.NoopScope) require.NoError(err) // since 10KB of blobs are in store, `c` gets evicted to put the store back within its capacity. require.NotContains(store.List(), cKey) - require.Equal([]string{dKey, eKey}, store.evictionOrder()) + require.Equal([]string{dKey, eKey}, store.impl.evictionOrder()) require.Contains(store.List(), aKey) }) @@ -239,7 +257,7 @@ func TestCrashRecovery(t *testing.T) { require.NoError(cF.Close()) // Assume that the application restarts here. - _, err := NewDiskStore(5*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) + _, err := NewStore(&Config{5 * memsize.KB, rootDir, true, _defaultShardLength}, tally.NoopScope) require.ErrorContains(err, "cannot evict enough, the unevictable/incomplete blobs are using up all the space") }) } @@ -252,7 +270,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { firstData := fillWithRandomData(t, f, 2*memsize.KB) // First crash. - store, err := NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) + store, err := NewStore(&Config{10 * memsize.KB, rootDir, true, _defaultShardLength}, tally.NoopScope) require.NoError(err) f, err = store.Open(key) @@ -266,7 +284,7 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { require.NoError(err) // Second crash. - store, err = NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) + store, err = NewStore(&Config{10 * memsize.KB, rootDir, true, _defaultShardLength}, tally.NoopScope) require.NoError(err) wantData := make([]byte, 5*memsize.KB) @@ -279,14 +297,14 @@ func TestIncompleteBlobDownloadResumedAfterMultipleCrashes(t *testing.T) { require.NoError(err) require.Equal(wantData, data) // The declared size is recovered correctly both times, without leaking or duplicating reservation. - require.Equal(4*memsize.KB, store.size) + require.Equal(4*memsize.KB, store.impl.size) require.NoError(store.MarkComplete(key)) // Third crash. - store, err = NewDiskStore(10*memsize.KB, rootDir, true, _defaultShardLength, tally.NoopScope) + store, err = NewStore(&Config{10 * memsize.KB, rootDir, true, _defaultShardLength}, tally.NoopScope) require.NoError(err) // now that the blob is complete, its actual size should be rebooted through stat, instead of trusting the _size sidecar file. - require.Equal(5*memsize.KB, store.size) + require.Equal(5*memsize.KB, store.impl.size) f, err = store.ScopeComplete().Open(key) require.NoError(err) defer func(f io.Closer) { require.NoError(f.Close()) }(f) @@ -310,7 +328,7 @@ func TestStoreWorksWhenFileSizeNotCorrect(t *testing.T) { require.NoError(underreportedF.Close()) require.NoError(store.MarkComplete(underreportedKey)) require.NoError(store.BanEviction(underreportedKey)) - require.Equal(8*memsize.KB, store.size) + require.Equal(8*memsize.KB, store.impl.size) // Even though only 1KB is actually used on disk, the store enforces capacity based on the // declared 8KB, so a 3KB blob doesn't fit alongside it (there's nothing evictable to make room). @@ -323,17 +341,17 @@ func TestStoreWorksWhenFileSizeNotCorrect(t *testing.T) { require.NoError(overreportedF.Close()) require.NoError(store.MarkComplete(overreportedKey)) // The store still only accounts for the declared 2KB, not the actual 5KB written. - require.Equal(10*memsize.KB, store.size) + require.Equal(10*memsize.KB, store.impl.size) // Deleting releases exactly the declared size that was reserved, not the actual bytes found on // disk, keeping reservation accounting self-consistent in both directions. require.NoError(store.Delete(overreportedKey)) - require.Equal(8*memsize.KB, store.size) + require.Equal(8*memsize.KB, store.impl.size) require.NoError(store.Delete(underreportedKey)) - require.Equal(uint64(0), store.size) + require.Equal(uint64(0), store.impl.size) } -func fillWithRandomData(t *testing.T, f FileReadWriter, sizeBytes uint64) []byte { +func fillWithRandomData(t *testing.T, f storelib.FileReadWriter, sizeBytes uint64) []byte { data := make([]byte, sizeBytes) _, err := rand.Read(data) require.NoError(t, err) diff --git a/lib/store/pather.go b/lib/store/disk/pather.go similarity index 79% rename from lib/store/pather.go rename to lib/store/disk/pather.go index 1f870b516..86c3b50f3 100644 --- a/lib/store/pather.go +++ b/lib/store/disk/pather.go @@ -1,4 +1,4 @@ -package store +package disk import ( "fmt" @@ -50,16 +50,21 @@ func (p *pather) sidecarFilePath(key string, complete bool, sidecarFileName stri return filepath.Join(dirPath, sidecarFileName) } -func (p *pather) rebootKeys(subDir string) ([]string, error) { +func (p *pather) rebootKeys(complete bool) ([]string, error) { + subDirName := _incompleteSubDir + if complete { + subDirName = _completeSubDir + } + dir := filepath.Join(p.dir, subDirName) keys := make([]string, 0) - ok, err := exists(subDir) + ok, err := exists(dir) if err != nil { return nil, fmt.Errorf("exists: %w", err) } if !ok { return []string{}, nil } - err = filepath.WalkDir(subDir, func(path string, entry fs.DirEntry, err error) error { + err = filepath.WalkDir(dir, func(path string, entry fs.DirEntry, err error) error { if err != nil { return err } @@ -67,7 +72,10 @@ func (p *pather) rebootKeys(subDir string) ([]string, error) { if !entry.IsDir() { return nil } - relPath, err := filepath.Rel(subDir, path) + if path == dir { + return nil + } + relPath, err := filepath.Rel(dir, path) if err != nil { return err } @@ -82,7 +90,7 @@ func (p *pather) rebootKeys(subDir string) ([]string, error) { return fs.SkipDir }) if err != nil { - return nil, fmt.Errorf("walk through subdir '%v' to reboot blob keys: %w", subDir, err) + return nil, fmt.Errorf("walk through dir '%v' to reboot blob keys: %w", dir, err) } return keys, nil } diff --git a/lib/store/pather_test.go b/lib/store/disk/pather_test.go similarity index 97% rename from lib/store/pather_test.go rename to lib/store/disk/pather_test.go index 3d456dacc..6490e4c6b 100644 --- a/lib/store/pather_test.go +++ b/lib/store/disk/pather_test.go @@ -1,4 +1,4 @@ -package store +package disk import ( "os" @@ -12,7 +12,7 @@ import ( func TestPather(t *testing.T) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) - t.Cleanup(func() { _ = os.RemoveAll(rootDir) }) + t.Cleanup(func() { require.NoError(t, os.RemoveAll(rootDir)) }) shardLength := 2 shardedPather := newPather(rootDir, shardLength) unshardedPather := newPather(rootDir, 0) diff --git a/lib/store/scoped_store.go b/lib/store/disk/scoped_store.go similarity index 53% rename from lib/store/scoped_store.go rename to lib/store/disk/scoped_store.go index bf66a2777..75cbf5389 100644 --- a/lib/store/scoped_store.go +++ b/lib/store/disk/scoped_store.go @@ -1,18 +1,19 @@ -package store +package disk import ( "errors" "os" "github.com/uber-go/tally" + storelib "github.com/uber/kraken/lib/store" "github.com/uber/kraken/lib/store/metadata" ) // ErrOutOfScope is returned when the provided key is in the store, but not in the store's [blobScope], // e.g. if the blob is incomplete, but ScopeComplete was called. -var ErrOutOfScope = errors.New("the blob is in DiskStore but filtered by the selected scope") +var ErrOutOfScope = errors.New("the blob is in Store but filtered by the selected scope") -// DiskStore is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. +// Store is a key-value, persistent, thread-safe, LRU store for blobs and their [metadata.Metadata]. // // - Supports pagination of blobs during reading/writing, such that blobs don't need to be fully loaded into memory. // @@ -22,91 +23,91 @@ var ErrOutOfScope = errors.New("the blob is in DiskStore but filtered by the sel // // - Supports (un-)marking blobs as non-evictable (may be needed when that data must be written back to remote storage). // -// - Crash-resistant - all state is restored upon restart (check [newDiskStore] for details). +// - Crash-resistant - all state is restored upon restart (check [newStore] for details). // -// - Uses directory sharding to speed up disk performance. -type DiskStore struct { - *diskStore +// - Supports directory sharding to speed up disk performance. +type Store struct { + impl *store scope blobScope } -// NewDiskStore initializes a new [*DiskStore]. If the store has been initialized in the same +// NewStore initializes a new [*Store]. If the store has been initialized in the same // directory before, its state is recovered from disk with the following caveats: // // - `rebootIncompleteBlobs` configures whether incomplete blobs are evicted or rebooted on restart. // // - If the store's size is bigger than its capacity (e.g. configured capacity has been reduced or files have been leaked), // it evicts blobs until size is within capacity. -func NewDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, shardLength int, metrics tally.Scope) (*DiskStore, error) { - diskStore, err := newDiskStore(capacityBytes, rootDir, rebootIncompleteBlobs, shardLength, metrics) +func NewStore(config *Config, metrics tally.Scope) (*Store, error) { + s, err := newStore(config, metrics) if err != nil { return nil, err } - return &DiskStore{ - diskStore: diskStore, - scope: blobScopeAny, + return &Store{ + impl: s, + scope: blobScopeAny, }, nil } // Create adds a new, incomplete blob to the store and reserves space for it. // Incomplete entries cannot be automatically evicted. MarkComplete must be called once the blob is complete. -// DiskStore uses `sizeBytes` for its eviction logic even if the blob's real size differs. -func (s *diskStore) Create(key string, sizeBytes uint64) (FileReadWriter, error) { - return s.create(key, sizeBytes) +// The store uses `sizeBytes` for its eviction logic even if the blob's real size differs. +func (s *Store) Create(key string, sizeBytes uint64) (storelib.FileReadWriter, error) { + return s.impl.Create(key, sizeBytes) } // Open returns an FD to a file in the store. [os.ErrNotExist] is returned on missing entry. -func (s *DiskStore) Open(key string) (FileReadWriter, error) { return s.open(key, s.scope) } +func (s *Store) Open(key string) (storelib.FileReadWriter, error) { return s.impl.Open(key, s.scope) } // Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExist] if the blob is not found. -func (s *DiskStore) Stat(key string) (os.FileInfo, error) { return s.stat(key, s.scope) } +func (s *Store) Stat(key string) (os.FileInfo, error) { return s.impl.Stat(key, s.scope) } // MarkComplete marks a blob as fully written. It enlists the blob for LRU eviction (unless BanEviction has been called). // Additionally, other store APIs may filter blobs based on completeness. -func (s *diskStore) MarkComplete(key string) error { return s.markComplete(key) } +func (s *Store) MarkComplete(key string) error { return s.impl.MarkComplete(key) } // Delete removes a blob and its [metadata.Metadata] from the store. -func (s *DiskStore) Delete(key string) error { return s.delete(key, s.scope) } +func (s *Store) Delete(key string) error { return s.impl.Delete(key, s.scope) } // List returns the keys of all blobs (except those out of scope). -func (s *DiskStore) List() []string { return s.list(s.scope) } +func (s *Store) List() []string { return s.impl.list(s.scope) } // BanEviction marks a blob as unevictable by LRU eviction. It is idempotent. // Needed when e.g. blobs must be written back to GCS/S3 and eviction before that is unacceptable. -func (s *DiskStore) BanEviction(key string) error { return s.banEviction(key, s.scope) } +func (s *Store) BanEviction(key string) error { return s.impl.BanEviction(key, s.scope) } // UnbanEviction removes the effect of BanEviction for a blob. It is idempotent. -func (s *DiskStore) UnbanEviction(key string) error { return s.unbanEviction(key, s.scope) } +func (s *Store) UnbanEviction(key string) error { return s.impl.UnbanEviction(key, s.scope) } // SetMetadata atomically sets the respective metadata for a blob. -func (s *DiskStore) SetMetadata(key string, md metadata.Metadata) error { - return s.setMetadata(key, md, s.scope) +func (s *Store) SetMetadata(key string, md metadata.Metadata) error { + return s.impl.SetMetadata(key, md, s.scope) } // GetMetadata populates `md` if the metadata is present. Returns [os.ErrNotExist] if key is not in store. -func (s *DiskStore) GetMetadata(key string, md metadata.Metadata) (ok bool, err error) { - return s.getMetadata(key, md, s.scope) +func (s *Store) GetMetadata(key string, md metadata.Metadata) (ok bool, err error) { + return s.impl.GetMetadata(key, md, s.scope) } // DeleteMetadata removes any metadata of a blob with `md`'s suffix, if present. -func (s *DiskStore) DeleteMetadata(key string, md metadata.Metadata) error { - return s.deleteMetadata(key, md, s.scope) +func (s *Store) DeleteMetadata(key string, md metadata.Metadata) error { + return s.impl.DeleteMetadata(key, md, s.scope) } // ListMetadata returns all [metadata.Metadata] of key. -func (s *DiskStore) ListMetadata(key string) ([]metadata.Metadata, error) { - return s.listMetadata(key, s.scope) +func (s *Store) ListMetadata(key string) ([]metadata.Metadata, error) { + return s.impl.ListMetadata(key, s.scope) } // WriteAtMetadata implements [io.WriterAt] for the metadata file on disk. -func (s *DiskStore) WriteAtMetadata(key string, md metadata.Metadata, p []byte, off int64) error { - return s.writeAtMetadata(key, md, p, off, s.scope) +func (s *Store) WriteAtMetadata(key string, md metadata.Metadata, p []byte, off int64) error { + return s.impl.WriteAtMetadata(key, md, p, off, s.scope) } -// ScopeComplete scopes [DiskStore]'s APIs such that they can only operate on complete blobs. +// ScopeComplete scopes [Store]'s APIs such that they can only operate on complete blobs (except MarkComplete and Create). // [ErrOutOfScope] is returned if the user tries to operate on an incomplete blob. -func (s *DiskStore) ScopeComplete() *DiskStore { return &DiskStore{s.diskStore, blobScopeComplete} } +func (s *Store) ScopeComplete() *Store { return &Store{s.impl, blobScopeComplete} } -// ScopeIncomplete scopes [DiskStore]'s APIs such that they can only operate on incomplete blobs. +// ScopeIncomplete scopes [Store]'s APIs such that they can only operate on incomplete blobs. // [ErrOutOfScope] is returned if the user tries to operate on a complete blob. -func (s *DiskStore) ScopeIncomplete() *DiskStore { return &DiskStore{s.diskStore, blobScopeIncomplete} } +func (s *Store) ScopeIncomplete() *Store { return &Store{s.impl, blobScopeIncomplete} } diff --git a/lib/store/disk_store.go b/lib/store/disk/store.go similarity index 74% rename from lib/store/disk_store.go rename to lib/store/disk/store.go index ce5495652..6c3ba7ada 100644 --- a/lib/store/disk_store.go +++ b/lib/store/disk/store.go @@ -1,4 +1,4 @@ -package store +package disk import ( "container/list" @@ -12,6 +12,7 @@ import ( "time" "github.com/uber-go/tally" + storelib "github.com/uber/kraken/lib/store" "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/closers" "github.com/uber/kraken/utils/log" @@ -21,7 +22,7 @@ import ( // the set of blobs that the store's APIs can operate on. type blobScope int -// flags to scope [diskStore]'s APIs to a subset of blobs. +// flags to scope [store]'s APIs to a subset of blobs. const ( blobScopeAny blobScope = iota blobScopeComplete @@ -38,69 +39,63 @@ const ( var _syncEvictionLatencyBuckets = tally.MustMakeExponentialDurationBuckets(100*time.Millisecond, 1.4, 15) -// diskStore implements the APIs of [DiskStore]. [diskStore]'s APIs expose the [blobScope] arg, -// while [DiskStore]'s APIs omit that arg (cleaner interface) and instead expose other APIs to scope the whole store. +// store implements the APIs of [Store]. [store]'s APIs expose the [blobScope] arg, +// while [Store]'s APIs omit that arg (cleaner interface) and instead expose other APIs to scope the whole store. // -// Check [DiskStore]'s comments for details on functionality. -type diskStore struct { +// Check [Store]'s comments for details on functionality. +type store struct { capacity uint64 - size uint64 // includes both used and reserved space. + size uint64 // Includes both actively-used and not yet used, but reserved space. blobs map[string]*blob // TODO - consider whether it's better to use struct instead of pointer to reduce GC stress. - evictQueue *list.List // Back is most recently used, front is the next to evict. - // synchronizes mem state access and syscalls to the fs in the APIs (opening, moving files, etc.) - mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. - // If enabled, incomplete blobs are rebooted in the store upon restart (or after a crash), allowing users to continue using the blob. - // If disabled, incomplete blobs are discarded (usually done to prevent leaks). - rebootIncompleteBlobs bool - log *zap.SugaredLogger - metrics tally.Scope + evictQueue *list.List // Front is the next to evict. + // Synchronizes 1) mem state and 2) disk state. Only acquired by public methods. + mu sync.RWMutex // TODO - evaluate whether the read-to-write ratio is more appropriate for a [sync.Mutex] instead. + config *Config + log *zap.SugaredLogger + metrics tally.Scope *pather } type blob struct { - node *list.Element // value of [list.Element] is [string]. + node *list.Element // Value of [list.Element] is [string]. size uint64 complete bool evictionBanned bool } -func newDiskStore(capacityBytes uint64, rootDir string, rebootIncompleteBlobs bool, shardLength int, metrics tally.Scope) (*diskStore, error) { - // TODO - create a Config struct. - // TODO - consider how to support blob mutation, which might be needed by build-index for tag mutation. - // TODO - move disk store files into their own directory and package. - +func newStore(config *Config, metrics tally.Scope) (*store, error) { log := log.Default().With("module", "disk_store") - ok, err := existsPersistedStore(rootDir) + ok, err := existsPersistedStore(config.RootDir) if err != nil { err = fmt.Errorf("could not check if previously-left persisted state exists on disk: %w", err) log.With("error", err).Error("Failed to initialize disk store") return nil, err } if !ok { - log.Info("Initialized a new, empty DiskStore (did not find any previously persisted state to reboot for DiskStore)") - return &diskStore{ - capacity: capacityBytes, - size: 0, - blobs: make(map[string]*blob), - evictQueue: list.New(), - log: log, - metrics: metrics, - pather: newPather(rootDir, shardLength), - rebootIncompleteBlobs: rebootIncompleteBlobs, + log.Info("Initialized a new, empty Store (did not find any previously persisted state to reboot for Store)") + return &store{ + capacity: config.CapacityBytes, + size: 0, + blobs: make(map[string]*blob), + evictQueue: list.New(), + config: config, + pather: newPather(config.RootDir, config.ShardLength), + log: log, + metrics: metrics, }, nil } - store, err := rebootPersistedStore(capacityBytes, rootDir, rebootIncompleteBlobs, shardLength, log, metrics) + store, err := rebootPersistedStore(config, log, metrics) if err != nil { err = fmt.Errorf("reboot persisted state into memory: %w", err) log.With("error", err).Error("Failed to initialize disk store") return nil, err } - log.With("num_blobs", len(store.blobs)).Info("Successfully rebooted DiskStore's previously left state on disk") + log.With("num_blobs", len(store.blobs)).Info("Successfully rebooted Store's previously left state on disk") return store, nil } -func (s *diskStore) open(key string, scope blobScope) (FileReadWriter, error) { +func (s *store) Open(key string, scope blobScope) (storelib.FileReadWriter, error) { s.mu.Lock() defer s.mu.Unlock() @@ -120,12 +115,12 @@ func (s *diskStore) open(key string, scope blobScope) (FileReadWriter, error) { if err != nil { return nil, fmt.Errorf("open: %w", err) } - return newReadWriter(f), nil + return storelib.NewReadWriter(f), nil } -func (s *diskStore) stat(key string, scope blobScope) (os.FileInfo, error) { - s.mu.Lock() - defer s.mu.Unlock() +func (s *store) Stat(key string, scope blobScope) (os.FileInfo, error) { + s.mu.RLock() + defer s.mu.RUnlock() b, ok := s.blobs[key] if !ok { @@ -134,11 +129,12 @@ func (s *diskStore) stat(key string, scope blobScope) (os.FileInfo, error) { if err := isOutOfScope(b, scope); err != nil { return nil, err } + // TODO - consider whether this should update the access time of the blob. blobPath := s.blobPath(key, b.complete) return os.Stat(blobPath) } -func (s *diskStore) create(key string, sizeBytes uint64) (FileReadWriter, error) { +func (s *store) Create(key string, sizeBytes uint64) (storelib.FileReadWriter, error) { // TODO - we might want some TTI on uploads to the store, after which we cancel the upload, e.g. 1min without the client uploading more data. s.mu.Lock() defer s.mu.Unlock() @@ -169,7 +165,7 @@ func (s *diskStore) create(key string, sizeBytes uint64) (FileReadWriter, error) return nil, fmt.Errorf("open file: %w", err) } - if s.rebootIncompleteBlobs { + if s.config.RebootIncompleteBlobs { err = s.persistBlobSize(key, sizeBytes) if err != nil { // Fail-open: the blob will be discarded upon reboot if incomplete. @@ -184,10 +180,10 @@ func (s *diskStore) create(key string, sizeBytes uint64) (FileReadWriter, error) evictionBanned: false, } - return newReadWriter(f), nil + return storelib.NewReadWriter(f), nil } -func (s *diskStore) persistBlobSize(key string, sizeBytes uint64) error { +func (s *store) persistBlobSize(key string, sizeBytes uint64) error { blobSizeFilePath := s.sidecarFilePath(key, _incompleteBlob, _blobSizeFileName) blobSizeF, err := os.OpenFile(blobSizeFilePath, os.O_WRONLY|os.O_CREATE|os.O_EXCL, _defaultFilePerm) if err != nil { @@ -201,7 +197,7 @@ func (s *diskStore) persistBlobSize(key string, sizeBytes uint64) error { return nil } -func (s *diskStore) reserveSpace(space uint64) error { +func (s *store) reserveSpace(space uint64) error { // TODO - benchmark and consider whether async eviction makes more sense. startTime := time.Now() for s.size+space > s.capacity { @@ -210,12 +206,12 @@ func (s *diskStore) reserveSpace(space uint64) error { "unevictable_bytes", s.size, "required_space", space, "capacity", s.capacity, - ).Error("Cannot evict enough data to free space for new entry to DiskStore. The unevictable/incomplete blobs are using up all the space") + ).Error("Cannot evict enough data to free space for new entry to Store. The unevictable/incomplete blobs are using up all the space") return errors.New("cannot evict enough, the unevictable/incomplete blobs are using up all the space") } toEvictNode := s.evictQueue.Front() - toEvictKey := toEvictNode.Value.(string) + toEvictKey := toEvictNode.Value.(string) //nolint:errcheck // We only ever store string in this value. err := s.deleteFromDisk(toEvictKey, _completeBlob) if err != nil { @@ -233,9 +229,9 @@ func (s *diskStore) reserveSpace(space uint64) error { return nil } -func (s *diskStore) releaseSpace(space uint64) { +func (s *store) releaseSpace(space uint64) { if space > s.size { - s.log.Error("Invariant violation - DiskStore wants to release more disk space than actually reserved. Failing open by releasing all reserved space.") + s.log.Error("Invariant violation - Store wants to release more disk space than actually reserved. Failing open by releasing all reserved space.") s.size = 0 return } @@ -243,12 +239,12 @@ func (s *diskStore) releaseSpace(space uint64) { } // Fully deletes the disk state of a blob, including metadata. Works on any blob. -func (s *diskStore) deleteFromDisk(key string, complete bool) error { +func (s *store) deleteFromDisk(key string, complete bool) error { dir := s.dirPath(key, complete) return os.RemoveAll(dir) } -func (s *diskStore) markComplete(key string) error { +func (s *store) MarkComplete(key string) error { s.mu.Lock() defer s.mu.Unlock() @@ -284,7 +280,7 @@ func (s *diskStore) markComplete(key string) error { // Best-effort attempt to remove immovable metadata. Fail-open on failure // to avoid an inconsistent state, as failure only costs a negligible amount // of disk until the blob is evicted. -func (s *diskStore) tryDeleteImmovableMetadata(key string) { +func (s *store) tryDeleteImmovableMetadata(key string) { mdList, err := s.listMetadataNoLock(key, blobScopeAny) if err != nil { err = fmt.Errorf("list metadata: %w", err) @@ -305,7 +301,7 @@ func (s *diskStore) tryDeleteImmovableMetadata(key string) { } } -func (s *diskStore) checkDiskIfUnevictable(key string, complete bool) (bool, error) { +func (s *store) checkDiskIfUnevictable(key string, complete bool) (bool, error) { flagBlobPath := s.sidecarFilePath(key, complete, _evictionBannedFileName) unevictable, err := exists(flagBlobPath) if err != nil { @@ -314,7 +310,7 @@ func (s *diskStore) checkDiskIfUnevictable(key string, complete bool) (bool, err return unevictable, nil } -func (s *diskStore) delete(key string, scope blobScope) error { +func (s *store) Delete(key string, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -339,7 +335,7 @@ func (s *diskStore) delete(key string, scope blobScope) error { return nil } -func (s *diskStore) list(scope blobScope) []string { +func (s *store) list(scope blobScope) []string { s.mu.RLock() defer s.mu.RUnlock() @@ -353,7 +349,7 @@ func (s *diskStore) list(scope blobScope) []string { return res } -func (s *diskStore) banEviction(key string, scope blobScope) error { +func (s *store) BanEviction(key string, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -385,7 +381,7 @@ func (s *diskStore) banEviction(key string, scope blobScope) error { return nil } -func (s *diskStore) unbanEviction(key string, scope blobScope) error { +func (s *store) UnbanEviction(key string, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -415,7 +411,7 @@ func (s *diskStore) unbanEviction(key string, scope blobScope) error { return nil } -func (s *diskStore) setMetadata(key string, md metadata.Metadata, scope blobScope) error { +func (s *store) SetMetadata(key string, md metadata.Metadata, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -450,7 +446,7 @@ func (s *diskStore) setMetadata(key string, md metadata.Metadata, scope blobScop return nil } -func (s *diskStore) getMetadata(key string, md metadata.Metadata, scope blobScope) (ok bool, err error) { +func (s *store) GetMetadata(key string, md metadata.Metadata, scope blobScope) (ok bool, err error) { s.mu.RLock() defer s.mu.RUnlock() @@ -482,7 +478,7 @@ func (s *diskStore) getMetadata(key string, md metadata.Metadata, scope blobScop return true, nil } -func (s *diskStore) deleteMetadata(key string, md metadata.Metadata, scope blobScope) error { +func (s *store) DeleteMetadata(key string, md metadata.Metadata, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -505,14 +501,14 @@ func (s *diskStore) deleteMetadata(key string, md metadata.Metadata, scope blobS return nil } -func (s *diskStore) listMetadata(key string, scope blobScope) ([]metadata.Metadata, error) { +func (s *store) ListMetadata(key string, scope blobScope) ([]metadata.Metadata, error) { s.mu.RLock() defer s.mu.RUnlock() return s.listMetadataNoLock(key, scope) } -func (s *diskStore) listMetadataNoLock(key string, scope blobScope) ([]metadata.Metadata, error) { +func (s *store) listMetadataNoLock(key string, scope blobScope) ([]metadata.Metadata, error) { b, ok := s.blobs[key] if !ok { return nil, os.ErrNotExist @@ -542,7 +538,7 @@ func (s *diskStore) listMetadataNoLock(key string, scope blobScope) ([]metadata. return res, nil } -func (s *diskStore) writeAtMetadata(key string, md metadata.Metadata, p []byte, off int64, scope blobScope) error { +func (s *store) WriteAtMetadata(key string, md metadata.Metadata, p []byte, off int64, scope blobScope) error { s.mu.Lock() defer s.mu.Unlock() @@ -573,13 +569,13 @@ func (s *diskStore) writeAtMetadata(key string, md metadata.Metadata, p []byte, } // used during testing -func (s *diskStore) evictionOrder() []string { +func (s *store) evictionOrder() []string { s.mu.RLock() defer s.mu.RUnlock() evictionOrder := make([]string, 0) for curr := s.evictQueue.Front(); curr != nil; curr = curr.Next() { - currKey := curr.Value.(string) + currKey := curr.Value.(string) //nolint:errcheck // We only ever store string in this value. evictionOrder = append(evictionOrder, currKey) } return evictionOrder diff --git a/lib/store/disk_store_test.go b/lib/store/disk/store_test.go similarity index 94% rename from lib/store/disk_store_test.go rename to lib/store/disk/store_test.go index 5bc74129a..fcbcf85a6 100644 --- a/lib/store/disk_store_test.go +++ b/lib/store/disk/store_test.go @@ -1,4 +1,4 @@ -package store +package disk import ( "bytes" @@ -16,6 +16,7 @@ import ( "github.com/stretchr/testify/require" "github.com/uber-go/tally" "github.com/uber/kraken/core" + storelib "github.com/uber/kraken/lib/store" "github.com/uber/kraken/lib/store/metadata" "github.com/uber/kraken/utils/memsize" ) @@ -24,17 +25,23 @@ const ( _defaultShardLength = 2 ) -func newTestStore(t *testing.T, capacity uint64, rebootIncompleteBlobs bool) (res *DiskStore, rootDir string) { +func newTestStore(t *testing.T, capacity uint64, rebootIncompleteBlobs bool) (res *Store, rootDir string) { rootDir, err := os.MkdirTemp("/tmp", "kraken-disk-store") require.NoError(t, err) - t.Cleanup(func() { _ = os.RemoveAll(rootDir) }) + t.Cleanup(func() { require.NoError(t, os.RemoveAll(rootDir)) }) + config := &Config{ + CapacityBytes: capacity, + RootDir: rootDir, + RebootIncompleteBlobs: rebootIncompleteBlobs, + ShardLength: _defaultShardLength, + } - store, err := NewDiskStore(capacity, rootDir, rebootIncompleteBlobs, _defaultShardLength, tally.NoopScope) + store, err := NewStore(config, tally.NoopScope) require.NoError(t, err) return store, rootDir } -func newTestFile(t *testing.T, store *DiskStore, size uint64) (f FileReadWriter, key string) { +func newTestFile(t *testing.T, store *Store, size uint64) (f storelib.FileReadWriter, key string) { require := require.New(t) key = core.DigestFixture().Hex() f, err := store.Create(key, size) @@ -43,9 +50,9 @@ func newTestFile(t *testing.T, store *DiskStore, size uint64) (f FileReadWriter, } // does not count 1) the directories for sharding, 2) metadata files, and 3) the _eviction_banned flag file. -func numBlobsOnDisk(t *testing.T, store *DiskStore) int { +func numBlobsOnDisk(t *testing.T, store *Store) int { numBlobs := 0 - err := filepath.WalkDir(store.dir, func(path string, _ fs.DirEntry, err error) error { + err := filepath.WalkDir(store.impl.dir, func(path string, _ fs.DirEntry, err error) error { if err != nil { return err } @@ -61,7 +68,7 @@ func numBlobsOnDisk(t *testing.T, store *DiskStore) int { return numBlobs } -func TestDiskStore(t *testing.T) { +func TestStore(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB, false) @@ -81,7 +88,7 @@ func TestDiskStore(t *testing.T) { require.Equal(int64(memsize.KB), n) require.NoError(err) } - require.Equal(10*memsize.KB, store.size) + require.Equal(10*memsize.KB, store.impl.size) f, err := store.Create(core.DigestFixture().Hex(), memsize.B) require.EqualError(err, "reserve space: cannot evict enough, the unevictable/incomplete blobs are using up all the space") @@ -125,7 +132,7 @@ func TestEviction(t *testing.T) { e, eKey := newTestFile(t, store, 1*memsize.KB) require.NoError(e.Close()) - require.Equal(24*memsize.KB, store.size) + require.Equal(24*memsize.KB, store.impl.size) require.Equal(5, numBlobsOnDisk(t, store)) // incomplete files cannot be evicted and adding 2KB would result in overreservation. _, err := store.Create(core.DigestFixture().Hex(), 2*memsize.KB) @@ -139,7 +146,7 @@ func TestEviction(t *testing.T) { require.NoError(store.BanEviction(dKey)) // e is banned from eviction before it even becomes complete. require.NoError(store.BanEviction(eKey)) - require.Equal(24*memsize.KB, store.size) + require.Equal(24*memsize.KB, store.impl.size) // Add f (4KB) which should evict c to make space, as d and e are unevictable and c was accessed last (the MarkComplete call). f, fKey := newTestFile(t, store, 4*memsize.KB) require.NoError(f.Close()) @@ -147,14 +154,14 @@ func TestEviction(t *testing.T) { keys := store.List() require.NotContains(keys, cKey) // new size == 23KB == 24KB - 5KB (c) + 4KB (f) - require.Equal(23*memsize.KB, store.size) + require.Equal(23*memsize.KB, store.impl.size) require.Equal(5, numBlobsOnDisk(t, store)) // Add g (1KB), which will not evict anything g, gKey := newTestFile(t, store, 1*memsize.KB) require.NoError(g.Close()) require.NoError(store.MarkComplete(gKey)) - require.Equal(24*memsize.KB, store.size) + require.Equal(24*memsize.KB, store.impl.size) require.Equal(6, numBlobsOnDisk(t, store)) // Add h (15KB), which evicts b and a: @@ -162,7 +169,7 @@ func TestEviction(t *testing.T) { require.NoError(h.Close()) require.NoError(store.MarkComplete(hKey)) // size == 24KB == 24KB + 15KB (h) - 5KB (b) - 10KB (a) - require.Equal(24*memsize.KB, store.size) + require.Equal(24*memsize.KB, store.impl.size) require.Equal(5, numBlobsOnDisk(t, store)) keys = store.List() require.NotContains(keys, bKey) @@ -182,7 +189,7 @@ func TestEviction(t *testing.T) { require.NoError(i.Close()) keys = store.List() require.NotContains(keys, fKey) - require.Equal(25*memsize.KB, store.size) + require.Equal(25*memsize.KB, store.impl.size) require.Equal(5, numBlobsOnDisk(t, store)) // eviction order: h(15KB), e(1KB), g(1KB), i(5KB); d(3KB) is unevictable @@ -191,7 +198,7 @@ func TestEviction(t *testing.T) { require.NoError(store.MarkComplete(jKey)) keys = store.List() require.NotContains(keys, hKey) - require.Equal(24*memsize.KB, store.size) + require.Equal(24*memsize.KB, store.impl.size) require.Equal(5, numBlobsOnDisk(t, store)) // eviction order: e(1KB), g(1KB), i(5KB), j(14KB); d(3KB) is unevictable @@ -200,7 +207,7 @@ func TestEviction(t *testing.T) { require.NoError(store.MarkComplete(kKey)) keys = store.List() require.NotContains(keys, eKey) - require.Equal(25*memsize.KB, store.size) + require.Equal(25*memsize.KB, store.impl.size) require.Equal(5, numBlobsOnDisk(t, store)) // eviction order: g(1KB), i(5KB), j(14KB), k(2KB); d(3KB) is unevictable @@ -209,7 +216,7 @@ func TestEviction(t *testing.T) { require.NoError(l.Close()) keys = store.List() require.NotContains(keys, gKey) - require.Equal(25*memsize.KB, store.size) + require.Equal(25*memsize.KB, store.impl.size) require.Equal(5, numBlobsOnDisk(t, store)) // eviction order: i(5KB), j(14KB), k(2KB), l(1KB); d(3KB) is unevictable @@ -217,7 +224,7 @@ func TestEviction(t *testing.T) { require.NoError(store.Delete(jKey)) require.NoError(store.Delete(lKey)) // evictionOrder: k(2KB); d(3KB) - require.Equal(5*memsize.KB, store.size) + require.Equal(5*memsize.KB, store.impl.size) require.Equal(2, numBlobsOnDisk(t, store)) } @@ -341,7 +348,7 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) require.Empty(store.List()) - require.Equal(uint64(0), store.size) + require.Equal(uint64(0), store.impl.size) require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("incomplete, unevictable blob", func(t *testing.T) { @@ -358,7 +365,7 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) require.Empty(store.List()) - require.Equal(uint64(0), store.size) + require.Equal(uint64(0), store.impl.size) require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("complete blob", func(t *testing.T) { @@ -375,7 +382,7 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) require.Empty(store.List()) - require.Equal(uint64(0), store.size) + require.Equal(uint64(0), store.impl.size) require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("complete, unevictable blob", func(t *testing.T) { @@ -393,7 +400,7 @@ func TestDelete(t *testing.T) { require.NoError(store.Delete(key)) require.Empty(store.List()) - require.Equal(uint64(0), store.size) + require.Equal(uint64(0), store.impl.size) require.Equal(0, numBlobsOnDisk(t, store)) }) t.Run("not found", func(t *testing.T) { @@ -420,7 +427,7 @@ func TestMarkComplete(t *testing.T) { require.NoError(store.MarkComplete(key)) require.Equal([]string{key}, store.ScopeComplete().List()) - require.Equal(uint64(100), store.size) + require.Equal(uint64(100), store.impl.size) _, err = store.ScopeComplete().Open(key) require.NoError(err) }) @@ -438,7 +445,7 @@ func TestMarkComplete(t *testing.T) { require.NoError(store.MarkComplete(key)) require.Equal([]string{key}, store.ScopeComplete().List()) - require.Equal(uint64(100), store.size) + require.Equal(uint64(100), store.impl.size) _, err = store.ScopeComplete().Open(key) require.NoError(err) }) @@ -585,8 +592,7 @@ func TestList(t *testing.T) { require := require.New(t) store, _ := newTestStore(t, 10*memsize.KB, false) - require.Empty(store.ScopeComplete().List()) - require.Empty(store.ScopeComplete().List()) + require.Empty(store.List()) incompleteBlobKey := core.DigestFixture().Hex() f, err := store.Create(incompleteBlobKey, 10*memsize.B) @@ -683,7 +689,7 @@ func TestMetadata(t *testing.T) { ok, err = store.GetMetadata(key, &readMd) require.NoError(err) require.False(ok) - mdFilePath := store.sidecarFilePath(key, _incompleteBlob, readMd.GetSuffix()) + mdFilePath := store.impl.sidecarFilePath(key, _incompleteBlob, readMd.GetSuffix()) // ensure the metadata file is deleted from disk _, err = os.Stat(mdFilePath) require.ErrorIs(err, os.ErrNotExist) @@ -778,7 +784,7 @@ func TestMetadata(t *testing.T) { md := metadata.NewTorrentMeta(core.MetaInfoFixture()) err = store.SetMetadata(keyA, md) require.NoError(err) - mdFilePath := store.sidecarFilePath(keyA, _completeBlob, md.GetSuffix()) + mdFilePath := store.impl.sidecarFilePath(keyA, _completeBlob, md.GetSuffix()) _, err = os.Stat(mdFilePath) require.NoError(err) @@ -840,6 +846,10 @@ func TestMetadata(t *testing.T) { require.NoError(err) require.NoError(f.Close()) require.NoError(store.BanEviction(key)) + // A sidecar of every blob when [Config.RebootIncompleteBlobs] is on. + sizeSidecarFilePath := store.impl.sidecarFilePath(key, _incompleteBlob, _blobSizeFileName) + _, err = os.Stat(sizeSidecarFilePath) + require.NoError(err) mdList, err := store.ListMetadata(key) require.NoError(err) @@ -1032,9 +1042,13 @@ func TestScopesDelete(t *testing.T) { f, key = newTestFile(t, store, 1*memsize.KB) require.NoError(f.Close()) require.NoError(store.Delete(key)) + _, err = store.Stat(key) + require.ErrorIs(err, os.ErrNotExist) f, key = newTestFile(t, store, 1*memsize.KB) require.NoError(f.Close()) require.NoError(store.MarkComplete(key)) require.NoError(store.Delete(key)) + _, err = store.Stat(key) + require.ErrorIs(err, os.ErrNotExist) } diff --git a/lib/store/file.go b/lib/store/file.go index 2f36a1d7b..1e10a941c 100644 --- a/lib/store/file.go +++ b/lib/store/file.go @@ -25,7 +25,7 @@ type FileReadWriter = base.FileReadWriter // FileReader is a read-only file. type FileReader = base.FileReader -func newReadWriter(f *os.File) FileReadWriter { +func NewReadWriter(f *os.File) FileReadWriter { return &rwImpl{ File: f, }