Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add in-memory caching pkg #1189

Merged
merged 10 commits into from
Mar 20, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions go.mod
Original file line number Diff line number Diff line change
Expand Up @@ -45,6 +45,7 @@ require (
github.com/lib/pq v1.10.7
github.com/mattn/go-sqlite3 v1.14.16
github.com/mholt/archiver/v4 v4.0.0-alpha.7
github.com/patrickmn/go-cache v2.1.0+incompatible
github.com/paulbellamy/ratecounter v0.2.0
github.com/petar-dambovaliev/aho-corasick v0.0.0-20211021192214-5ab2d9280aa9
github.com/pkg/errors v0.9.1
Expand Down
2 changes: 2 additions & 0 deletions go.sum
Original file line number Diff line number Diff line change
Expand Up @@ -283,6 +283,8 @@ github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1y
github.com/onsi/gomega v1.23.0 h1:/oxKu9c2HVap+F3PfKort2Hw5DEU+HGlW8n+tguWsys=
github.com/onsi/gomega v1.23.0/go.mod h1:Z/NWtiqwBrwUt4/2loMmHL63EDLnYHmVbuBpDr2vQAg=
github.com/op/go-logging v0.0.0-20160315200505-970db520ece7/go.mod h1:HzydrMdWErDVzsI23lYNej1Htcns9BCg93Dk0bBINWk=
github.com/patrickmn/go-cache v2.1.0+incompatible h1:HRMgzkcYKYpi3C8ajMPV8OFXaaRUnok+kx1WdO15EQc=
github.com/patrickmn/go-cache v2.1.0+incompatible/go.mod h1:3Qf8kWWT7OJRJbdiICTKqZju1ZixQ/KpMGzzAfe6+WQ=
github.com/paulbellamy/ratecounter v0.2.0 h1:2L/RhJq+HA8gBQImDXtLPrDXK5qAj6ozWVK/zFXVJGs=
github.com/paulbellamy/ratecounter v0.2.0/go.mod h1:Hfx1hDpSGoqxkVVpBi/IlYD7kChlfo5C6hzIHwPqfFE=
github.com/petar-dambovaliev/aho-corasick v0.0.0-20211021192214-5ab2d9280aa9 h1:lL+y4Xv20pVlCGyLzNHRC0I0rIHhIL1lTvHizoS/dU8=
Expand Down
20 changes: 20 additions & 0 deletions pkg/cache/cache.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,20 @@
// Package cache provides an interface which can be implemented by different cache types.
package cache

// Cache is used to store key/value pairs.
type Cache interface {
// Set stores the given key/value pair.
Set(string, string)
// Get returns the value for the given key and a boolean indicating if the key was found.
Get(string) (string, bool)
// Exists returns true if the given key exists in the cache.
Exists(string) bool
// Delete the given key from the cache.
Delete(string)
// Clear all key/value pairs from the cache.
Clear()
// Count the number of key/value pairs in the cache.
Count() int
// Contents returns all keys in the cache encoded as a string.
Contents() string
}
85 changes: 85 additions & 0 deletions pkg/cache/memory/memory.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,85 @@
package memory

import (
"strings"
"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
)

// Cache is a wrapper around the go-cache library.
type Cache struct {
c *cache.Cache
}

// New constructs a new in-memory cache.
func New() *Cache {
c := cache.New(expirationInterval, purgeInterval)
return &Cache{c: c}
}

// 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)}
}

c := cache.NewFrom(expirationInterval, purgeInterval, items)
return &Cache{c: c}
}

// Set adds a key-value pair to the cache.
func (c *Cache) Set(key, value string) {
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
}

// Exists returns true if the given key exists in the cache.
func (c *Cache) Exists(key string) bool {
_, ok := c.c.Get(key)
return ok
}

// Delete removes the key-value pair from the cache.
func (c *Cache) Delete(key string) {
c.c.Delete(key)
}

// Clear removes all key-value pairs from the cache.
func (c *Cache) Clear() {
c.c.Flush()
}

// Count returns the number of key-value pairs in the cache.
func (c *Cache) Count() int {
return c.c.ItemCount()
}

