Skip to content

Commit 719ea29

Browse files
owenrumneyowenrumney
authored
chore: add context to the cache interface (#9565)
1 parent 3dd0ebb commit 719ea29

File tree

28 files changed

+198
-184
lines changed

28 files changed

+198
-184
lines changed

‎internal/cachetest/cache.go‎

Lines changed: 13 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cachetest
22

33
import (
4+
"context"
45
"errors"
56
"testing"
67

@@ -42,39 +43,39 @@ func NewErrorCache(opts ErrorCacheOptions) *ErrorCache {
4243
}
4344
}
4445

45-
func (c *ErrorCache) MissingBlobs(artifactID string, blobIDs []string) (bool, []string, error) {
46+
func (c *ErrorCache) MissingBlobs(ctx context.Context, artifactID string, blobIDs []string) (bool, []string, error) {
4647
if c.opts.MissingBlobs {
4748
return false, nil, errors.New("MissingBlobs failed")
4849
}
49-
return c.MemoryCache.MissingBlobs(artifactID, blobIDs)
50+
return c.MemoryCache.MissingBlobs(ctx, artifactID, blobIDs)
5051
}
5152

52-
func (c *ErrorCache) PutArtifact(artifactID string, artifactInfo types.ArtifactInfo) error {
53+
func (c *ErrorCache) PutArtifact(ctx context.Context, artifactID string, artifactInfo types.ArtifactInfo) error {
5354
if c.opts.PutArtifact {
5455
return errors.New("PutArtifact failed")
5556
}
56-
return c.MemoryCache.PutArtifact(artifactID, artifactInfo)
57+
return c.MemoryCache.PutArtifact(ctx, artifactID, artifactInfo)
5758
}
5859

59-
func (c *ErrorCache) PutBlob(artifactID string, blobInfo types.BlobInfo) error {
60+
func (c *ErrorCache) PutBlob(ctx context.Context, artifactID string, blobInfo types.BlobInfo) error {
6061
if c.opts.PutBlob {
6162
return errors.New("PutBlob failed")
6263
}
63-
return c.MemoryCache.PutBlob(artifactID, blobInfo)
64+
return c.MemoryCache.PutBlob(ctx, artifactID, blobInfo)
6465
}
6566

66-
func (c *ErrorCache) GetArtifact(artifactID string) (types.ArtifactInfo, error) {
67+
func (c *ErrorCache) GetArtifact(ctx context.Context, artifactID string) (types.ArtifactInfo, error) {
6768
if c.opts.GetArtifact {
6869
return types.ArtifactInfo{}, errors.New("GetArtifact failed")
6970
}
70-
return c.MemoryCache.GetArtifact(artifactID)
71+
return c.MemoryCache.GetArtifact(ctx, artifactID)
7172
}
7273

73-
func (c *ErrorCache) GetBlob(blobID string) (types.BlobInfo, error) {
74+
func (c *ErrorCache) GetBlob(ctx context.Context, blobID string) (types.BlobInfo, error) {
7475
if c.opts.GetBlob {
7576
return types.BlobInfo{}, errors.New("GetBlob failed")
7677
}
77-
return c.MemoryCache.GetBlob(blobID)
78+
return c.MemoryCache.GetBlob(ctx, blobID)
7879
}
7980

8081
func NewCache(t *testing.T, setUpCache func(t *testing.T) cache.Cache) cache.Cache {
@@ -85,7 +86,7 @@ func NewCache(t *testing.T, setUpCache func(t *testing.T) cache.Cache) cache.Cac
8586
}
8687

8788
func AssertArtifact(t *testing.T, c cache.Cache, wantArtifact WantArtifact) {
88-
gotArtifact, err := c.GetArtifact(wantArtifact.ID)
89+
gotArtifact, err := c.GetArtifact(t.Context(), wantArtifact.ID)
8990
require.NoError(t, err, "artifact not found")
9091
assert.Equal(t, wantArtifact.ArtifactInfo, gotArtifact, wantArtifact.ID)
9192
}
@@ -100,7 +101,7 @@ func AssertBlobs(t *testing.T, c cache.Cache, wantBlobs []WantBlob) {
100101
}
101102

102103
for _, want := range wantBlobs {
103-
got, err := c.GetBlob(want.ID)
104+
got, err := c.GetBlob(t.Context(), want.ID)
104105
require.NoError(t, err, "blob not found")
105106

106107
for i := range got.Misconfigurations {

‎pkg/cache/cache.go‎

Lines changed: 9 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,8 @@
11
package cache
22

33
import (
4+
"context"
5+
46
"github.com/aquasecurity/trivy/pkg/fanal/types"
57
)
68

@@ -21,29 +23,29 @@ type Cache interface {
2123
// ArtifactCache uses local or remote cache
2224
type ArtifactCache interface {
2325
// MissingBlobs returns missing blob IDs such as layer IDs in cache
24-
MissingBlobs(artifactID string, blobIDs []string) (missingArtifact bool, missingBlobIDs []string, err error)
26+
MissingBlobs(ctx context.Context, artifactID string, blobIDs []string) (missingArtifact bool, missingBlobIDs []string, err error)
2527

2628
// PutArtifact stores artifact information such as image metadata in cache
27-
PutArtifact(artifactID string, artifactInfo types.ArtifactInfo) (err error)
29+
PutArtifact(ctx context.Context, artifactID string, artifactInfo types.ArtifactInfo) (err error)
2830

2931
// PutBlob stores blob information such as layer information in local cache
30-
PutBlob(blobID string, blobInfo types.BlobInfo) (err error)
32+
PutBlob(ctx context.Context, blobID string, blobInfo types.BlobInfo) (err error)
3133

3234
// DeleteBlobs removes blobs by IDs
33-
DeleteBlobs(blobIDs []string) error
35+
DeleteBlobs(ctx context.Context, blobIDs []string) error
3436
}
3537

3638
// LocalArtifactCache always uses local cache
3739
type LocalArtifactCache interface {
3840
// GetArtifact gets artifact information such as image metadata from local cache
39-
GetArtifact(artifactID string) (artifactInfo types.ArtifactInfo, err error)
41+
GetArtifact(ctx context.Context, artifactID string) (artifactInfo types.ArtifactInfo, err error)
4042

4143
// GetBlob gets blob information such as layer data from local cache
42-
GetBlob(blobID string) (blobInfo types.BlobInfo, err error)
44+
GetBlob(ctx context.Context, blobID string) (blobInfo types.BlobInfo, err error)
4345

4446
// Close closes the local database
4547
Close() (err error)
4648

4749
// Clear deletes the local database
48-
Clear() (err error)
50+
Clear(ctx context.Context) (err error)
4951
}

‎pkg/cache/fs.go‎

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cache
22

33
import (
4+
"context"
45
"encoding/json"
56
"errors"
67
"os"
@@ -63,7 +64,7 @@ func NewFSCache(cacheDir string) (FSCache, error) {
6364
}
6465

6566
// GetBlob gets blob information such as layer data from local cache
66-
func (fs FSCache) GetBlob(blobID string) (types.BlobInfo, error) {
67+
func (fs FSCache) GetBlob(_ context.Context, blobID string) (types.BlobInfo, error) {
6768
var blobInfo types.BlobInfo
6869
err := fs.db.View(func(tx *bolt.Tx) error {
6970
var err error
@@ -91,7 +92,7 @@ func (fs FSCache) getBlob(blobBucket *bolt.Bucket, diffID string) (types.BlobInf
9192
}
9293

9394
// PutBlob stores blob information such as layer information in local cache
94-
func (fs FSCache) PutBlob(blobID string, blobInfo types.BlobInfo) error {
95+
func (fs FSCache) PutBlob(_ context.Context, blobID string, blobInfo types.BlobInfo) error {
9596
b, err := json.Marshal(blobInfo)
9697
if err != nil {
9798
return xerrors.Errorf("unable to marshal blob JSON (%s): %w", blobID, err)
@@ -111,7 +112,7 @@ func (fs FSCache) PutBlob(blobID string, blobInfo types.BlobInfo) error {
111112
}
112113

113114
// GetArtifact gets artifact information such as image metadata from local cache
114-
func (fs FSCache) GetArtifact(artifactID string) (types.ArtifactInfo, error) {
115+
func (fs FSCache) GetArtifact(_ context.Context, artifactID string) (types.ArtifactInfo, error) {
115116
var blob []byte
116117
err := fs.db.View(func(tx *bolt.Tx) error {
117118
artifactBucket := tx.Bucket([]byte(artifactBucket))
@@ -130,7 +131,7 @@ func (fs FSCache) GetArtifact(artifactID string) (types.ArtifactInfo, error) {
130131
}
131132

132133
// DeleteBlobs removes blobs by IDs
133-
func (fs FSCache) DeleteBlobs(blobIDs []string) error {
134+
func (fs FSCache) DeleteBlobs(_ context.Context, blobIDs []string) error {
134135
var errs error
135136
err := fs.db.Update(func(tx *bolt.Tx) error {
136137
bucket := tx.Bucket([]byte(blobBucket))
@@ -148,7 +149,7 @@ func (fs FSCache) DeleteBlobs(blobIDs []string) error {
148149
}
149150

150151
// PutArtifact stores artifact information such as image metadata in local cache
151-
func (fs FSCache) PutArtifact(artifactID string, artifactInfo types.ArtifactInfo) (err error) {
152+
func (fs FSCache) PutArtifact(_ context.Context, artifactID string, artifactInfo types.ArtifactInfo) (err error) {
152153
b, err := json.Marshal(artifactInfo)
153154
if err != nil {
154155
return xerrors.Errorf("unable to marshal artifact JSON (%s): %w", artifactID, err)
@@ -169,7 +170,7 @@ func (fs FSCache) PutArtifact(artifactID string, artifactInfo types.ArtifactInfo
169170
}
170171

171172
// MissingBlobs returns missing blob IDs such as layer IDs
172-
func (fs FSCache) MissingBlobs(artifactID string, blobIDs []string) (bool, []string, error) {
173+
func (fs FSCache) MissingBlobs(ctx context.Context, artifactID string, blobIDs []string) (bool, []string, error) {
173174
var missingArtifact bool
174175
var missingBlobIDs []string
175176
err := fs.db.View(func(tx *bolt.Tx) error {
@@ -192,7 +193,7 @@ func (fs FSCache) MissingBlobs(artifactID string, blobIDs []string) (bool, []str
192193
}
193194

194195
// get artifact info
195-
artifactInfo, err := fs.GetArtifact(artifactID)
196+
artifactInfo, err := fs.GetArtifact(ctx, artifactID)
196197
if err != nil {
197198
// error means cache missed artifact info
198199
return true, missingBlobIDs, nil
@@ -212,7 +213,7 @@ func (fs FSCache) Close() error {
212213
}
213214

214215
// Clear removes the database
215-
func (fs FSCache) Clear() error {
216+
func (fs FSCache) Clear(_ context.Context) error {
216217
if err := fs.Close(); err != nil {
217218
return err
218219
}

‎pkg/cache/fs_test.go‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -100,11 +100,11 @@ func TestFSCache_GetBlob(t *testing.T) {
100100
fs, err := NewFSCache(tmpDir)
101101
require.NoError(t, err)
102102
defer func() {
103-
_ = fs.Clear()
103+
_ = fs.Clear(t.Context())
104104
_ = fs.Close()
105105
}()
106106

107-
got, err := fs.GetBlob(tt.args.layerID)
107+
got, err := fs.GetBlob(t.Context(), tt.args.layerID)
108108
assert.Equal(t, tt.wantErr, err != nil, err)
109109
assert.Equal(t, tt.want, got)
110110
})
@@ -276,15 +276,15 @@ func TestFSCache_PutBlob(t *testing.T) {
276276
fs, err := NewFSCache(tmpDir)
277277
require.NoError(t, err)
278278
defer func() {
279-
_ = fs.Clear()
279+
_ = fs.Clear(t.Context())
280280
_ = fs.Close()
281281
}()
282282

283283
if strings.HasPrefix(tt.name, "sad") {
284284
require.NoError(t, fs.Close())
285285
}
286286

287-
err = fs.PutBlob(tt.args.diffID, tt.args.layerInfo)
287+
err = fs.PutBlob(t.Context(), tt.args.diffID, tt.args.layerInfo)
288288
if tt.wantErr != "" {
289289
require.ErrorContains(t, err, tt.wantErr, tt.name)
290290
return
@@ -356,11 +356,11 @@ func TestFSCache_PutArtifact(t *testing.T) {
356356
fs, err := NewFSCache(tmpDir)
357357
require.NoError(t, err)
358358
defer func() {
359-
_ = fs.Clear()
359+
_ = fs.Clear(t.Context())
360360
_ = fs.Close()
361361
}()
362362

363-
err = fs.PutArtifact(tt.args.imageID, tt.args.imageConfig)
363+
err = fs.PutArtifact(t.Context(), tt.args.imageID, tt.args.imageConfig)
364364
if tt.wantErr != "" {
365365
require.ErrorContains(t, err, tt.wantErr, tt.name)
366366
return
@@ -475,11 +475,11 @@ func TestFSCache_MissingBlobs(t *testing.T) {
475475
fs, err := NewFSCache(tmpDir)
476476
require.NoError(t, err)
477477
defer func() {
478-
_ = fs.Clear()
478+
_ = fs.Clear(t.Context())
479479
_ = fs.Close()
480480
}()
481481

482-
gotMissingImage, gotMissingLayerIDs, err := fs.MissingBlobs(tt.args.imageID, tt.args.layerIDs)
482+
gotMissingImage, gotMissingLayerIDs, err := fs.MissingBlobs(t.Context(), tt.args.imageID, tt.args.layerIDs)
483483
if tt.wantErr != "" {
484484
assert.ErrorContains(t, err, tt.wantErr, tt.name)
485485
return

‎pkg/cache/memory.go‎

Lines changed: 10 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package cache
22

33
import (
4+
"context"
45
"sync"
56

67
"golang.org/x/xerrors"
@@ -20,27 +21,27 @@ func NewMemoryCache() *MemoryCache {
2021
}
2122

2223
// PutArtifact stores the artifact information in the memory cache
23-
func (c *MemoryCache) PutArtifact(artifactID string, artifactInfo types.ArtifactInfo) error {
24+
func (c *MemoryCache) PutArtifact(_ context.Context, artifactID string, artifactInfo types.ArtifactInfo) error {
2425
c.artifacts.Store(artifactID, artifactInfo)
2526
return nil
2627
}
2728

2829
// PutBlob stores the blob information in the memory cache
29-
func (c *MemoryCache) PutBlob(blobID string, blobInfo types.BlobInfo) error {
30+
func (c *MemoryCache) PutBlob(_ context.Context, blobID string, blobInfo types.BlobInfo) error {
3031
c.blobs.Store(blobID, blobInfo)
3132
return nil
3233
}
3334

3435
// DeleteBlobs removes the specified blobs from the memory cache
35-
func (c *MemoryCache) DeleteBlobs(blobIDs []string) error {
36+
func (c *MemoryCache) DeleteBlobs(_ context.Context, blobIDs []string) error {
3637
for _, blobID := range blobIDs {
3738
c.blobs.Delete(blobID)
3839
}
3940
return nil
4041
}
4142

4243
// GetArtifact retrieves the artifact information from the memory cache
43-
func (c *MemoryCache) GetArtifact(artifactID string) (types.ArtifactInfo, error) {
44+
func (c *MemoryCache) GetArtifact(_ context.Context, artifactID string) (types.ArtifactInfo, error) {
4445
info, ok := c.artifacts.Load(artifactID)
4546
if !ok {
4647
return types.ArtifactInfo{}, xerrors.Errorf("artifact (%s) not found in memory cache", artifactID)
@@ -53,7 +54,7 @@ func (c *MemoryCache) GetArtifact(artifactID string) (types.ArtifactInfo, error)
5354
}
5455

5556
// GetBlob retrieves the blob information from the memory cache
56-
func (c *MemoryCache) GetBlob(blobID string) (types.BlobInfo, error) {
57+
func (c *MemoryCache) GetBlob(_ context.Context, blobID string) (types.BlobInfo, error) {
5758
info, ok := c.blobs.Load(blobID)
5859
if !ok {
5960
return types.BlobInfo{}, xerrors.Errorf("blob (%s) not found in memory cache", blobID)
@@ -66,16 +67,16 @@ func (c *MemoryCache) GetBlob(blobID string) (types.BlobInfo, error) {
6667
}
6768

6869
// MissingBlobs determines the missing artifact and blob information in the memory cache
69-
func (c *MemoryCache) MissingBlobs(artifactID string, blobIDs []string) (bool, []string, error) {
70+
func (c *MemoryCache) MissingBlobs(ctx context.Context, artifactID string, blobIDs []string) (bool, []string, error) {
7071
var missingArtifact bool
7172
var missingBlobIDs []string
7273

73-
if _, err := c.GetArtifact(artifactID); err != nil {
74+
if _, err := c.GetArtifact(ctx, artifactID); err != nil {
7475
missingArtifact = true
7576
}
7677

7778
for _, blobID := range blobIDs {
78-
if _, err := c.GetBlob(blobID); err != nil {
79+
if _, err := c.GetBlob(ctx, blobID); err != nil {
7980
missingBlobIDs = append(missingBlobIDs, blobID)
8081
}
8182
}
@@ -91,7 +92,7 @@ func (c *MemoryCache) Close() error {
9192
}
9293

9394
// Clear clears the artifact and blob information from the memory cache
94-
func (c *MemoryCache) Clear() error {
95+
func (c *MemoryCache) Clear(_ context.Context) error {
9596
c.artifacts = sync.Map{}
9697
c.blobs = sync.Map{}
9798
return nil

0 commit comments

Comments
 (0)