From d4f1e88b653141d183c3402c36deefd7e6866457 Mon Sep 17 00:00:00 2001 From: testinginprod <98415576+testinginprod@users.noreply.github.com> Date: Tue, 20 Jun 2023 15:44:06 +0200 Subject: [PATCH] feat(collections): Add Clear method on Map and KeySet (#16618) Co-authored-by: unknown unknown --- collections/CHANGELOG.md | 4 +++ collections/iter.go | 29 +++++++++++++------- collections/keyset.go | 6 +++++ collections/map.go | 57 +++++++++++++++++++++++++++++++++++++--- collections/map_test.go | 35 ++++++++++++++++++++++++ 5 files changed, 118 insertions(+), 13 deletions(-) diff --git a/collections/CHANGELOG.md b/collections/CHANGELOG.md index 101b5146f3..39169cdb5c 100644 --- a/collections/CHANGELOG.md +++ b/collections/CHANGELOG.md @@ -31,6 +31,10 @@ Ref: https://keepachangelog.com/en/1.0.0/ ## [Unreleased] +### Features + +* [#16074](https://github.com/cosmos/cosmos-sdk/pull/16607) - Introduces `Clear` method for `Map` and `KeySet` + ## [v0.2.0](https://github.com/cosmos/cosmos-sdk/releases/tag/collections%2Fv0.2.0) ### Features diff --git a/collections/iter.go b/collections/iter.go index 38a21a22a3..7958e3c27e 100644 --- a/collections/iter.go +++ b/collections/iter.go @@ -130,39 +130,48 @@ func (r *Range[K]) RangeValues() (start, end *RangeKey[K], order Order, err erro return r.start, r.end, r.order, nil } -// iteratorFromRanger generates an Iterator instance, with the proper prefixing and ranging. -// a nil Ranger can be seen as an ascending iteration over all the possible keys. -func iteratorFromRanger[K, V any](ctx context.Context, m Map[K, V], r Ranger[K]) (iter Iterator[K, V], err error) { +// parseRangeInstruction converts a Ranger into start bytes, end bytes and order of a store iteration. +func parseRangeInstruction[K any](prefix []byte, keyCodec codec.KeyCodec[K], r Ranger[K]) ([]byte, []byte, Order, error) { var ( start *RangeKey[K] end *RangeKey[K] order = OrderAscending + err error ) if r != nil { start, end, order, err = r.RangeValues() if err != nil { - return iter, err + return nil, nil, 0, err } } - startBytes := m.prefix + startBytes := prefix if start != nil { - startBytes, err = encodeRangeBound(m.prefix, m.kc, start) + startBytes, err = encodeRangeBound(prefix, keyCodec, start) if err != nil { - return iter, err + return nil, nil, 0, err } } var endBytes []byte if end != nil { - endBytes, err = encodeRangeBound(m.prefix, m.kc, end) + endBytes, err = encodeRangeBound(prefix, keyCodec, end) if err != nil { - return iter, err + return nil, nil, 0, err } } else { - endBytes = nextBytesPrefixKey(m.prefix) + endBytes = nextBytesPrefixKey(prefix) } + return startBytes, endBytes, order, nil +} +// iteratorFromRanger generates an Iterator instance, with the proper prefixing and ranging. +// a nil Ranger can be seen as an ascending iteration over all the possible keys. +func iteratorFromRanger[K, V any](ctx context.Context, m Map[K, V], r Ranger[K]) (iter Iterator[K, V], err error) { + startBytes, endBytes, order, err := parseRangeInstruction(m.prefix, m.kc, r) + if err != nil { + return Iterator[K, V]{}, err + } return newIterator(ctx, startBytes, endBytes, order, m) } diff --git a/collections/keyset.go b/collections/keyset.go index 408086b672..74f7d54044 100644 --- a/collections/keyset.go +++ b/collections/keyset.go @@ -57,6 +57,12 @@ func (k KeySet[K]) Walk(ctx context.Context, ranger Ranger[K], walkFunc func(key return (Map[K, NoValue])(k).Walk(ctx, ranger, func(key K, value NoValue) (bool, error) { return walkFunc(key) }) } +// Clear clears the KeySet using the provided Ranger. Refer to Map.Clear for +// behavioral documentation. +func (k KeySet[K]) Clear(ctx context.Context, ranger Ranger[K]) error { + return (Map[K, NoValue])(k).Clear(ctx, ranger) +} + func (k KeySet[K]) KeyCodec() codec.KeyCodec[K] { return (Map[K, NoValue])(k).KeyCodec() } func (k KeySet[K]) ValueCodec() codec.ValueCodec[NoValue] { return (Map[K, NoValue])(k).ValueCodec() } diff --git a/collections/map.go b/collections/map.go index 9ada53c494..2a7bb36a0f 100644 --- a/collections/map.go +++ b/collections/map.go @@ -5,7 +5,6 @@ import ( "fmt" "cosmossdk.io/collections/codec" - "cosmossdk.io/core/store" ) @@ -65,8 +64,7 @@ func (m Map[K, V]) Set(ctx context.Context, key K, value V) error { } kvStore := m.sa(ctx) - kvStore.Set(bytesKey, valueBytes) - return nil + return kvStore.Set(bytesKey, valueBytes) } // Get returns the value associated with the provided key, @@ -150,6 +148,59 @@ func (m Map[K, V]) Walk(ctx context.Context, ranger Ranger[K], walkFunc func(key return nil } +// Clear clears the collection contained within the provided key range. +// A nil ranger equals to clearing the whole collection. In case the collection +// is empty no error will be returned. +// NOTE: this API needs to be used with care, considering that as of today +// cosmos-sdk stores the deletion records to be committed in a memory cache, +// clearing a lot of data might make the node go OOM. +func (m Map[K, V]) Clear(ctx context.Context, ranger Ranger[K]) error { + startBytes, endBytes, _, err := parseRangeInstruction(m.prefix, m.kc, ranger) + if err != nil { + return err + } + return deleteDomain(m.sa(ctx), startBytes, endBytes) +} + +const clearBatchSize = 10000 + +// deleteDomain deletes the domain of an iterator, the key difference +// is that it uses batches to clear the store meaning that it will read +// the keys within the domain close the iterator and then delete them. +func deleteDomain(s store.KVStore, start, end []byte) error { + for { + iter, err := s.Iterator(start, end) + if err != nil { + return err + } + + keys := make([][]byte, 0, clearBatchSize) + for ; iter.Valid() && len(keys) < clearBatchSize; iter.Next() { + keys = append(keys, iter.Key()) + } + + // we close the iterator here instead of deferring + err = iter.Close() + if err != nil { + return err + } + + for _, key := range keys { + err = s.Delete(key) + if err != nil { + return err + } + } + + // If we've retrieved less than the batchSize, we're done. + if len(keys) < clearBatchSize { + break + } + } + + return nil +} + // IterateRaw iterates over the collection. The iteration range is untyped, it uses raw // bytes. The resulting Iterator is typed. // A nil start iterates from the first key contained in the collection. diff --git a/collections/map_test.go b/collections/map_test.go index 2887693d53..a011083a42 100644 --- a/collections/map_test.go +++ b/collections/map_test.go @@ -1,6 +1,7 @@ package collections import ( + "context" "testing" "github.com/stretchr/testify/require" @@ -36,6 +37,40 @@ func TestMap(t *testing.T) { require.False(t, has) } +func TestMap_Clear(t *testing.T) { + makeTest := func() (context.Context, Map[uint64, uint64]) { + sk, ctx := deps() + m := NewMap(NewSchemaBuilder(sk), NewPrefix(0), "test", Uint64Key, Uint64Value) + for i := uint64(0); i < clearBatchSize*2; i++ { + require.NoError(t, m.Set(ctx, i, i)) + } + return ctx, m + } + + t.Run("nil ranger", func(t *testing.T) { + ctx, m := makeTest() + err := m.Clear(ctx, nil) + require.NoError(t, err) + _, err = m.Iterate(ctx, nil) + require.ErrorIs(t, err, ErrInvalidIterator) + }) + + t.Run("custom ranger", func(t *testing.T) { + ctx, m := makeTest() + // delete from 0 to 100 + err := m.Clear(ctx, new(Range[uint64]).StartInclusive(0).EndInclusive(100)) + require.NoError(t, err) + + iter, err := m.Iterate(ctx, nil) + require.NoError(t, err) + keys, err := iter.Keys() + require.NoError(t, err) + require.Len(t, keys, clearBatchSize*2-101) + require.Equal(t, keys[0], uint64(101)) + require.Equal(t, keys[len(keys)-1], uint64(clearBatchSize*2-1)) + }) +} + func TestMap_IterateRaw(t *testing.T) { sk, ctx := deps() // safety check to ensure prefix boundaries are not crossed