// Contents returns all key-value pairs in the cache encodes as a string.
func (c *Cache) Contents() string {
items := c.c.Items()
res := make([]string, 0, len(items))
for k := range items {
res = append(res, k)
}
return strings.Join(res, ",")
}
ahrav marked this conversation as resolved.
Show resolved Hide resolved
Comment on lines +77 to +85
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggestion (optional): If you do want to add convenience methods for actually serializing/deserializing the data, we could use strings.Replacer to escape commas (assuming it has an acceptable benchmark).

For example:

func unescape(s string) string {
	replacer := strings.NewReplacer(`\,`, ",")
	return replacer.Replace(s)
}

func escape(s string) string {
	replacer := strings.NewReplacer(",", `\,`)
	return replacer.Replace(s)
}

153 changes: 153 additions & 0 deletions pkg/cache/memory/memory_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,153 @@
package memory

import (
"fmt"
"sort"
"strings"
"testing"

"github.com/google/go-cmp/cmp"

logContext "github.com/trufflesecurity/trufflehog/v3/pkg/context"
)

func TestCache(t *testing.T) {
c := New()

// Test set and get.
c.Set("key1", "key1")
v, ok := c.Get("key1")
if !ok || v != "key1" {
t.Fatalf("Unexpected value for key1: %v, %v", v, ok)
}

// Test exists.
if !c.Exists("key1") {
t.Fatalf("Expected key1 to exist")
}

// Test the count.
if c.Count() != 1 {
t.Fatalf("Unexpected count: %d", c.Count())
}

// Test delete.
c.Delete("key1")
v, ok = c.Get("key1")
if ok || v != "" {
t.Fatalf("Unexpected value for key1 after delete: %v, %v", v, ok)
}

// Test clear.
c.Set("key10", "key10")
c.Clear()
v, ok = c.Get("key10")
if ok || v != "" {
t.Fatalf("Unexpected value for key10 after clear: %v, %v", v, ok)
}

// Test contents.
keys := []string{"key1", "key2", "key3"}
for _, k := range keys {
c.Set(k, k)
}

items := c.Contents()
sort.Strings(keys)
res := strings.Split(items, ",")
sort.Strings(res)

if len(keys) != len(res) {
t.Fatalf("Unexpected length of items: %d", len(res))
}
if !cmp.Equal(keys, res) {
t.Fatalf("Unexpected items: %v", res)
}
}

func TestCache_NewWithData(t *testing.T) {
c := NewWithData(logContext.Background(), []string{"key1", "key2", "key3"})

// Test the count.
if c.Count() != 3 {
t.Fatalf("Unexpected count: %d", c.Count())
}

// Test contents.
keys := []string{"key1", "key2", "key3"}
items := c.Contents()
sort.Strings(keys)
res := strings.Split(items, ",")
sort.Strings(res)

if len(keys) != len(res) {
t.Fatalf("Unexpected length of items: %d", len(res))
}
if !cmp.Equal(keys, res) {
t.Fatalf("Unexpected items: %v", res)
}
}

func setupBenchmarks(b *testing.B) *Cache {
b.Helper()

c := New()

for i := 0; i < 500_000; i++ {
key := fmt.Sprintf("key%d", i)
c.Set(key, key)
}

return c
}

func BenchmarkSet(b *testing.B) {
c := New()

for i := 0; i < b.N; i++ {
key := fmt.Sprintf("key%d", i)
c.Set(key, key)
}
}

func BenchmarkGet(b *testing.B) {
c := setupBenchmarks(b)
b.ResetTimer()

for i := 0; i < b.N; i++ {
key := fmt.Sprintf("key%d", i)
c.Get(key)
}
}

func BenchmarkDelete(b *testing.B) {
c := setupBenchmarks(b)
b.ResetTimer()

for i := 0; i < b.N; i++ {
key := fmt.Sprintf("key%d", i)
c.Delete(key)
}
}

func BenchmarkCount(b *testing.B) {
c := setupBenchmarks(b)
b.ResetTimer()

for i := 0; i < b.N; i++ {
c.Count()
}
}

func BenchmarkContents(b *testing.B) {
c := setupBenchmarks(b)
b.ResetTimer()

var s string

for i := 0; i < b.N; i++ {
s = c.Contents()
}

_ = s
}