feat(store): Add memory LRU cache

lib/store/disk/config.go

@@ -12,5 +12,5 @@ type Config struct {
 	// 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
+	ShardLength int // TODO - add a mechanism to migrate in-place from 1 shardLength to another.
 }

lib/store/disk/file.go

@@ -0,0 +1,47 @@
+package disk
+
+import (
+	storelib "github.com/uber/kraken/lib/store"
+	"os"
+)
+
+func newFile(f *os.File) *File {
+	return &File{
+		fd: f,
+	}
+}
+
+var _ storelib.FileReadWriter = &File{}
+
+// File represends an open file descriptor to a blob in [Store].
+type File struct {
+	fd *os.File
+}
+
+func (f *File) Read(p []byte) (n int, err error)               { return f.fd.Read(p) }
+func (f *File) ReadAt(p []byte, off int64) (n int, err error)  { return f.fd.ReadAt(p, off) }
+func (f *File) Seek(off int64, whence int) (int64, error)      { return f.fd.Seek(off, whence) }
+func (f *File) Write(p []byte) (n int, err error)              { return f.fd.Write(p) }
+func (f *File) WriteAt(p []byte, off int64) (n int, err error) { return f.fd.WriteAt(p, off) }
+func (f *File) Close() error                                   { return f.fd.Close() }
+
+// Size returns the number of bytes the file contains.
+func (f *File) Size() int64 {
+	info, err := f.fd.Stat()
+	if err != nil {
+		return 0
+	}
+	return info.Size()
+}
+
+// Cancel is supposed to remove any written content.
+// In this implementation file is not actually removed, but rather closed, which is fine, as there won't be key collisions when creating new files.
+func (f *File) Cancel() error {
+	return f.fd.Close()
+}
+
+// Commit is supposed to flush all content for buffered writer.
+func (f *File) Commit() error {
+	// TODO - consider whether we should do f.Sync() before closing the file.
+	return f.fd.Close()
+}

lib/store/disk/scoped_store.go

@@ -1,18 +1,13 @@
 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 Store but filtered by the selected scope")
-
 // 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.
