Skip to content

Commit

Permalink
More tests, fix uncovered bugs.
Browse files Browse the repository at this point in the history
Signed-off-by: Cody Littley <cody@eigenlabs.org>
  • Loading branch information
cody-littley committed Sep 3, 2024
1 parent 8697492 commit e23c352
Show file tree
Hide file tree
Showing 3 changed files with 81 additions and 36 deletions.
12 changes: 9 additions & 3 deletions lightnode/chunkgroup/assignment_queue.go
Original file line number Diff line number Diff line change
Expand Up @@ -53,7 +53,7 @@ type assignmentQueue struct {

// The number of elements in the queue. Tracked separately since the heap and NodeIdSet
// may contain removed nodes that have not yet been fully garbage collected.
size uint
size uint32
}

// newAssignmentQueue creates a new priority queue.
Expand All @@ -67,7 +67,7 @@ func newAssignmentQueue() *assignmentQueue {
}

// Size returns the number of elements in the priority queue.
func (queue *assignmentQueue) Size() uint {
func (queue *assignmentQueue) Size() uint32 {
return queue.size
}

Expand Down Expand Up @@ -124,6 +124,7 @@ func (queue *assignmentQueue) Remove(key assignmentKey) {
queue.size--

queue.assignmentSet[key] = false
queue.collectGarbage()
}

// collectGarbage removes all nodes that have been removed from the queue but have not yet been fully deleted.
Expand All @@ -134,14 +135,18 @@ func (queue *assignmentQueue) collectGarbage() {
}

// sanity check to prevent infinite loops
maxIterations := len(queue.heap.data)
maxIterations := len(queue.heap.data) + 1

for {
maxIterations--
if maxIterations < 0 {
panic("garbage collection did not terminate")
}

if len(queue.heap.data) == 0 {
return
}

next := queue.heap.data[0]

notRemoved, present := queue.assignmentSet[next.key]
Expand All @@ -155,5 +160,6 @@ func (queue *assignmentQueue) collectGarbage() {
}

heap.Pop(queue.heap)
delete(queue.assignmentSet, next.key)
}
}
62 changes: 31 additions & 31 deletions lightnode/chunkgroup/assignment_queue_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -31,17 +31,17 @@ func randomAssignment(nextShuffleTime time.Time) *chunkGroupAssignment {

func TestEmptyQueue(t *testing.T) {
queue := newAssignmentQueue()
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())
assert.Nil(t, queue.Pop())
assert.Nil(t, queue.Peek())
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())
}

func TestInOrderInsertion(t *testing.T) {
tu.InitializeRandom()

queue := newAssignmentQueue()
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())

startTime := tu.RandomTime()
numberOfElements := uint(100)
Expand All @@ -59,33 +59,33 @@ func TestInOrderInsertion(t *testing.T) {
queue.Push(element)
}

assert.Equal(t, i+1, queue.Size())
assert.Equal(t, uint32(i+1), queue.Size())
}

// Pop elements in order.
for i := uint(0); i < numberOfElements; i++ {
preview := queue.Peek()
assert.Equal(t, expectedOrder[i], preview)
assert.Equal(t, numberOfElements-i, queue.Size())
assert.Equal(t, uint32(numberOfElements-i), queue.Size())

next := queue.Pop()
assert.Equal(t, expectedOrder[i], next)
assert.Equal(t, numberOfElements-i-1, queue.Size())
assert.Equal(t, uint32(numberOfElements-i-1), queue.Size())
}
}

func TestReverseOrderInsertion(t *testing.T) {
tu.InitializeRandom()

queue := newAssignmentQueue()
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())

startTime := tu.RandomTime()
numberOfElements := uint(100)
numberOfElements := uint32(100)
expectedOrder := make([]*chunkGroupAssignment, 0, numberOfElements)

// Generate the elements that will eventually be inserted.
for i := uint(0); i < numberOfElements; i++ {
for i := uint32(0); i < numberOfElements; i++ {
element := randomAssignment(startTime.Add(time.Second * time.Duration(i)))
expectedOrder = append(expectedOrder, element)
}
Expand All @@ -100,11 +100,11 @@ func TestReverseOrderInsertion(t *testing.T) {
queue.Push(expectedOrder[i])
}

assert.Equal(t, numberOfElements-uint(i), queue.Size())
assert.Equal(t, numberOfElements-uint32(i), queue.Size())
}

