actions/trufflehog
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Extend memory cache (#2275)

Browse Source 
* Extend memory cache to allow for configuring custom expiration and purge interval

* use any for value type

* fix test

* fix test

* address comments

* address

* make new construct more clear

* reduce duplication

* fix test
This commit is contained in:
ahrav
2024-01-11 08:20:37 -08:00
committed by GitHub
parent 5c84375aa1
commit 677238c96c
5 changed files with 83 additions and 41 deletions
Download Patch File Download Diff File Expand all files Collapse all files
pkg/cache/cache.go
Vendored
+3 -3
View File
@@ -4,9 +4,9 @@ package cache
// Cache is used to store key/value pairs.
type Cache interface {
// Set stores the given key/value pair.
Set(string, string)
Set(string, any)
// Get returns the value for the given key and a boolean indicating if the key was found.
Get(string) (string, bool)
Get(string) (any, bool)
// Exists returns true if the given key exists in the cache.
Exists(string) bool
// Delete the given key from the cache.
@@ -18,7 +18,7 @@ type Cache interface {
// Keys returns all keys in the cache.
Keys() []string
// Values returns all values in the cache.
Values() []string
Values() []any
// Contents returns all keys in the cache encoded as a string.
Contents() string
}
pkg/cache/memory/memory.go
Vendored
+54 -29
View File
@@ -5,52 +5,77 @@ import (
"time"
"github.com/patrickmn/go-cache"
"github.com/trufflesecurity/trufflehog/v3/pkg/context"
)
const (
expirationInterval = 12 * time.Hour
purgeInterval = 13 * time.Hour
defaultExpiration = cache.DefaultExpiration
defaultExpirationInterval = 12 * time.Hour
defaultPurgeInterval = 13 * time.Hour
defaultExpiration = cache.DefaultExpiration
)
// Cache is a wrapper around the go-cache library.
// Cache wraps the go-cache library to provide an in-memory key-value store.
type Cache struct {
c *cache.Cache
c *cache.Cache
expiration time.Duration
purgeInterval time.Duration
}
// New constructs a new in-memory cache.
func New() *Cache {
c := cache.New(expirationInterval, purgeInterval)
return &Cache{c: c}
// CacheOption defines a function type used for configuring a Cache.
type CacheOption func(*Cache)
// WithExpirationInterval returns a CacheOption to set the expiration interval of cache items.
// The interval determines the duration a cached item remains in the cache before it is expired.
func WithExpirationInterval(interval time.Duration) CacheOption {
return func(c *Cache) { c.expiration = interval }
}
// WithPurgeInterval returns a CacheOption to set the interval at which the cache purges expired items.
// Regular purging helps in freeing up memory by removing stale entries.
func WithPurgeInterval(interval time.Duration) CacheOption {
return func(c *Cache) { c.purgeInterval = interval }
}
// New constructs a new in-memory cache instance with optional configurations.
// By default, it sets the expiration and purge intervals to 12 and 13 hours, respectively.
// These defaults can be overridden using the functional options: WithExpirationInterval and WithPurgeInterval.
func New(opts ...CacheOption) *Cache {
return NewWithData(nil, opts...)
}
// CacheEntry represents a single entry in the cache, consisting of a key and its corresponding value.
type CacheEntry struct {
// Key is the unique identifier for the entry.
Key string
// Value is the data stored in the entry.
Value any
}
// NewWithData constructs a new in-memory cache with existing data.
func NewWithData(ctx context.Context, data []string) *Cache {
ctx.Logger().V(3).Info("Loading cache", "num-items", len(data))
items := make(map[string]cache.Item, len(data))
for _, d := range data {
items[d] = cache.Item{Object: d, Expiration: int64(defaultExpiration)}
// It also accepts CacheOption parameters to override default configuration values.
func NewWithData(data []CacheEntry, opts ...CacheOption) *Cache {
instance := &Cache{expiration: defaultExpirationInterval, purgeInterval: defaultPurgeInterval}
for _, opt := range opts {
opt(instance)
}
c := cache.NewFrom(expirationInterval, purgeInterval, items)
return &Cache{c: c}
// Convert data slice to map required by go-cache.
items := make(map[string]cache.Item, len(data))
for _, d := range data {
items[d.Key] = cache.Item{Object: d.Value, Expiration: int64(defaultExpiration)}
}
instance.c = cache.NewFrom(instance.expiration, instance.purgeInterval, items)
return instance
}
// Set adds a key-value pair to the cache.
func (c *Cache) Set(key, value string) {
func (c *Cache) Set(key string, value any) {
c.c.Set(key, value, defaultExpiration)
}
// Get returns the value for the given key.
func (c *Cache) Get(key string) (string, bool) {
res, ok := c.c.Get(key)
if !ok {
return "", ok
}
return res.(string), ok
func (c *Cache) Get(key string) (any, bool) {
return c.c.Get(key)
}
// Exists returns true if the given key exists in the cache.
@@ -85,11 +110,11 @@ func (c *Cache) Keys() []string {
}
// Values returns all values in the cache.
func (c *Cache) Values() []string {
func (c *Cache) Values() []any {
items := c.c.Items()
res := make([]string, 0, len(items))
res := make([]any, 0, len(items))
for _, v := range items {
res = append(res, v.Object.(string))
res = append(res, v.Object)
}
return res
}
pkg/cache/memory/memory_test.go
Vendored
+8 -6
View File
@@ -7,8 +7,6 @@ import (
"testing"
"github.com/google/go-cmp/cmp"
logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
)
func TestCache(t *testing.T) {
@@ -34,7 +32,7 @@ func TestCache(t *testing.T) {
// Test delete.
c.Delete("key1")
v, ok = c.Get("key1")
if ok || v != "" {
if ok || v != nil {
t.Fatalf("Unexpected value for key1 after delete: %v, %v", v, ok)
}
@@ -42,7 +40,7 @@ func TestCache(t *testing.T) {
c.Set("key10", "key10")
c.Clear()
v, ok = c.Get("key10")
if ok || v != "" {
if ok || v != nil {
t.Fatalf("Unexpected value for key10 after clear: %v, %v", v, ok)
}
@@ -60,7 +58,10 @@ func TestCache(t *testing.T) {
}
// Test getting only the values.
vals := c.Values()
vals := make([]string, 0, c.Count())
for _, v := range c.Values() {
vals = append(vals, v.(string))
}
sort.Strings(vals)
sort.Strings(values)
if !cmp.Equal(values, vals) {
@@ -82,7 +83,8 @@ func TestCache(t *testing.T) {
}
func TestCache_NewWithData(t *testing.T) {
c := NewWithData(logContext.Background(), []string{"key1", "key2", "key3"})
data := []CacheEntry{{"key1", "value1"}, {"key2", "value2"}, {"key3", "value3"}}
c := NewWithData(data)
// Test the count.
if c.Count() != 3 {
pkg/sources/gcs/gcs.go
+9 -2
View File
@@ -97,7 +97,7 @@ func newPersistableCache(increment int, cache cache.Cache, p *sources.Progress)
// Set overrides the cache Set method of the cache to enable the persistence
// of the cache contents the Progress of the source at given increments.
func (c *persistableCache) Set(key, val string) {
func (c *persistableCache) Set(key string, val any) {
c.Cache.Set(key, val)
if ok, contents := c.shouldPersist(); ok {
c.Progress.EncodedResumeInfo = contents
@@ -297,7 +297,14 @@ func (s *Source) Chunks(ctx context.Context, chunksChan chan *sources.Chunk, _ .
func (s *Source) setupCache(ctx context.Context) *persistableCache {
var c cache.Cache
if s.Progress.EncodedResumeInfo != "" {
c = memory.NewWithData(ctx, strings.Split(s.Progress.EncodedResumeInfo, ","))
keys := strings.Split(s.Progress.EncodedResumeInfo, ",")
entries := make([]memory.CacheEntry, len(keys))
for i, val := range keys {
entries[i] = memory.CacheEntry{Key: val, Value: val}
}
c = memory.NewWithData(entries)
ctx.Logger().V(3).Info("Loaded cache", "num_entries", len(entries))
} else {
c = memory.New()
}
pkg/sources/github/github.go
+9 -1
View File
@@ -472,7 +472,15 @@ func (s *Source) enumerate(ctx context.Context, apiEndpoint string) (*github.Cli
return nil, errors.Errorf("Invalid configuration given for source. Name: %s, Type: %s", s.name, s.Type())
}
s.repos = s.filteredRepoCache.Values()
s.repos = make([]string, 0, s.filteredRepoCache.Count())
for _, repo := range s.filteredRepoCache.Values() {
r, ok := repo.(string)
if !ok {
ctx.Logger().Error(fmt.Errorf("type assertion failed"), "unexpected value in cache", "repo", repo)
continue
}
s.repos = append(s.repos, r)
}
githubReposEnumerated.WithLabelValues(s.name).Set(float64(len(s.repos)))
s.log.Info("Completed enumeration", "num_repos", len(s.repos), "num_orgs", s.orgsCache.Count(), "num_members", len(s.memberCache))