@@ -28,7 +23,7 @@ var ErrOutOfScope = errors.New("the blob is in Store but filtered by the selecte
 //   - Supports directory sharding to speed up disk performance.
 type Store struct {
 	impl  *store
-	scope blobScope
+	scope storelib.BlobScope
 }
 
 // NewStore initializes a new [*Store]. If the store has been initialized in the same
@@ -45,19 +40,22 @@ func NewStore(config *Config, metrics tally.Scope) (*Store, error) {
 	}
 	return &Store{
 		impl:  s,
-		scope: blobScopeAny,
+		scope: storelib.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.
 // 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) {
+func (s *Store) Create(key string, sizeBytes uint64) (*File, 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 *Store) Open(key string) (storelib.FileReadWriter, error) { return s.impl.Open(key, s.scope) }
+func (s *Store) Open(key string) (*File, error) { return s.impl.Open(key, s.scope) }
+
+// Has checks if the blob is in the store.
+func (s *Store) Has(key string) (inStore bool, inScope bool) { return s.impl.Has(key, s.scope) }
 
 // Stat returns [os.FileInfo] about the blob. Returns [os.ErrNotExist] if the blob is not found.
 func (s *Store) Stat(key string) (os.FileInfo, error) { return s.impl.Stat(key, s.scope) }
@@ -66,7 +64,7 @@ func (s *Store) Stat(key string) (os.FileInfo, error) { return s.impl.Stat(key,
 // Additionally, other store APIs may filter blobs based on completeness.
 func (s *Store) MarkComplete(key string) error { return s.impl.MarkComplete(key) }
 
-// Delete removes a blob and its [metadata.Metadata] from the store.
+// Delete removes a blob and its [metadata.Metadata] from the store. Returns [os.ErrNotExist] on missing blob.
 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).
@@ -89,7 +87,7 @@ func (s *Store) GetMetadata(key string, md metadata.Metadata) (ok bool, err erro
 	return s.impl.GetMetadata(key, md, s.scope)
 }
 
-// DeleteMetadata removes any metadata of a blob with `md`'s suffix, if present.
+// DeleteMetadata removes a blob's metadata. No error returned if the metadata is not present.
 func (s *Store) DeleteMetadata(key string, md metadata.Metadata) error {
 	return s.impl.DeleteMetadata(key, md, s.scope)
 }
@@ -105,9 +103,9 @@ func (s *Store) WriteAtMetadata(key string, md metadata.Metadata, p []byte, off
 }
 
 // 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 *Store) ScopeComplete() *Store { return &Store{s.impl, blobScopeComplete} }
+// [storelib.ErrOutOfScope] is returned if the user tries to operate on an incomplete blob.
+func (s *Store) ScopeComplete() *Store { return &Store{s.impl, storelib.BlobScopeComplete} }
 
 // 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 *Store) ScopeIncomplete() *Store { return &Store{s.impl, blobScopeIncomplete} }
+// [storelib.ErrOutOfScope] is returned if the user tries to operate on a complete blob.
+func (s *Store) ScopeIncomplete() *Store { return &Store{s.impl, storelib.BlobScopeIncomplete} }

lib/store/disk/store.go

@@ -19,16 +19,6 @@ import (
 	"go.uber.org/zap"
 )
 
-// the set of blobs that the store's APIs can operate on.
-type blobScope int
-
-// flags to scope [store]'s APIs to a subset of blobs.
-const (
-	blobScopeAny blobScope = iota
-	blobScopeComplete
-	blobScopeIncomplete
-)
-
 const (
 	_completeBlob           = true
 	_incompleteBlob         = false
@@ -39,7 +29,7 @@ const (
 
 var _syncEvictionLatencyBuckets = tally.MustMakeExponentialDurationBuckets(100*time.Millisecond, 1.4, 15)
 
-// store implements the APIs of [Store]. [store]'s APIs expose the [blobScope] arg,
+// store implements the APIs of [Store]. [store]'s APIs expose the [storelib.BlobScope] arg,
 // while [Store]'s APIs omit that arg (cleaner interface) and instead expose other APIs to scope the whole store.
 //
 // Check [Store]'s comments for details on functionality.
@@ -73,7 +63,7 @@ func newStore(config *Config, metrics tally.Scope) (*store, error) {
 	}
 	if !ok {
 		log.Info("Initialized a new, empty Store (did not find any previously persisted state to reboot for Store)")
-		return &store{
+		store := &store{
 			capacity:   config.CapacityBytes,
 			size:       0,
 			blobs:      make(map[string]*blob),
@@ -82,7 +72,10 @@ func newStore(config *Config, metrics tally.Scope) (*store, error) {
 			pather:     newPather(config.RootDir, config.ShardLength),
 			log:        log,
 			metrics:    metrics,
-		}, nil
+		}
+
+		store.emitUsageMetrics()
+		return store, nil
 	}
 
 	store, err := rebootPersistedStore(config, log, metrics)
@@ -92,10 +85,12 @@ func newStore(config *Config, metrics tally.Scope) (*store, error) {
 		return nil, err
 	}
 	log.With("num_blobs", len(store.blobs)).Info("Successfully rebooted Store's previously left state on disk")
+
+	store.emitUsageMetrics()
 	return store, nil
 }
 
-func (s *store) Open(key string, scope blobScope) (storelib.FileReadWriter, error) {
+func (s *store) Open(key string, scope storelib.BlobScope) (*File, error) {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -115,10 +110,24 @@ func (s *store) Open(key string, scope blobScope) (storelib.FileReadWriter, erro
 	if err != nil {
 		return nil, fmt.Errorf("open: %w", err)
 	}
-	return storelib.NewReadWriter(f), nil
+	return newFile(f), nil
 }
 
-func (s *store) Stat(key string, scope blobScope) (os.FileInfo, error) {
+func (s *store) Has(key string, scope storelib.BlobScope) (inStore bool, inScope bool) {
+	s.mu.RLock()
+	defer s.mu.RUnlock()
+
+	b, ok := s.blobs[key]
+	if !ok {
+		return false, false
+	}
+	if err := isOutOfScope(b, scope); err != nil {
+		return true, false
+	}
+	return true, true
+}
+
+func (s *store) Stat(key string, scope storelib.BlobScope) (os.FileInfo, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
@@ -134,7 +143,7 @@ func (s *store) Stat(key string, scope blobScope) (os.FileInfo, error) {
 	return os.Stat(blobPath)
 }
 
-func (s *store) Create(key string, sizeBytes uint64) (storelib.FileReadWriter, error) {
+func (s *store) Create(key string, sizeBytes uint64) (*File, 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()
@@ -180,7 +189,8 @@ func (s *store) Create(key string, sizeBytes uint64) (storelib.FileReadWriter, e
 		evictionBanned: false,
 	}
 
-	return storelib.NewReadWriter(f), nil
+	s.emitUsageMetrics()
+	return newFile(f), nil
 }
 
 func (s *store) persistBlobSize(key string, sizeBytes uint64) error {
@@ -231,7 +241,7 @@ func (s *store) reserveSpace(space uint64) error {
 
 func (s *store) releaseSpace(space uint64) {
 	if space > s.size {
-		s.log.Error("Invariant violation - Store wants to release more disk space than actually reserved. Failing open by releasing all reserved space.")
+		s.log.Error("Invariant violation - disk.Store wants to release more space than actually reserved. Failing open by setting store.size = 0")
 		s.size = 0
 		return
 	}
@@ -281,7 +291,7 @@ func (s *store) MarkComplete(key string) error {
 // to avoid an inconsistent state, as failure only costs a negligible amount
 // of disk until the blob is evicted.
 func (s *store) tryDeleteImmovableMetadata(key string) {
-	mdList, err := s.listMetadataNoLock(key, blobScopeAny)
+	mdList, err := s.listMetadataNoLock(key, storelib.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")
@@ -310,7 +320,7 @@ func (s *store) checkDiskIfUnevictable(key string, complete bool) (bool, error)
 	return unevictable, nil
 }
 
-func (s *store) Delete(key string, scope blobScope) error {
+func (s *store) Delete(key string, scope storelib.BlobScope) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -332,10 +342,11 @@ func (s *store) Delete(key string, scope blobScope) error {
 	delete(s.blobs, key)
 	s.releaseSpace(b.size)
 
+	s.emitUsageMetrics()
 	return nil
 }
 
-func (s *store) list(scope blobScope) []string {
+func (s *store) list(scope storelib.BlobScope) []string {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
@@ -349,7 +360,7 @@ func (s *store) list(scope blobScope) []string {
 	return res
 }
 
-func (s *store) BanEviction(key string, scope blobScope) error {
+func (s *store) BanEviction(key string, scope storelib.BlobScope) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -381,7 +392,7 @@ func (s *store) BanEviction(key string, scope blobScope) error {
 	return nil
 }
 
-func (s *store) UnbanEviction(key string, scope blobScope) error {
+func (s *store) UnbanEviction(key string, scope storelib.BlobScope) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -411,7 +422,7 @@ func (s *store) UnbanEviction(key string, scope blobScope) error {
 	return nil
 }
 
-func (s *store) SetMetadata(key string, md metadata.Metadata, scope blobScope) error {
+func (s *store) SetMetadata(key string, md metadata.Metadata, scope storelib.BlobScope) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -446,7 +457,7 @@ func (s *store) SetMetadata(key string, md metadata.Metadata, scope blobScope) e
 	return nil
 }
 
-func (s *store) GetMetadata(key string, md metadata.Metadata, scope blobScope) (ok bool, err error) {
+func (s *store) GetMetadata(key string, md metadata.Metadata, scope storelib.BlobScope) (ok bool, err error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
@@ -478,7 +489,8 @@ func (s *store) GetMetadata(key string, md metadata.Metadata, scope blobScope) (
 	return true, nil
 }
 
-func (s *store) DeleteMetadata(key string, md metadata.Metadata, scope blobScope) error {
+func (s *store) DeleteMetadata(key string, md metadata.Metadata, scope storelib.BlobScope) error {
+	// TODO - change interface to take `mdSuffix string` instead of `md metadata.Metadata`, just like memory.Store.`
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -501,14 +513,14 @@ func (s *store) DeleteMetadata(key string, md metadata.Metadata, scope blobScope
 	return nil
 }
 
-func (s *store) ListMetadata(key string, scope blobScope) ([]metadata.Metadata, error) {
+func (s *store) ListMetadata(key string, scope storelib.BlobScope) ([]metadata.Metadata, error) {
 	s.mu.RLock()
 	defer s.mu.RUnlock()
 
 	return s.listMetadataNoLock(key, scope)
 }
 
-func (s *store) listMetadataNoLock(key string, scope blobScope) ([]metadata.Metadata, error) {
+func (s *store) listMetadataNoLock(key string, scope storelib.BlobScope) ([]metadata.Metadata, error) {
 	b, ok := s.blobs[key]
 	if !ok {
 		return nil, os.ErrNotExist
@@ -538,7 +550,7 @@ func (s *store) listMetadataNoLock(key string, scope blobScope) ([]metadata.Meta
 	return res, nil
 }
 
-func (s *store) 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 storelib.BlobScope) error {
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
@@ -592,9 +604,14 @@ 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
+func isOutOfScope(b *blob, scope storelib.BlobScope) error {
+	if (b.complete && scope == storelib.BlobScopeIncomplete) || (!b.complete && scope == storelib.BlobScopeComplete) {
+		return storelib.ErrOutOfScope
 	}
 	return nil
 }
+
+func (s *store) emitUsageMetrics() {
+	s.metrics.Gauge("num_entries").Update(float64(len(s.blobs)))
+	s.metrics.Gauge("size_bytes").Update(float64(s.size))
+}