// Pop elements in order.
for i := uint(0); i < numberOfElements; i++ {
for i := uint32(0); i < numberOfElements; i++ {
preview := queue.Peek()
assert.Equal(t, expectedOrder[i], preview)
assert.Equal(t, numberOfElements-i, queue.Size())
Expand All @@ -119,7 +119,7 @@ func TestRandomInsertion(t *testing.T) {
tu.InitializeRandom()

queue := newAssignmentQueue()
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())

startTime := tu.RandomTime()
numberOfElements := uint(100)
Expand All @@ -143,26 +143,26 @@ func TestRandomInsertion(t *testing.T) {
queue.Push(expectedOrder[perm[i]])
}

assert.Equal(t, i+1, queue.Size())
assert.Equal(t, uint32(i+1), queue.Size())
}

// Pop elements in order.
for i := uint(0); i < numberOfElements; i++ {
preview := queue.Peek()
assert.Equal(t, expectedOrder[i], preview)
assert.Equal(t, numberOfElements-i, queue.Size())
assert.Equal(t, uint32(numberOfElements-i), queue.Size())

next := queue.Pop()
assert.Equal(t, expectedOrder[i], next)
assert.Equal(t, numberOfElements-i-1, queue.Size())
assert.Equal(t, uint32(numberOfElements-i-1), queue.Size())
}
}

func TestPeriodicRemoval(t *testing.T) {
tu.InitializeRandom()

queue := newAssignmentQueue()
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())

startTime := tu.RandomTime()
numberOfElements := uint(100)
Expand Down Expand Up @@ -192,7 +192,7 @@ func TestPeriodicRemoval(t *testing.T) {
queue.Push(expectedOrder[perm[i]])
}

assert.Equal(t, i+1, queue.Size())
assert.Equal(t, uint32(i+1), queue.Size())
}

removalCount := uint(0)
Expand All @@ -209,27 +209,27 @@ func TestPeriodicRemoval(t *testing.T) {
}

removalCount++
assert.Equal(t, numberOfElements-removalCount, queue.Size())
assert.Equal(t, uint32(numberOfElements-removalCount), queue.Size())
}
}

// Pop elements in order. We shouldn't see the removed elements.
for i := uint(0); i < (numberOfElements - removalCount); i++ {
preview := queue.Peek()
assert.Equal(t, expectedOrderWithRemovals[i], preview)
assert.Equal(t, numberOfElements-i-removalCount, queue.Size())
assert.Equal(t, uint32(numberOfElements-i-removalCount), queue.Size())

next := queue.Pop()
assert.Equal(t, expectedOrderWithRemovals[i], next)
assert.Equal(t, numberOfElements-i-removalCount-1, queue.Size())
assert.Equal(t, uint32(numberOfElements-i-removalCount-1), queue.Size())
}
}

func TestContiguousRemoval(t *testing.T) {
tu.InitializeRandom()

queue := newAssignmentQueue()
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())

startTime := tu.RandomTime()
numberOfElements := uint(100)
Expand Down Expand Up @@ -259,7 +259,7 @@ func TestContiguousRemoval(t *testing.T) {
queue.Push(expectedOrder[perm[i]])
}

assert.Equal(t, i+1, queue.Size())
assert.Equal(t, uint32(i+1), queue.Size())
}

removalCount := uint(0)
Expand All @@ -276,27 +276,27 @@ func TestContiguousRemoval(t *testing.T) {
}

removalCount++
assert.Equal(t, numberOfElements-removalCount, queue.Size())
assert.Equal(t, uint32(numberOfElements-removalCount), queue.Size())
}
}

// Pop elements in order. We shouldn't see the removed elements.
for i := uint(0); i < (numberOfElements - removalCount); i++ {
preview := queue.Peek()
assert.Equal(t, expectedOrderWithRemovals[i], preview)
assert.Equal(t, numberOfElements-i-removalCount, queue.Size())
assert.Equal(t, uint32(numberOfElements-i-removalCount), queue.Size())

next := queue.Pop()
assert.Equal(t, expectedOrderWithRemovals[i], next)
assert.Equal(t, numberOfElements-i-removalCount-1, queue.Size())
assert.Equal(t, uint32(numberOfElements-i-removalCount-1), queue.Size())
}
}

func TestRemoveFollowedByPush(t *testing.T) {
tu.InitializeRandom()

queue := newAssignmentQueue()
assert.Equal(t, uint(0), queue.Size())
assert.Equal(t, uint32(0), queue.Size())

startTime := tu.RandomTime()
numberOfElements := uint(100)
Expand All @@ -320,7 +320,7 @@ func TestRemoveFollowedByPush(t *testing.T) {
queue.Push(expectedOrder[perm[i]])
}

assert.Equal(t, i+1, queue.Size())
assert.Equal(t, uint32(i+1), queue.Size())
}

removalCount := uint(0)
Expand All @@ -337,7 +337,7 @@ func TestRemoveFollowedByPush(t *testing.T) {
}

