diff --git a/watchtower/wtdb/range_index.go b/watchtower/wtdb/range_index.go new file mode 100644 index 000000000..77d09f645 --- /dev/null +++ b/watchtower/wtdb/range_index.go @@ -0,0 +1,619 @@ +package wtdb + +import ( + "fmt" + "sync" +) + +// rangeItem represents the start and end values of a range. +type rangeItem struct { + start uint64 + end uint64 +} + +// RangeIndexOption describes the signature of a functional option that can be +// used to modify the behaviour of a RangeIndex. +type RangeIndexOption func(*RangeIndex) + +// WithSerializeUint64Fn is a functional option that can be used to set the +// function to be used to do the serialization of a uint64 into a byte slice. +func WithSerializeUint64Fn(fn func(uint64) ([]byte, error)) RangeIndexOption { + return func(index *RangeIndex) { + index.serializeUint64 = fn + } +} + +// RangeIndex can be used to keep track of which numbers have been added to a +// set. It does so by keeping track of a sorted list of rangeItems. Each +// rangeItem has a start and end value of a range where all values in-between +// have been added to the set. It works well in situations where it is expected +// numbers in the set are not sparse. +type RangeIndex struct { + // set is a sorted list of rangeItem. + set []rangeItem + + // mu is used to ensure safe access to set. + mu sync.Mutex + + // serializeUint64 is the function that can be used to convert a uint64 + // to a byte slice. + serializeUint64 func(uint64) ([]byte, error) +} + +// NewRangeIndex constructs a new RangeIndex. An initial set of ranges may be +// passed to the function in the form of a map. +func NewRangeIndex(ranges map[uint64]uint64, + opts ...RangeIndexOption) (*RangeIndex, error) { + + index := &RangeIndex{ + serializeUint64: defaultSerializeUint64, + set: make([]rangeItem, 0), + } + + // Apply any functional options. + for _, o := range opts { + o(index) + } + + for s, e := range ranges { + if err := index.addRange(s, e); err != nil { + return nil, err + } + } + + return index, nil +} + +// addRange can be used to add an entire new range to the set. This method +// should only ever be called by NewRangeIndex to initialise the in-memory +// structure and so the RangeIndex mutex is not held during this method. +func (a *RangeIndex) addRange(start, end uint64) error { + // Check that the given range is valid. + if start > end { + return fmt.Errorf("invalid range. Start height %d is larger "+ + "than end height %d", start, end) + } + + // min is a helper closure that will return the minimum of two uint64s. + min := func(a, b uint64) uint64 { + if a < b { + return a + } + + return b + } + + // max is a helper closure that will return the maximum of two uint64s. + max := func(a, b uint64) uint64 { + if a > b { + return a + } + + return b + } + + // Collect the ranges that fall before and after the new range along + // with the start and end values of the new range. + var before, after []rangeItem + for _, x := range a.set { + // If the new start value can't extend the current ranges end + // value, then the two cannot be merged. The range is added to + // the group of ranges that fall before the new range. + if x.end+1 < start { + before = append(before, x) + continue + } + + // If the current ranges start value does not follow on directly + // from the new end value, then the two cannot be merged. The + // range is added to the group of ranges that fall after the new + // range. + if end+1 < x.start { + after = append(after, x) + continue + } + + // Otherwise, there is an overlap and so the two can be merged. + start = min(start, x.start) + end = max(end, x.end) + } + + // Re-construct the range index set. + a.set = append(append(before, rangeItem{ + start: start, + end: end, + }), after...) + + return nil +} + +// IsInIndex returns true if the given number is in the range set. +func (a *RangeIndex) IsInIndex(n uint64) bool { + a.mu.Lock() + defer a.mu.Unlock() + + _, isCovered := a.lowerBoundIndex(n) + + return isCovered +} + +// NumInSet returns the number of items covered by the range set. +func (a *RangeIndex) NumInSet() uint64 { + a.mu.Lock() + defer a.mu.Unlock() + + var numItems uint64 + for _, r := range a.set { + numItems += r.end - r.start + 1 + } + + return numItems +} + +// MaxHeight returns the highest number covered in the range. +func (a *RangeIndex) MaxHeight() uint64 { + a.mu.Lock() + defer a.mu.Unlock() + + if len(a.set) == 0 { + return 0 + } + + return a.set[len(a.set)-1].end +} + +// GetAllRanges returns a copy of the range set in the form of a map. +func (a *RangeIndex) GetAllRanges() map[uint64]uint64 { + a.mu.Lock() + defer a.mu.Unlock() + + cp := make(map[uint64]uint64, len(a.set)) + for _, item := range a.set { + cp[item.start] = item.end + } + + return cp +} + +// lowerBoundIndex returns the index of the RangeIndex that is most appropriate +// for the new value, n. In other words, it returns the index of the rangeItem +// set of the range where the start value is the highest start value in the set +// that is still lower than or equal to the given number, n. The returned +// boolean is true if the given number is already covered in the RangeIndex. +// A returned index of -1 indicates that no lower bound range exists in the set. +// Since the most likely case is that the new number will just extend the +// highest range, a check is first done to see if this is the case which will +// make the methods' computational complexity O(1). Otherwise, a binary search +// is done which brings the computational complexity to O(log N). +func (a *RangeIndex) lowerBoundIndex(n uint64) (int, bool) { + // If the set is empty, then there is no such index and the value + // definitely is not in the set. + if len(a.set) == 0 { + return -1, false + } + + // In most cases, the last index item will be the one we want. So just + // do a quick check on that index first to avoid doing the binary + // search. + lastIndex := len(a.set) - 1 + lastRange := a.set[lastIndex] + if lastRange.start <= n { + return lastIndex, lastRange.end >= n + } + + // Otherwise, do a binary search to find the index of interest. + var ( + low = 0 + high = len(a.set) - 1 + rangeIndex = -1 + ) + for { + mid := (low + high) / 2 //nolint: gomnd + currentRange := a.set[mid] + + switch { + case currentRange.start > n: + // If the start of the range is greater than n, we can + // completely cut out that entire part of the array. + high = mid + + case currentRange.start < n: + // If the range already includes the given height, we + // can stop searching now. + if currentRange.end >= n { + return mid, true + } + + // If the start of the range is smaller than n, we can + // store this as the new best index to return. + rangeIndex = mid + + // If low and mid are already equal, then increment low + // by 1. Exit if this means that low is now greater than + // high. + if low == mid { + low = mid + 1 + if low > high { + return rangeIndex, false + } + } else { + low = mid + } + + continue + + default: + // If the height is equal to the start value of the + // current range that mid is pointing to, then the + // height is already covered. + return mid, true + } + + // Exit if we have checked all the ranges. + if low == high { + break + } + } + + return rangeIndex, false +} + +// KVStore is an interface representing a key-value store. +type KVStore interface { + // Put saves the specified key/value pair to the store. Keys that do not + // already exist are added and keys that already exist are overwritten. + Put(key, value []byte) error + + // Delete removes the specified key from the bucket. Deleting a key that + // does not exist does not return an error. + Delete(key []byte) error +} + +// Add adds a single number to the range set. It first attempts to apply the +// necessary changes to the passed KV store and then only if this succeeds, will +// the changes be applied to the in-memory structure. +func (a *RangeIndex) Add(newHeight uint64, kv KVStore) error { + a.mu.Lock() + defer a.mu.Unlock() + + // Compute the changes that will need to be applied to both the sorted + // rangeItem array representation and the key-value store representation + // of the range index. + arrayChanges, kvStoreChanges := a.getChanges(newHeight) + + // First attempt to apply the KV store changes. Only if this succeeds + // will we apply the changes to our in-memory range index structure. + err := a.applyKVChanges(kv, kvStoreChanges) + if err != nil { + return err + } + + // Since the DB changes were successful, we can now commit the + // changes to our in-memory representation of the range set. + a.applyArrayChanges(arrayChanges) + + return nil +} + +// applyKVChanges applies the given set of kvChanges to a KV store. It is +// assumed that a transaction is being held on the kv store so that if any +// of the actions of the function fails, the changes will be reverted. +func (a *RangeIndex) applyKVChanges(kv KVStore, changes *kvChanges) error { + // Exit early if there are no changes to apply. + if kv == nil || changes == nil { + return nil + } + + // Check if any range pair needs to be deleted. + if changes.deleteKVKey != nil { + del, err := a.serializeUint64(*changes.deleteKVKey) + if err != nil { + return err + } + + if err := kv.Delete(del); err != nil { + return err + } + } + + start, err := a.serializeUint64(changes.key) + if err != nil { + return err + } + + end, err := a.serializeUint64(changes.value) + if err != nil { + return err + } + + return kv.Put(start, end) +} + +// applyArrayChanges applies the given arrayChanges to the in-memory RangeIndex +// itself. This should only be done once the persisted kv store changes have +// already been applied. +func (a *RangeIndex) applyArrayChanges(changes *arrayChanges) { + if changes == nil { + return + } + + if changes.indexToDelete != nil { + a.set = append( + a.set[:*changes.indexToDelete], + a.set[*changes.indexToDelete+1:]..., + ) + } + + if changes.newIndex != nil { + switch { + case *changes.newIndex == 0: + a.set = append([]rangeItem{{ + start: changes.start, + end: changes.end, + }}, a.set...) + + case *changes.newIndex == len(a.set): + a.set = append(a.set, rangeItem{ + start: changes.start, + end: changes.end, + }) + + default: + a.set = append( + a.set[:*changes.newIndex+1], + a.set[*changes.newIndex:]..., + ) + a.set[*changes.newIndex] = rangeItem{ + start: changes.start, + end: changes.end, + } + } + + return + } + + if changes.indexToEdit != nil { + a.set[*changes.indexToEdit] = rangeItem{ + start: changes.start, + end: changes.end, + } + } +} + +// arrayChanges encompasses the diff to apply to the sorted rangeItem array +// representation of a range index. Such a diff will either include adding a +// new range or editing an existing range. If an existing range is edited, then +// the diff might also include deleting an index (this will be the case if the +// editing of the one range results in the merge of another range). +type arrayChanges struct { + start uint64 + end uint64 + + // newIndex, if set, is the index of the in-memory range array where a + // new range, [start:end], should be added. newIndex should never be + // set at the same time as indexToEdit or indexToDelete. + newIndex *int + + // indexToDelete, if set, is the index of the sorted rangeItem array + // that should be deleted. This should be applied before reading the + // index value of indexToEdit. This should not be set at the same time + // as newIndex. + indexToDelete *int + + // indexToEdit is the index of the in-memory range array that should be + // edited. The range at this index will be changed to [start:end]. This + // should only be read after indexToDelete index has been deleted. + indexToEdit *int +} + +// kvChanges encompasses the diff to apply to a KV-store representation of a +// range index. A kv-store diff for the addition of a single number to the range +// index will include either a brand new key-value pair or the altering of the +// value of an existing key. Optionally, the diff may also include the deletion +// of an existing key. A deletion will be required if the addition of the new +// number results in the merge of two ranges. +type kvChanges struct { + key uint64 + value uint64 + + // deleteKVKey, if set, is the key of the kv store representation that + // should be deleted. + deleteKVKey *uint64 +} + +// getChanges will calculate and return the changes that need to be applied to +// both the sorted-rangeItem-array representation and the key-value store +// representation of the range index. +func (a *RangeIndex) getChanges(n uint64) (*arrayChanges, *kvChanges) { + // If the set is empty then a new range item is added. + if len(a.set) == 0 { + // For the array representation, a new range [n:n] is added to + // the first index of the array. + firstIndex := 0 + ac := &arrayChanges{ + newIndex: &firstIndex, + start: n, + end: n, + } + + // For the KV representation, a new [n:n] pair is added. + kvc := &kvChanges{ + key: n, + value: n, + } + + return ac, kvc + } + + // Find the index of the lower bound range to the new number. + indexOfRangeBelow, alreadyCovered := a.lowerBoundIndex(n) + + switch { + // The new number is already covered by the range index. No changes are + // required. + case alreadyCovered: + return nil, nil + + // No lower bound index exists. + case indexOfRangeBelow < 0: + // Check if the very first range can be merged into this new + // one. + if n+1 == a.set[0].start { + // If so, the two ranges can be merged and so the start + // value of the range is n and the end value is the end + // of the existing first range. + start := n + end := a.set[0].end + + // For the array representation, we can just edit the + // first entry of the array + editIndex := 0 + ac := &arrayChanges{ + indexToEdit: &editIndex, + start: start, + end: end, + } + + // For the KV store representation, we add a new kv pair + // and delete the range with the key equal to the start + // value of the range we are merging. + kvKeyToDelete := a.set[0].start + kvc := &kvChanges{ + key: start, + value: end, + deleteKVKey: &kvKeyToDelete, + } + + return ac, kvc + } + + // Otherwise, we add a new index. + + // For the array representation, a new range [n:n] is added to + // the first index of the array. + newIndex := 0 + ac := &arrayChanges{ + newIndex: &newIndex, + start: n, + end: n, + } + + // For the KV representation, a new [n:n] pair is added. + kvc := &kvChanges{ + key: n, + value: n, + } + + return ac, kvc + + // A lower range does exist, and it can be extended to include this new + // number. + case a.set[indexOfRangeBelow].end+1 == n: + start := a.set[indexOfRangeBelow].start + end := n + indexToChange := indexOfRangeBelow + + // If there are no intervals above this one or if there are, but + // they can't be merged into this one then we just need to edit + // this interval. + if indexOfRangeBelow == len(a.set)-1 || + a.set[indexOfRangeBelow+1].start != n+1 { + + // For the array representation, we just edit the index. + ac := &arrayChanges{ + indexToEdit: &indexToChange, + start: start, + end: end, + } + + // For the key-value representation, we just overwrite + // the end value at the existing start key. + kvc := &kvChanges{ + key: start, + value: end, + } + + return ac, kvc + } + + // There is a range above this one that we need to merge into + // this one. + delIndex := indexOfRangeBelow + 1 + end = a.set[delIndex].end + + // For the array representation, we delete the range above this + // one and edit this range to include the end value of the range + // above. + ac := &arrayChanges{ + indexToDelete: &delIndex, + indexToEdit: &indexToChange, + start: start, + end: end, + } + + // For the kv representation, we tweak the end value of an + // existing key and delete the key of the range we are deleting. + deleteKey := a.set[delIndex].start + kvc := &kvChanges{ + key: start, + value: end, + deleteKVKey: &deleteKey, + } + + return ac, kvc + + // A lower range does exist, but it can't be extended to include this + // new number, and so we need to add a new range after the lower bound + // range. + default: + newIndex := indexOfRangeBelow + 1 + + // If there are no ranges above this new one or if there are, + // but they can't be merged into this new one, then we can just + // add the new one as is. + if newIndex == len(a.set) || a.set[newIndex].start != n+1 { + ac := &arrayChanges{ + newIndex: &newIndex, + start: n, + end: n, + } + + kvc := &kvChanges{ + key: n, + value: n, + } + + return ac, kvc + } + + // Else, we merge the above index. + start := n + end := a.set[newIndex].end + toEdit := newIndex + + // For the array representation, we edit the range above to + // include the new start value. + ac := &arrayChanges{ + indexToEdit: &toEdit, + start: start, + end: end, + } + + // For the kv representation, we insert the new start-end key + // value pair and delete the key using the old start value. + delKey := a.set[newIndex].start + kvc := &kvChanges{ + key: start, + value: end, + deleteKVKey: &delKey, + } + + return ac, kvc + } +} + +func defaultSerializeUint64(i uint64) ([]byte, error) { + var b [8]byte + byteOrder.PutUint64(b[:], i) + return b[:], nil +} diff --git a/watchtower/wtdb/range_index_test.go b/watchtower/wtdb/range_index_test.go new file mode 100644 index 000000000..b03b9a410 --- /dev/null +++ b/watchtower/wtdb/range_index_test.go @@ -0,0 +1,349 @@ +package wtdb + +import ( + "fmt" + "math/rand" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +// TestRangeIndex tests that the RangeIndex works as expected. +func TestRangeIndex(t *testing.T) { + t.Parallel() + + assertRanges := func(index *RangeIndex, kvStore *mockKVStore, + v map[uint64]uint64) { + + require.EqualValues(t, v, index.GetAllRanges()) + require.EqualValues(t, v, kvStore.kv) + } + + t.Run("test zero value height", func(t *testing.T) { + t.Parallel() + + kvStore := newMockKVStore(nil) + index, err := NewRangeIndex(nil) + require.NoError(t, err) + + // Since zero values are tricky, assert that the empty index + // does not include zero. + require.False(t, index.IsInIndex(0)) + + // Now add zero to the index. + _ = index.Add(0, kvStore) + assertRanges(index, kvStore, map[uint64]uint64{0: 0}) + require.True(t, index.IsInIndex(0)) + }) + + t.Run("add duplicates", func(t *testing.T) { + t.Parallel() + + kvStore := newMockKVStore(nil) + index, err := NewRangeIndex(nil) + require.NoError(t, err) + + require.False(t, index.IsInIndex(1)) + + // Add 1 to the range. + err = index.Add(1, kvStore) + require.NoError(t, err) + + assertRanges(index, kvStore, map[uint64]uint64{1: 1}) + require.EqualValues(t, 1, index.MaxHeight()) + require.True(t, index.IsInIndex(1)) + + // Add 1 again and assert that nothing has changed. + err = index.Add(1, kvStore) + require.NoError(t, err) + + assertRanges(index, kvStore, map[uint64]uint64{1: 1}) + require.EqualValues(t, 1, index.MaxHeight()) + require.True(t, index.IsInIndex(1)) + }) + + t.Run("extend an existing range", func(t *testing.T) { + t.Parallel() + + kvStore := newMockKVStore(nil) + index, err := NewRangeIndex(nil) + require.NoError(t, err) + + assertRanges(index, kvStore, map[uint64]uint64{}) + + // Add 2. + _ = index.Add(2, kvStore) + assertRanges(index, kvStore, map[uint64]uint64{ + 2: 2, + }) + + // Add 3, 4 and 5 and assert that these just extend the existing + // range by incrementing its end value. + _ = index.Add(3, kvStore) + _ = index.Add(4, kvStore) + _ = index.Add(5, kvStore) + assertRanges(index, kvStore, map[uint64]uint64{ + 2: 5, + }) + + // Now add 1 and 0 and assert that these just extend the + // existing range by decrementing its start value. + _ = index.Add(1, kvStore) + _ = index.Add(0, kvStore) + assertRanges(index, kvStore, map[uint64]uint64{ + 0: 5, + }) + + // Assert various other properties of the current range. + require.True(t, index.IsInIndex(3)) + require.EqualValues(t, 5, index.MaxHeight()) + }) + + t.Run("add new ranges above and below", func(t *testing.T) { + t.Parallel() + + // Initialise the index with an initial range. + initialState := map[uint64]uint64{ + 4: 10, + } + + kvStore := newMockKVStore(initialState) + index, err := NewRangeIndex(initialState) + require.NoError(t, err) + + // Add 2 and 12. This should create two new ranges in the index. + _ = index.Add(12, kvStore) + _ = index.Add(2, kvStore) + assertRanges(index, kvStore, map[uint64]uint64{ + 2: 2, + 4: 10, + 12: 12, + }) + + // Assert various other properties of the current range. + require.EqualValues(t, 12, index.MaxHeight()) + require.False(t, index.IsInIndex(3)) + require.False(t, index.IsInIndex(11)) + require.True(t, index.IsInIndex(2)) + require.True(t, index.IsInIndex(5)) + require.True(t, index.IsInIndex(12)) + }) + + t.Run("merging two ranges", func(t *testing.T) { + t.Parallel() + + // Initialise the index with an initial set of ranges. + initialState := map[uint64]uint64{ + 2: 2, + 4: 10, + 12: 12, + } + + kvStore := newMockKVStore(initialState) + index, err := NewRangeIndex(initialState) + require.NoError(t, err) + + // Adding 3 should merge the first and second ranges. + _ = index.Add(3, kvStore) + assertRanges(index, kvStore, map[uint64]uint64{ + 2: 10, + 12: 12, + }) + + // Adding 11 should merge the first and second ranges. + _ = index.Add(11, kvStore) + assertRanges(index, kvStore, map[uint64]uint64{ + 2: 12, + }) + + // Assert various other properties of the current range. + require.EqualValues(t, 12, index.MaxHeight()) + require.True(t, index.IsInIndex(2)) + require.True(t, index.IsInIndex(5)) + require.True(t, index.IsInIndex(12)) + require.False(t, index.IsInIndex(1)) + }) + + t.Run("failure applying KV store updates", func(t *testing.T) { + t.Parallel() + + kvStore := newMockKVStore(nil) + index, err := NewRangeIndex(nil) + require.NoError(t, err) + + assertRanges(index, kvStore, map[uint64]uint64{}) + + // Ensure that the kv store will return an error when its + // methods are called. + kvStore.setError(fmt.Errorf("db error")) + + // Now attempt to add a new item to the range. + err = index.Add(20, kvStore) + require.Error(t, err) + + // Assert that the update failed for both the kv store and the + // array store. + assertRanges(index, kvStore, map[uint64]uint64{}) + + // Now let the kv store again not return an error. + kvStore.setError(nil) + + // Again attempt to add a new item to the range. + err = index.Add(20, kvStore) + require.NoError(t, err) + + // It should now succeed. + assertRanges(index, kvStore, map[uint64]uint64{ + 20: 20, + }) + }) + + t.Run("initialising with different ranges", func(t *testing.T) { + t.Parallel() + + tests := []struct { + name string + input map[uint64]uint64 + expectErr bool + expectedIndex map[uint64]uint64 + }{ + { + name: "invalid ranges", + input: map[uint64]uint64{ + 1: 5, + 6: 4, + }, + expectErr: true, + }, + { + name: "non-overlapping ranges", + input: map[uint64]uint64{ + 1: 2, + 4: 6, + 8: 20, + }, + expectedIndex: map[uint64]uint64{ + 1: 2, + 4: 6, + 8: 20, + }, + }, + { + name: "merge-able ranges", + input: map[uint64]uint64{ + 1: 2, + 3: 6, + 7: 20, + }, + expectedIndex: map[uint64]uint64{ + 1: 20, + }, + }, + { + name: "overlapping ranges", + input: map[uint64]uint64{ + 1: 4, + 3: 7, + 6: 20, + }, + expectedIndex: map[uint64]uint64{ + 1: 20, + }, + }, + } + + for _, test := range tests { + index, err := NewRangeIndex(test.input) + if test.expectErr { + require.Error(t, err) + continue + } + require.NoError(t, err) + + require.EqualValues( + t, test.expectedIndex, index.GetAllRanges(), + ) + } + }) + + t.Run("test large number of random inserts", func(t *testing.T) { + t.Parallel() + + size := 10000 + + // Construct an array with values from 0 to size. + arr := make([]uint64, size) + for i := 0; i < size; i++ { + arr[i] = uint64(i) + } + + // Shuffle the array so that the values are not added in order. + rand.Seed(time.Now().UnixNano()) + rand.Shuffle(len(arr), func(i, j int) { + arr[i], arr[j] = arr[j], arr[i] + }) + + kvStore := newMockKVStore(nil) + index, err := NewRangeIndex(nil) + require.NoError(t, err) + + // Now add each item in the array to the index. + for _, n := range arr { + _ = index.Add(n, kvStore) + } + + // Assert that in the end, there is only a single range. + assertRanges(index, kvStore, map[uint64]uint64{ + 0: uint64(size - 1), + }) + + require.EqualValues(t, uint64(size-1), index.MaxHeight()) + }) +} + +type mockKVStore struct { + kv map[uint64]uint64 + + err error +} + +func newMockKVStore(initialRanges map[uint64]uint64) *mockKVStore { + if initialRanges != nil { + return &mockKVStore{ + kv: initialRanges, + } + } + + return &mockKVStore{ + kv: make(map[uint64]uint64), + } +} + +func (m *mockKVStore) setError(err error) { + m.err = err +} + +func (m *mockKVStore) Put(key, value []byte) error { + if m.err != nil { + return m.err + } + + k := byteOrder.Uint64(key) + v := byteOrder.Uint64(value) + + m.kv[k] = v + + return nil +} + +func (m *mockKVStore) Delete(key []byte) error { + if m.err != nil { + return m.err + } + + k := byteOrder.Uint64(key) + delete(m.kv, k) + + return nil +}