github.com/onflow/flow-go@v0.35.7-crescendo-preview.23-atree-inlining/module/mempool/queue/heroStore_test.go (about)

     1  package queue_test
     2  
     3  import (
     4  	"sync"
     5  	"testing"
     6  	"time"
     7  
     8  	"github.com/stretchr/testify/require"
     9  
    10  	"github.com/onflow/flow-go/engine"
    11  	"github.com/onflow/flow-go/module/mempool/queue"
    12  	"github.com/onflow/flow-go/module/metrics"
    13  	"github.com/onflow/flow-go/utils/unittest"
    14  )
    15  
    16  func TestHeroStore_Sequential(t *testing.T) {
    17  	sizeLimit := 100
    18  
    19  	// sequentially put and get
    20  	store := queue.NewHeroStore(uint32(sizeLimit), unittest.Logger(), metrics.NewNoopCollector())
    21  
    22  	messages := unittest.EngineMessageFixtures(sizeLimit)
    23  	for i := 0; i < sizeLimit; i++ {
    24  		require.True(t, store.Put(messages[i]))
    25  
    26  		// duplicate put should fail
    27  		require.False(t, store.Put(messages[i]))
    28  	}
    29  
    30  	// once store meets the size limit, any extra put should fail.
    31  	for i := 0; i < 100; i++ {
    32  		require.False(t, store.Put(unittest.EngineMessageFixture()))
    33  	}
    34  
    35  	// getting entities sequentially.
    36  	for i := 0; i < sizeLimit; i++ {
    37  		msg, ok := store.Get()
    38  		require.True(t, ok)
    39  		require.Equal(t, messages[i], msg)
    40  	}
    41  }
    42  
    43  // TestHeroStore_Concurrent evaluates correctness of store implementation against concurrent put and get.
    44  func TestHeroStore_Concurrent(t *testing.T) {
    45  	sizeLimit := 100
    46  	store := queue.NewHeroStore(uint32(sizeLimit), unittest.Logger(), metrics.NewNoopCollector())
    47  
    48  	// initially there should be nothing to pop
    49  	msg, ok := store.Get()
    50  	require.False(t, ok)
    51  	require.Nil(t, msg)
    52  
    53  	putWG := &sync.WaitGroup{}
    54  	putWG.Add(sizeLimit)
    55  
    56  	messages := unittest.EngineMessageFixtures(sizeLimit)
    57  	// putting messages concurrently.
    58  	for _, m := range messages {
    59  		m := m
    60  		go func() {
    61  			require.True(t, store.Put(m))
    62  			putWG.Done()
    63  		}()
    64  	}
    65  	unittest.RequireReturnsBefore(t, putWG.Wait, 100*time.Millisecond, "could not put all messages on time")
    66  
    67  	// once store meets the size limit, any extra put should fail.
    68  	putWG.Add(sizeLimit)
    69  	for i := 0; i < sizeLimit; i++ {
    70  		go func() {
    71  			require.False(t, store.Put(unittest.EngineMessageFixture()))
    72  			putWG.Done()
    73  		}()
    74  	}
    75  	unittest.RequireReturnsBefore(t, putWG.Wait, 100*time.Millisecond, "could not put all messages on time")
    76  
    77  	getWG := &sync.WaitGroup{}
    78  	getWG.Add(sizeLimit)
    79  	matchLock := &sync.Mutex{}
    80  
    81  	// pop-ing entities concurrently.
    82  	for i := 0; i < sizeLimit; i++ {
    83  		go func() {
    84  			msg, ok := store.Get()
    85  			require.True(t, ok)
    86  
    87  			matchLock.Lock()
    88  			matchAndRemoveMessages(t, messages, msg)
    89  			matchLock.Unlock()
    90  
    91  			getWG.Done()
    92  		}()
    93  	}
    94  	unittest.RequireReturnsBefore(t, getWG.Wait, 100*time.Millisecond, "could not get all messages on time")
    95  
    96  	// store must be empty after getting all
    97  	msg, ok = store.Get()
    98  	require.False(t, ok)
    99  	require.Nil(t, msg)
   100  }
   101  
   102  // matchAndRemove checks existence of the msg in the "messages" array, and if a match is found, it is removed.
   103  // If no match is found for a message, it fails the test.
   104  func matchAndRemoveMessages(t *testing.T, messages []*engine.Message, msg *engine.Message) []*engine.Message {
   105  	for i, m := range messages {
   106  		if m.OriginID == msg.OriginID {
   107  			require.Equal(t, m, msg)
   108  			// removes the matched message from the list
   109  			messages = append(messages[:i], messages[i+1:]...)
   110  			return messages
   111  		}
   112  	}
   113  
   114  	// no message found in the list to match
   115  	require.Fail(t, "could not find a match for the message")
   116  	return nil
   117  }