removalCount++
assert.Equal(t, numberOfElements-removalCount, queue.Size())
assert.Equal(t, uint32(numberOfElements-removalCount), queue.Size())
}
}

Expand All @@ -346,18 +346,18 @@ func TestRemoveFollowedByPush(t *testing.T) {
if i%7 == 0 {
queue.Push(expectedOrder[i])
removalCount--
assert.Equal(t, numberOfElements-removalCount, queue.Size())
assert.Equal(t, uint32(numberOfElements-removalCount), queue.Size())
}
}

// Pop elements in order.
for i := uint(0); i < (numberOfElements - removalCount); i++ {
preview := queue.Peek()
assert.Equal(t, expectedOrder[i], preview)
assert.Equal(t, numberOfElements-i, queue.Size())
assert.Equal(t, uint32(numberOfElements-i), queue.Size())

next := queue.Pop()
assert.Equal(t, expectedOrder[i], next)
assert.Equal(t, numberOfElements-i-1, queue.Size())
assert.Equal(t, uint32(numberOfElements-i-1), queue.Size())
}
}
43 changes: 41 additions & 2 deletions lightnode/chunkgroup/map_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,8 +17,31 @@ func randomRegistration(startTime time.Time, lastRegistration time.Time) *lightn
tu.RandomTimeInRange(startTime, lastRegistration))
}

// TODO test that deletes things and ensures there is no garbage left in maps
// TODO take into account when a light node was registered
// Delete everything from the map and verify that the internal data structures have been cleaned up.
func verifyMapDeletion(
t *testing.T,
now time.Time,
m *Map,
expectedRegistrations *map[uint64]*lightnode.Registration) {

// Delete remaining elements.
for id := range *expectedRegistrations {
m.Remove(id)
}

assert.Equal(t, uint32(0), m.Size())
assert.Equal(t, 0, len(m.assignmentMap))
assert.Equal(t, 0, len(m.lightNodes))

assert.Equal(t, m.chunkGroupCount, uint32(len(m.chunkGroups)))
for chunkIndex := uint32(0); chunkIndex < m.chunkGroupCount; chunkIndex++ {
assert.Equal(t, 0, len(m.chunkGroups[chunkIndex]))
}

assert.Equal(t, uint32(0), m.shuffleQueue.Size())
assert.Equal(t, 0, len(m.shuffleQueue.assignmentSet))
assert.Equal(t, 0, len(m.shuffleQueue.heap.data))
}

func TestAddRemoveGetOneAssignment(t *testing.T) {
tu.InitializeRandom()
Expand Down Expand Up @@ -78,6 +101,8 @@ func TestAddRemoveGetOneAssignment(t *testing.T) {
assert.Equal(t, registration, cgMap.Get(id))
}
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

func TestAddRemoveGetMultipleAssignments(t *testing.T) {
Expand Down Expand Up @@ -138,6 +163,8 @@ func TestAddRemoveGetMultipleAssignments(t *testing.T) {
assert.Equal(t, registration, cgMap.Get(id))
}
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

func TestChunkGroupCalculationsSingleAssignment(t *testing.T) {
Expand Down Expand Up @@ -230,6 +257,8 @@ func TestChunkGroupCalculationsSingleAssignment(t *testing.T) {

assert.Equal(t, count, nodesReported)
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

func TestChunkGroupCalculationsMultipleAssignments(t *testing.T) {
Expand Down Expand Up @@ -348,6 +377,8 @@ func TestChunkGroupCalculationsMultipleAssignments(t *testing.T) {
assert.Equal(t, expectedNodes, nodes)
}
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

func TestGetRandomNodeSingleAssignment(t *testing.T) {
Expand Down Expand Up @@ -425,6 +456,8 @@ func TestGetRandomNodeSingleAssignment(t *testing.T) {
}
}
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

func TestGetRandomNodeMultipleAssignments(t *testing.T) {
Expand Down Expand Up @@ -511,6 +544,8 @@ func TestGetRandomNodeMultipleAssignments(t *testing.T) {
}
}
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

func TestSingleChunkGroupSingleAssignment(t *testing.T) {
Expand Down Expand Up @@ -602,6 +637,8 @@ func TestSingleChunkGroupSingleAssignment(t *testing.T) {

assert.Equal(t, count, nodesReported)
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

func TestSingleChunkGroupMultipleAssignments(t *testing.T) {
Expand Down Expand Up @@ -693,4 +730,6 @@ func TestSingleChunkGroupMultipleAssignments(t *testing.T) {

assert.Equal(t, count, nodesReported)
}

verifyMapDeletion(t, startTime, &cgMap, &expectedRegistrations)
}

0 comments on commit e23c352

Please sign in to comment.