add RemoveByRange

Author: zhangchiqing

storage/operation/badgerimpl/writer.go

@@ -103,7 +103,8 @@ func (b *ReaderBatchWriter) Delete(key []byte) error {
 
 // DeleteByRange removes all keys with a prefix that falls within the
 // range [start, end], both inclusive.
-// No errors expected during normal operation
+// It returns error if endPrefix < startPrefix
+// no other errors are expected during normal operation
 func (b *ReaderBatchWriter) DeleteByRange(globalReader storage.Reader, startPrefix, endPrefix []byte) error {
 	err := operation.IterateKeysInPrefixRange(startPrefix, endPrefix, func(key []byte) error {
 		err := b.batch.Delete(key)

storage/operation/pebbleimpl/writer.go

@@ -1,6 +1,7 @@
 package pebbleimpl
 
 import (
+	"bytes"
 	"fmt"
 
 	"github.com/cockroachdb/pebble"
@@ -102,8 +103,13 @@ func (b *ReaderBatchWriter) Delete(key []byte) error {
 
 // DeleteByRange removes all keys with a prefix that falls within the
 // range [start, end], both inclusive.
-// No errors expected during normal operation
+// It returns error if endPrefix < startPrefix
+// no other errors are expected during normal operation
 func (b *ReaderBatchWriter) DeleteByRange(globalReader storage.Reader, startPrefix, endPrefix []byte) error {
+	if bytes.Compare(startPrefix, endPrefix) > 0 {
+		return fmt.Errorf("startPrefix key must be less than or equal to endPrefix key")
+	}
+
 	// DeleteRange takes the prefix range with start (inclusive) and end (exclusive, note: not inclusive).
 	// therefore, we need to increment the endPrefix to make it inclusive.
 	start, end, hasUpperBound := storage.StartEndPrefixToLowerUpperBound(startPrefix, endPrefix)

storage/operation/reads.go

@@ -17,11 +17,11 @@ import (
 // return (false, err) if running into any exception, the iteration should be stopped.
 type CheckFunc func(key []byte) (bool, error)
 
-// createFunc returns a pointer to an initialized entity that we can potentially
+// CreateFunc returns a pointer to an initialized entity that we can potentially
 // decode the next value into during a badger DB iteration.
 type CreateFunc func() interface{}
 
-// handleFunc is a function that starts the processing of the current key-value
+// HandleFunc is a function that starts the processing of the current key-value
 // pair during a badger iteration. It should be called after the key was checked
 // and the entity was decoded.
 // No errors are expected during normal operation. Any errors will halt the iteration.

storage/operation/writes.go

@@ -1,6 +1,9 @@
 package operation
 
 import (
+	"bytes"
+	"fmt"
+
 	"github.com/vmihailenco/msgpack"
 
 	"github.com/onflow/flow-go/module/irrecoverable"
@@ -43,13 +46,22 @@ func Remove(key []byte) func(storage.Writer) error {
 	}
 }
 
-// RemoveByPrefix removes all keys with the given prefix defined by [startPrefix, endPrefix] (both inclusive).
-// If no keys exist with the given prefix, this is a no-op.
+// RemoveByPrefix removes all keys with the given prefix
 // Error returns:
 // * generic error in case of unexpected database error
 func RemoveByPrefix(reader storage.Reader, key []byte) func(storage.Writer) error {
+	return RemoveByRange(reader, key, key)
+}
+
+// RemoveByRange removes all keys with a prefix that falls within the range [start, end], both inclusive.
+// It returns error if endPrefix < startPrefix
+// no other errors are expected during normal operation
+func RemoveByRange(reader storage.Reader, startPrefix []byte, endPrefix []byte) func(storage.Writer) error {
 	return func(w storage.Writer) error {
-		err := w.DeleteByRange(reader, key, key)
+		if bytes.Compare(startPrefix, endPrefix) > 0 {
+			return fmt.Errorf("startPrefix key must be less than or equal to endPrefix key")
+		}
+		err := w.DeleteByRange(reader, startPrefix, endPrefix)
 		if err != nil {
 			return irrecoverable.NewExceptionf("could not delete item: %w", err)
 		}

storage/operation/writes_test.go

@@ -192,7 +192,7 @@ func TestConcurrentRemove(t *testing.T) {
 	})
 }
 
-func TestRemoveRange(t *testing.T) {
+func TestRemoveByPrefix(t *testing.T) {
 	dbtest.RunWithStorages(t, func(t *testing.T, r storage.Reader, withWriter dbtest.WithWriter) {
 
 		// Define the prefix
@@ -244,6 +244,59 @@ func TestRemoveRange(t *testing.T) {
 	})
 }
 
+func TestRemoveByRange(t *testing.T) {
+	dbtest.RunWithStorages(t, func(t *testing.T, r storage.Reader, withWriter dbtest.WithWriter) {
+
+		startPrefix, endPrefix := []byte{0x10}, []byte{0x12}
+		// Create a range of keys around the boundaries of the prefix
+		keys := [][]byte{
+			{0x09, 0xff},
+			// within the range
+			{0x10, 0x00},
+			{0x10, 0x50},
+			{0x10, 0xff},
+			{0x11},
+			{0x12},
+			{0x12, 0x00},
+			{0x12, 0xff},
+			// after end -> not included in range
+			{0x13},
+			{0x1A, 0xff},
+		}
+
+		// Keys expected to be in the prefix range
+		includeStart, includeEnd := 1, 7
+
+		// Insert the keys into the storage
+		require.NoError(t, withWriter(func(writer storage.Writer) error {
+			for _, key := range keys {
+				value := []byte{0x00} // value are skipped, doesn't matter
+				err := operation.Upsert(key, value)(writer)
+				if err != nil {
+					return err
+				}
+			}
+			return nil
+		}))
+
+		// Remove the keys in the prefix range
+		require.NoError(t, withWriter(operation.RemoveByRange(r, startPrefix, endPrefix)))
+
+		// Verify that the keys in the prefix range have been removed
+		for i, key := range keys {
+			var exists bool
+			require.NoError(t, operation.Exists(key, &exists)(r))
+			t.Logf("key %x exists: %t", key, exists)
+
+			deleted := includeStart <= i && i <= includeEnd
+
+			// An item that was not deleted must exist
+			require.Equal(t, !deleted, exists,
+				"expected key %x to be %s", key, map[bool]string{true: "deleted", false: "not deleted"})
+		}
+	})
+}
+
 func TestRemoveFrom(t *testing.T) {
 	dbtest.RunWithStorages(t, func(t *testing.T, r storage.Reader, withWriter dbtest.WithWriter) {