github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/engine/common/synchronization/request_heap_test.go (about)

     1  package synchronization
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/stretchr/testify/require"
     7  
     8  	"github.com/onflow/flow-go/engine"
     9  	"github.com/onflow/flow-go/model/flow"
    10  	"github.com/onflow/flow-go/utils/unittest"
    11  )
    12  
    13  // TestRequestQueue_Get tests that after pushing multiple items we can drain the queue using Get method
    14  func TestRequestQueue_Get(t *testing.T) {
    15  	q := NewRequestHeap(100)
    16  	items := 20
    17  	messages := make(map[flow.Identifier]*engine.Message)
    18  	for i := 0; i < items; i++ {
    19  		msg := &engine.Message{
    20  			OriginID: unittest.IdentifierFixture(),
    21  			Payload:  unittest.IdentifierFixture(),
    22  		}
    23  		messages[msg.OriginID] = msg
    24  		require.True(t, q.Put(msg))
    25  	}
    26  
    27  	for i := 0; i < items; i++ {
    28  		msg, ok := q.Get()
    29  		require.True(t, ok)
    30  		expected, ok := messages[msg.OriginID]
    31  		require.True(t, ok)
    32  		require.Equal(t, expected, msg)
    33  	}
    34  
    35  	// at this point queue should be empty
    36  	_, ok := q.Get()
    37  	require.False(t, ok)
    38  }
    39  
    40  // TestRequestQueue_Put tests that putting an item into queue overwrites previous one
    41  func TestRequestQueue_Put(t *testing.T) {
    42  	q := NewRequestHeap(100)
    43  	msg := &engine.Message{
    44  		OriginID: unittest.IdentifierFixture(),
    45  		Payload:  unittest.IdentifierFixture(),
    46  	}
    47  	require.True(t, q.Put(msg))
    48  	require.Equal(t, msg, q.requests[msg.OriginID])
    49  
    50  	newMsg := &engine.Message{
    51  		OriginID: msg.OriginID,
    52  		Payload:  unittest.BlockFixture(),
    53  	}
    54  
    55  	// this should overwrite message
    56  	require.True(t, q.Put(newMsg))
    57  	require.Equal(t, newMsg, q.requests[msg.OriginID])
    58  }
    59  
    60  // TestRequestQueue_PutAtMaxCapacity tests that putting an item over max capacity results in successful eject and put
    61  func TestRequestQueue_PutAtMaxCapacity(t *testing.T) {
    62  	limit := uint(10)
    63  	q := NewRequestHeap(limit)
    64  
    65  	messages := make(map[flow.Identifier]*engine.Message)
    66  	var lastMsg *engine.Message // tracks the last message that we added to the heap
    67  	for i := uint(0); i < limit+1; i++ {
    68  		lastMsg = &engine.Message{
    69  			OriginID: unittest.IdentifierFixture(),
    70  			Payload:  unittest.IdentifierFixture(),
    71  		}
    72  		messages[lastMsg.OriginID] = lastMsg
    73  		require.True(t, q.Put(lastMsg))
    74  	}
    75  
    76  	// We have inserted 11 elements into the heap with capacity 10. The heap should now store
    77  	// 10 of these elements. By convention, the last-inserted element should be stored (no
    78  	// ejecting the just stored element).
    79  	lastMessagePopped := false
    80  	for k := uint(0); k < limit; k++ {
    81  		m, ok := q.Get()
    82  		require.True(t, ok)
    83  
    84  		// in the following code segment, we check that:
    85  		//  - only previously inserted elements are popped from the heap
    86  		//  - each element is only popped once
    87  		expectedMessage, found := messages[m.OriginID]
    88  		require.True(t, found)
    89  		require.Equal(t, m, expectedMessage)
    90  		delete(messages, m.OriginID)
    91  
    92  		// We expect that the most-recently inserted element is popped from the heap eventually,
    93  		// because this element has zero probability of being ejected itself
    94  		if m == lastMsg {
    95  			lastMessagePopped = true
    96  		}
    97  	}
    98  
    99  	// We have popped 10 elements from a heap with capacity 10, so the heap should now be empty:
   100  	m, ok := q.Get()
   101  	require.False(t, ok)
   102  	require.Nil(t, m)
   103  
   104  	// We have removed 10 out of 11 elements from `messages`. What remains in the map
   105  	// is that one message that the heap ejected. By convention, the heap should always eject first
   106  	// before accepting a new element. Therefore, the last-inserted element ('lastMsg') should have
   107  	// been popped from the heap in the for-loop, i.e. `lastMessagePopped` is expected to be true.
   108  	require.Equal(t, 1, len(messages))
   109  	require.True(t, lastMessagePopped)
   110  }