diff --git a/CHANGELOG.md b/CHANGELOG.md index abfbf2fbd0..447b78a1e2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -8,6 +8,8 @@ All notable changes to this project will be documented in this file. ### Added - `parse_big_decimal` bloblang method for Kafka Connect / Debezium decimal decoding @aratz-lasa +- `ListableCache` interface added to the service package. This is an optional interface for caches that are able to enumerate their keys, detected via type assertion on caches obtained from `*Resources.AccessCache` @ecordell +- `memory`, `file`, `aws_s3` and `gcp_cloud_storage` caches now implement `ListableCache` @ecordell ## 1.19.0 - 2026-07-10 diff --git a/internal/component/cache/cache_metrics.go b/internal/component/cache/cache_metrics.go index 950f2a9a80..b10c918b47 100644 --- a/internal/component/cache/cache_metrics.go +++ b/internal/component/cache/cache_metrics.go @@ -35,13 +35,14 @@ type metricsCache struct { } // MetricsForCache wraps a cache with a struct that adds standard metrics over -// each method. +// each method. If the cache implements KeyLister then the returned cache will +// also implement it. func MetricsForCache(c V1, stats metrics.Type) V1 { cacheSuccess := stats.GetCounterVec("cache_success", "operation") cacheError := stats.GetCounterVec("cache_error", "operation") cacheLatency := stats.GetTimerVec("cache_latency_ns", "operation") - return &metricsCache{ + mCache := &metricsCache{ c: c, sig: shutdown.NewSignaller(), mGetNotFound: stats.GetCounterVec("cache_not_found", "operation").With("get"), @@ -62,6 +63,19 @@ func MetricsForCache(c V1, stats metrics.Type) V1 { mDelSuccess: cacheSuccess.With("delete"), mDelLatency: cacheLatency.With("delete"), } + + kl, ok := c.(KeyLister) + if !ok { + return mCache + } + return &listableMetricsCache{ + metricsCache: mCache, + kl: kl, + + mListKeysError: cacheError.With("list_keys"), + mListKeysSuccess: cacheSuccess.With("list_keys"), + mListKeysLatency: cacheLatency.With("list_keys"), + } } func (a *metricsCache) Get(ctx context.Context, key string) ([]byte, error) { @@ -135,3 +149,26 @@ func (a *metricsCache) Delete(ctx context.Context, key string) error { func (a *metricsCache) Close(ctx context.Context) error { return a.c.Close(ctx) } + +//------------------------------------------------------------------------------ + +type listableMetricsCache struct { + *metricsCache + kl KeyLister + + mListKeysError metrics.StatCounter + mListKeysSuccess metrics.StatCounter + mListKeysLatency metrics.StatTimer +} + +func (a *listableMetricsCache) ListKeys(ctx context.Context) ([]string, error) { + started := time.Now() + keys, err := a.kl.ListKeys(ctx) + a.mListKeysLatency.Timing(int64(time.Since(started))) + if err != nil { + a.mListKeysError.Incr(1) + } else { + a.mListKeysSuccess.Incr(1) + } + return keys, err +} diff --git a/internal/component/cache/cache_metrics_test.go b/internal/component/cache/cache_metrics_test.go index 1aeb2e4ef7..a8d74097f8 100644 --- a/internal/component/cache/cache_metrics_test.go +++ b/internal/component/cache/cache_metrics_test.go @@ -6,6 +6,7 @@ import ( "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/warpstreamlabs/bento/internal/component" "github.com/warpstreamlabs/bento/internal/component/metrics" @@ -253,3 +254,52 @@ func TestCacheAirGapDelete(t *testing.T) { assert.NoError(t, err) assert.Equal(t, map[string]testCacheItem{}, rl.m) } + +type listableClosableCache struct { + *closableCache +} + +func (c *listableClosableCache) ListKeys(ctx context.Context) ([]string, error) { + if c.err != nil { + return nil, c.err + } + keys := make([]string, 0, len(c.m)) + for k := range c.m { + keys = append(keys, k) + } + return keys, nil +} + +func TestCacheAirGapListKeys(t *testing.T) { + ctx := context.Background() + rl := &listableClosableCache{ + closableCache: &closableCache{ + m: map[string]testCacheItem{ + "foo": { + b: []byte("bar"), + }, + "baz": { + b: []byte("qux"), + }, + }, + }, + } + agrl := MetricsForCache(rl, metrics.Noop()) + + kl, ok := agrl.(KeyLister) + require.True(t, ok) + + keys, err := kl.ListKeys(ctx) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"foo", "baz"}, keys) +} + +func TestCacheAirGapNotListable(t *testing.T) { + rl := &closableCache{ + m: map[string]testCacheItem{}, + } + agrl := MetricsForCache(rl, metrics.Noop()) + + _, ok := agrl.(KeyLister) + assert.False(t, ok) +} diff --git a/internal/component/cache/interface.go b/internal/component/cache/interface.go index 872bd2dc0a..d730e0cb58 100644 --- a/internal/component/cache/interface.go +++ b/internal/component/cache/interface.go @@ -11,6 +11,13 @@ type TTLItem struct { TTL *time.Duration } +// KeyLister is an optional interface implemented by caches that are capable +// of enumerating the keys they hold. +type KeyLister interface { + // ListKeys returns a slice of all keys currently held by the cache. + ListKeys(ctx context.Context) ([]string, error) +} + // V1 Defines a common interface of cache implementations. type V1 interface { // Get attempts to locate and return a cached value by its key, returns an diff --git a/internal/impl/aws/cache_s3.go b/internal/impl/aws/cache_s3.go index d48d2b1c05..e7a3fb389c 100644 --- a/internal/impl/aws/cache_s3.go +++ b/internal/impl/aws/cache_s3.go @@ -223,6 +223,40 @@ func (s *s3Cache) Delete(ctx context.Context, key string) (err error) { } } +func (s *s3Cache) ListKeys(ctx context.Context) ([]string, error) { + boff := s.boffPool.Get().(backoff.BackOff) + defer func() { + boff.Reset() + s.boffPool.Put(boff) + }() + + var keys []string + pager := s3.NewListObjectsV2Paginator(s.s3, &s3.ListObjectsV2Input{ + Bucket: &s.bucket, + }) + for pager.HasMorePages() { + page, err := pager.NextPage(ctx) + if err != nil { + wait := boff.NextBackOff() + if wait == backoff.Stop { + return nil, err + } + select { + case <-time.After(wait): + case <-ctx.Done(): + return nil, err + } + continue + } + for _, obj := range page.Contents { + if obj.Key != nil { + keys = append(keys, *obj.Key) + } + } + } + return keys, nil +} + func (s *s3Cache) Close(context.Context) error { return nil } diff --git a/internal/impl/aws/integration_s3_test.go b/internal/impl/aws/integration_s3_test.go index 6f38d7fbc4..43cfdeca3c 100644 --- a/internal/impl/aws/integration_s3_test.go +++ b/internal/impl/aws/integration_s3_test.go @@ -379,6 +379,7 @@ cache_resources: integration.CacheTestDoubleAdd(), integration.CacheTestDelete(), integration.CacheTestGetAndSet(1), + integration.CacheTestListKeys(5), ) suite.Run( t, template, diff --git a/internal/impl/gcp/cache_cloud_storage.go b/internal/impl/gcp/cache_cloud_storage.go index 60bddaeb10..3e4486e7be 100644 --- a/internal/impl/gcp/cache_cloud_storage.go +++ b/internal/impl/gcp/cache_cloud_storage.go @@ -7,6 +7,7 @@ import ( "time" "cloud.google.com/go/storage" + "google.golang.org/api/iterator" "github.com/warpstreamlabs/bento/public/service" ) @@ -129,6 +130,22 @@ func (c *gcpCloudStorageCache) Delete(ctx context.Context, key string) error { return c.bucketHandle.Object(key).Delete(ctx) } +func (c *gcpCloudStorageCache) ListKeys(ctx context.Context) ([]string, error) { + var keys []string + it := c.bucketHandle.Objects(ctx, nil) + for { + attrs, err := it.Next() + if errors.Is(err, iterator.Done) { + break + } + if err != nil { + return nil, err + } + keys = append(keys, attrs.Name) + } + return keys, nil +} + func (c *gcpCloudStorageCache) Close(ctx context.Context) error { return nil } diff --git a/internal/impl/io/cache_file.go b/internal/impl/io/cache_file.go index a412eb70e1..e65c054d94 100644 --- a/internal/impl/io/cache_file.go +++ b/internal/impl/io/cache_file.go @@ -3,6 +3,7 @@ package io import ( "context" "errors" + "fmt" "io/fs" "os" "path/filepath" @@ -88,6 +89,44 @@ func (f *fileCache) Delete(_ context.Context, key string) error { return f.mgr.FS().Remove(filepath.Join(f.dir, key)) } +func (f *fileCache) ListKeys(ctx context.Context) ([]string, error) { + var keys []string + var walk func(rel string) error + walk = func(rel string) error { + if err := ctx.Err(); err != nil { + return err + } + dir, err := f.mgr.FS().Open(filepath.Join(f.dir, rel)) + if err != nil { + return err + } + defer dir.Close() + + rd, ok := dir.(fs.ReadDirFile) + if !ok { + return fmt.Errorf("unable to list directory: %v", filepath.Join(f.dir, rel)) + } + entries, err := rd.ReadDir(-1) + if err != nil { + return err + } + for _, e := range entries { + if e.IsDir() { + if err := walk(filepath.Join(rel, e.Name())); err != nil { + return err + } + } else { + keys = append(keys, filepath.Join(rel, e.Name())) + } + } + return nil + } + if err := walk(""); err != nil { + return nil, err + } + return keys, nil +} + func (f *fileCache) Close(context.Context) error { return nil } diff --git a/internal/impl/io/cache_file_test.go b/internal/impl/io/cache_file_test.go index 65ae19ac86..eeee781120 100644 --- a/internal/impl/io/cache_file_test.go +++ b/internal/impl/io/cache_file_test.go @@ -3,6 +3,7 @@ package io import ( "context" "os" + "path/filepath" "testing" "github.com/stretchr/testify/assert" @@ -48,3 +49,42 @@ func TestFileCache(t *testing.T) { _, err = c.Get(tCtx, "foo") assert.Equal(t, service.ErrKeyNotFound, err) } + +func TestFileCacheListKeys(t *testing.T) { + dir := t.TempDir() + + tCtx := context.Background() + c := newFileCache(dir, service.MockResources()) + + var _ service.ListableCache = c + + require.NoError(t, c.Set(tCtx, "foo", []byte("1"), nil)) + require.NoError(t, c.Set(tCtx, "bar", []byte("2"), nil)) + + nestedKey := filepath.Join("nested", "baz") + require.NoError(t, os.MkdirAll(filepath.Join(dir, "nested"), 0o755)) + require.NoError(t, c.Set(tCtx, nestedKey, []byte("3"), nil)) + + keys, err := c.ListKeys(tCtx) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"foo", "bar", nestedKey}, keys) + + require.NoError(t, c.Delete(tCtx, "foo")) + + keys, err = c.ListKeys(tCtx) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"bar", nestedKey}, keys) +} + +func TestFileCacheListKeysCancelled(t *testing.T) { + dir := t.TempDir() + + c := newFileCache(dir, service.MockResources()) + require.NoError(t, c.Set(context.Background(), "foo", []byte("1"), nil)) + + ctx, cancel := context.WithCancel(context.Background()) + cancel() + + _, err := c.ListKeys(ctx) + require.ErrorIs(t, err, context.Canceled) +} diff --git a/internal/impl/pure/cache_memory.go b/internal/impl/pure/cache_memory.go index 07d05c4ec4..2787e32baf 100644 --- a/internal/impl/pure/cache_memory.go +++ b/internal/impl/pure/cache_memory.go @@ -239,6 +239,21 @@ func (m *memoryCache) Delete(_ context.Context, key string) error { return nil } +func (m *memoryCache) ListKeys(_ context.Context) ([]string, error) { + var keys []string + for _, shard := range m.shards { + shard.RLock() + for k, v := range shard.items { + // Simulate compaction by omitting keys with an expired ttl. + if !shard.isExpired(v) { + keys = append(keys, k) + } + } + shard.RUnlock() + } + return keys, nil +} + func (m *memoryCache) Close(context.Context) error { return nil } diff --git a/internal/impl/pure/cache_memory_test.go b/internal/impl/pure/cache_memory_test.go index 7451991239..f26386bda7 100644 --- a/internal/impl/pure/cache_memory_test.go +++ b/internal/impl/pure/cache_memory_test.go @@ -73,6 +73,49 @@ func TestMemoryCache(t *testing.T) { } } +func TestMemoryCacheListKeys(t *testing.T) { + ctx := context.Background() + + for _, nShards := range []int{1, 16} { + t.Run(fmt.Sprintf("%v shards", nShards), func(t *testing.T) { + c := newMemCache(0, 0, nShards, map[string]string{"foo": "1"}) + + var _ service.ListableCache = c + + require.NoError(t, c.Set(ctx, "bar", []byte("2"), nil)) + require.NoError(t, c.Add(ctx, "baz", []byte("3"), nil)) + + keys, err := c.ListKeys(ctx) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"foo", "bar", "baz"}, keys) + + require.NoError(t, c.Delete(ctx, "bar")) + + keys, err = c.ListKeys(ctx) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"foo", "baz"}, keys) + }) + } +} + +func TestMemoryCacheListKeysExpired(t *testing.T) { + ctx := context.Background() + + // A long compaction interval so that expired items are retained but + // considered expired by reads. + c := newMemCache(time.Hour, time.Hour, 1, nil) + + ttl := time.Millisecond + require.NoError(t, c.Set(ctx, "expires", []byte("1"), &ttl)) + require.NoError(t, c.Set(ctx, "remains", []byte("2"), nil)) + + <-time.After(time.Millisecond * 50) + + keys, err := c.ListKeys(ctx) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"remains"}, keys) +} + func TestMemoryCacheCompaction(t *testing.T) { defConf, err := memCacheConfig().ParseYAML(` default_ttl: 0s diff --git a/internal/manager/mock/cache.go b/internal/manager/mock/cache.go index d9ccac170f..dd218a3feb 100644 --- a/internal/manager/mock/cache.go +++ b/internal/manager/mock/cache.go @@ -66,6 +66,15 @@ func (c *Cache) Delete(ctx context.Context, key string) error { return nil } +// ListKeys returns all keys of the mock cache. +func (c *Cache) ListKeys(ctx context.Context) ([]string, error) { + keys := make([]string, 0, len(c.Values)) + for k := range c.Values { + keys = append(keys, k) + } + return keys, nil +} + // Close does nothing. func (c *Cache) Close(ctx context.Context) error { return nil diff --git a/public/service/cache.go b/public/service/cache.go index 658649b3e5..bec6d0da69 100644 --- a/public/service/cache.go +++ b/public/service/cache.go @@ -39,6 +39,19 @@ type Cache interface { Closer } +// ListableCache is an optional interface that can be implemented by caches +// that are capable of enumerating the keys they hold. Not all caches are able +// to support this, so in order to detect whether a cache obtained from +// *Resources.AccessCache is listable perform a type assertion against this +// interface. +type ListableCache interface { + Cache + + // ListKeys returns a slice of all keys currently held by the cache. The + // order of the returned keys is not guaranteed. + ListKeys(ctx context.Context) ([]string, error) +} + // CacheItem represents an individual cache item. type CacheItem struct { Key string @@ -66,7 +79,11 @@ type airGapCache struct { func newAirGapCache(c Cache, stats metrics.Type) cache.V1 { ag := &airGapCache{c: c, cm: nil} ag.cm, _ = c.(batchedCache) - return cache.MetricsForCache(ag, stats) + var v1 cache.V1 = ag + if lc, ok := c.(ListableCache); ok { + v1 = &listableAirGapCache{airGapCache: ag, l: lc} + } + return cache.MetricsForCache(v1, stats) } func (a *airGapCache) Get(ctx context.Context, key string) ([]byte, error) { @@ -117,6 +134,16 @@ func (a *airGapCache) Close(ctx context.Context) error { return a.c.Close(ctx) } +// Implements types.Cache and cache.KeyLister around a ListableCache. +type listableAirGapCache struct { + *airGapCache + l ListableCache +} + +func (a *listableAirGapCache) ListKeys(ctx context.Context) ([]string, error) { + return a.l.ListKeys(ctx) +} + //------------------------------------------------------------------------------ // Implements Cache around a types.Cache. @@ -124,8 +151,12 @@ type reverseAirGapCache struct { c cache.V1 } -func newReverseAirGapCache(c cache.V1) *reverseAirGapCache { - return &reverseAirGapCache{c} +func newReverseAirGapCache(c cache.V1) Cache { + rag := &reverseAirGapCache{c} + if kl, ok := c.(cache.KeyLister); ok { + return &reverseAirGapListableCache{reverseAirGapCache: rag, kl: kl} + } + return rag } func (r *reverseAirGapCache) Get(ctx context.Context, key string) ([]byte, error) { @@ -154,3 +185,14 @@ func (r *reverseAirGapCache) Delete(ctx context.Context, key string) error { func (r *reverseAirGapCache) Close(ctx context.Context) error { return r.c.Close(ctx) } + +// Implements ListableCache around a types.Cache that also implements +// cache.KeyLister. +type reverseAirGapListableCache struct { + *reverseAirGapCache + kl cache.KeyLister +} + +func (r *reverseAirGapListableCache) ListKeys(ctx context.Context) ([]string, error) { + return r.kl.ListKeys(ctx) +} diff --git a/public/service/cache_test.go b/public/service/cache_test.go index 8db4aa436f..60648424e1 100644 --- a/public/service/cache_test.go +++ b/public/service/cache_test.go @@ -7,6 +7,7 @@ import ( "time" "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" "github.com/warpstreamlabs/bento/internal/component" "github.com/warpstreamlabs/bento/internal/component/cache" @@ -283,6 +284,55 @@ func TestCacheAirGapAddWithTTL(t *testing.T) { assert.EqualError(t, err, "key already exists") } +type listableClosableCache struct { + *closableCache +} + +func (c *listableClosableCache) ListKeys(ctx context.Context) ([]string, error) { + if c.err != nil { + return nil, c.err + } + keys := make([]string, 0, len(c.m)) + for k := range c.m { + keys = append(keys, k) + } + return keys, nil +} + +func TestCacheAirGapListKeys(t *testing.T) { + ctx := context.Background() + rl := &listableClosableCache{ + closableCache: &closableCache{ + m: map[string]testCacheItem{ + "foo": { + b: []byte("bar"), + }, + "baz": { + b: []byte("qux"), + }, + }, + }, + } + agrl := newAirGapCache(rl, metrics.Noop()) + + kl, ok := agrl.(cache.KeyLister) + require.True(t, ok) + + keys, err := kl.ListKeys(ctx) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"foo", "baz"}, keys) +} + +func TestCacheAirGapNotListable(t *testing.T) { + rl := &closableCache{ + m: map[string]testCacheItem{}, + } + agrl := newAirGapCache(rl, metrics.Noop()) + + _, ok := agrl.(cache.KeyLister) + assert.False(t, ok) +} + func TestCacheAirGapDelete(t *testing.T) { ctx := context.Background() rl := &closableCache{ @@ -490,3 +540,82 @@ func TestCacheReverseAirGapDelete(t *testing.T) { assert.NoError(t, err) assert.Equal(t, map[string]testCacheItem{}, rl.m) } + +type listableClosableCacheType struct { + *closableCacheType +} + +func (c *listableClosableCacheType) ListKeys(ctx context.Context) ([]string, error) { + if c.err != nil { + return nil, c.err + } + keys := make([]string, 0, len(c.m)) + for k := range c.m { + keys = append(keys, k) + } + return keys, nil +} + +func TestCacheReverseAirGapListKeys(t *testing.T) { + rl := &listableClosableCacheType{ + closableCacheType: &closableCacheType{ + m: map[string]testCacheItem{ + "foo": { + b: []byte("bar"), + }, + "baz": { + b: []byte("qux"), + }, + }, + }, + } + agrl := newReverseAirGapCache(rl) + + lc, ok := agrl.(ListableCache) + require.True(t, ok) + + keys, err := lc.ListKeys(context.Background()) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"foo", "baz"}, keys) +} + +func TestCacheReverseAirGapNotListable(t *testing.T) { + rl := &closableCacheType{ + m: map[string]testCacheItem{}, + } + agrl := newReverseAirGapCache(rl) + + _, ok := agrl.(ListableCache) + assert.False(t, ok) +} + +// TestCacheListKeysRoundTrip covers the full journey of a ListableCache +// implementation registered as a plugin and then accessed as a resource, +// which passes through the air gap, metrics and reverse air gap wrappers. +func TestCacheListKeysRoundTrip(t *testing.T) { + rl := &listableClosableCache{ + closableCache: &closableCache{ + m: map[string]testCacheItem{ + "foo": { + b: []byte("bar"), + }, + }, + }, + } + agrl := newReverseAirGapCache(newAirGapCache(rl, metrics.Noop())) + + lc, ok := agrl.(ListableCache) + require.True(t, ok) + + keys, err := lc.ListKeys(context.Background()) + require.NoError(t, err) + assert.ElementsMatch(t, []string{"foo"}, keys) + + rlNotListable := &closableCache{ + m: map[string]testCacheItem{}, + } + agrlNotListable := newReverseAirGapCache(newAirGapCache(rlNotListable, metrics.Noop())) + + _, ok = agrlNotListable.(ListableCache) + assert.False(t, ok) +} diff --git a/public/service/integration/cache_test_definitions.go b/public/service/integration/cache_test_definitions.go index f949ff7c23..35b51e0a82 100644 --- a/public/service/integration/cache_test_definitions.go +++ b/public/service/integration/cache_test_definitions.go @@ -10,6 +10,7 @@ import ( "github.com/stretchr/testify/require" "github.com/warpstreamlabs/bento/internal/component" + componentcache "github.com/warpstreamlabs/bento/internal/component/cache" ) // CacheTestOpenClose checks that the cache can be started, an item added, and @@ -95,6 +96,34 @@ func CacheTestDelete() CacheTestDefinition { ) } +// CacheTestListKeys checks that a cache implementing the optional KeyLister +// interface enumerates the keys it holds after n items are set. +func CacheTestListKeys(n int) CacheTestDefinition { + return namedCacheTest( + "can list keys", + func(t *testing.T, env *cacheTestEnvironment) { + cache := initCache(t, env) + t.Cleanup(func() { + closeCache(t, cache) + }) + + kl, ok := cache.(componentcache.KeyLister) + require.True(t, ok, "cache does not support listing keys") + + expKeys := make([]string, 0, n) + for i := range n { + key := fmt.Sprintf("listkey%v", i) + require.NoError(t, cache.Set(env.ctx, key, fmt.Appendf(nil, "value%v", i), nil)) + expKeys = append(expKeys, key) + } + + keys, err := kl.ListKeys(env.ctx) + require.NoError(t, err) + assert.ElementsMatch(t, expKeys, keys) + }, + ) +} + // CacheTestGetAndSet checks that we can set and then get n items. func CacheTestGetAndSet(n int) CacheTestDefinition { return namedCacheTest(