Implement cachedAssetStore

This commit is contained in:
Brian Picciano 2022-05-07 18:22:57 -06:00
parent 3088c9d76c
commit 450136a0c0
2 changed files with 99 additions and 19 deletions

View File

@ -6,6 +6,7 @@ import (
"errors" "errors"
"fmt" "fmt"
"io" "io"
"sync"
) )
var ( var (
@ -84,3 +85,69 @@ func (s *assetStore) Delete(id string) error {
_, err := s.db.Exec(`DELETE FROM assets WHERE id = ?`, id) _, err := s.db.Exec(`DELETE FROM assets WHERE id = ?`, id)
return err return err
} }
////////////////////////////////////////////////////////////////////////////////
type cachedAssetStore struct {
inner AssetStore
m sync.Map
}
// NewCachedAssetStore wraps an AssetStore in an in-memory cache.
func NewCachedAssetStore(assetStore AssetStore) AssetStore {
return &cachedAssetStore{
inner: assetStore,
}
}
func (s *cachedAssetStore) Set(id string, from io.Reader) error {
buf := new(bytes.Buffer)
from = io.TeeReader(from, buf)
if err := s.inner.Set(id, from); err != nil {
return err
}
s.m.Store(id, buf.Bytes())
return nil
}
func (s *cachedAssetStore) Get(id string, into io.Writer) error {
if bodyI, ok := s.m.Load(id); ok {
if err, ok := bodyI.(error); ok {
return err
}
if _, err := io.Copy(into, bytes.NewReader(bodyI.([]byte))); err != nil {
return fmt.Errorf("writing body to io.Writer: %w", err)
}
return nil
}
buf := new(bytes.Buffer)
into = io.MultiWriter(into, buf)
if err := s.inner.Get(id, into); errors.Is(err, ErrAssetNotFound) {
s.m.Store(id, err)
return err
} else if err != nil {
return err
}
s.m.Store(id, buf.Bytes())
return nil
}
func (s *cachedAssetStore) Delete(id string) error {
if err := s.inner.Delete(id); err != nil {
return err
}
s.m.Delete(id)
return nil
}

View File

@ -12,14 +12,14 @@ type assetTestHarness struct {
store AssetStore store AssetStore
} }
func newAssetTestHarness(t *testing.T) assetTestHarness { func newAssetTestHarness(t *testing.T) *assetTestHarness {
db := NewInMemSQLDB() db := NewInMemSQLDB()
t.Cleanup(func() { db.Close() }) t.Cleanup(func() { db.Close() })
store := NewAssetStore(db) store := NewAssetStore(db)
return assetTestHarness{ return &assetTestHarness{
store: store, store: store,
} }
} }
@ -40,7 +40,8 @@ func (h *assetTestHarness) assertNotFound(t *testing.T, id string) {
func TestAssetStore(t *testing.T) { func TestAssetStore(t *testing.T) {
h := newAssetTestHarness(t) testAssetStore := func(t *testing.T, h *assetTestHarness) {
t.Helper()
h.assertNotFound(t, "foo") h.assertNotFound(t, "foo")
h.assertNotFound(t, "bar") h.assertNotFound(t, "bar")
@ -64,4 +65,16 @@ func TestAssetStore(t *testing.T) {
assert.NoError(t, h.store.Delete("bar")) assert.NoError(t, h.store.Delete("bar"))
h.assertNotFound(t, "foo") h.assertNotFound(t, "foo")
h.assertNotFound(t, "bar") h.assertNotFound(t, "bar")
}
t.Run("sql", func(t *testing.T) {
h := newAssetTestHarness(t)
testAssetStore(t, h)
})
t.Run("mem", func(t *testing.T) {
h := newAssetTestHarness(t)
h.store = NewCachedAssetStore(h.store)
testAssetStore(t, h)
})
} }