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

     1  package queue_test
     2  
     3  import (
     4  	"context"
     5  	"math/rand"
     6  	"sync"
     7  	"sync/atomic"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/stretchr/testify/assert"
    12  
    13  	"github.com/onflow/flow-go/module/metrics"
    14  	"github.com/onflow/flow-go/network"
    15  	"github.com/onflow/flow-go/network/queue"
    16  )
    17  
    18  // TestSingleQueueWorkers tests that a single worker can successfully read all elements from the queue
    19  func TestSingleQueueWorker(t *testing.T) {
    20  	testWorkers(t, 10, 100, 1)
    21  }
    22  
    23  // TestMultipleQueueWorkers tests that multiple workers can successfully read all elements from the queue
    24  func TestMultipleQueueWorkers(t *testing.T) {
    25  	testWorkers(t, 10, 100, rand.Intn(9)+2)
    26  
    27  }
    28  
    29  // testWorkers tests that with the given max priority, message count and worker count, a queue can be successfully read.
    30  // workerCnt should not be more than maxPriority for this test
    31  func testWorkers(t *testing.T, maxPriority int, messageCnt int, workerCnt int) {
    32  
    33  	assert.LessOrEqual(t, workerCnt, maxPriority)
    34  
    35  	ctx, cancel := context.WithCancel(context.Background())
    36  	defer cancel()
    37  	// the priority function just returns the message as the priority itself (message = priority)
    38  	var q network.MessageQueue = queue.NewMessageQueue(ctx, func(m interface{}) (queue.Priority, error) {
    39  		i, ok := m.(int)
    40  		assert.True(t, ok)
    41  		return queue.Priority(i), nil
    42  	},
    43  		metrics.NewNoopCollector())
    44  
    45  	var l sync.Mutex                                // protect comparisons with expectedPriority
    46  	messagesPerPriority := messageCnt / maxPriority // messages per priority
    47  	expectedPriority := maxPriority - 1             // when dequeing, the priority can be the current highest priority or one less
    48  	var callbackCnt int64                           //count the number of times the callback gets called
    49  	// callback checks if message is of expected priority
    50  	callback := func(data interface{}) {
    51  		actual := data.(int)
    52  		l.Lock()
    53  		assert.LessOrEqual(t, expectedPriority, actual)
    54  		atomic.AddInt64(&callbackCnt, 1)
    55  		if callbackCnt%int64(messagesPerPriority) == 0 {
    56  			expectedPriority--
    57  		}
    58  		l.Unlock()
    59  	}
    60  
    61  	// the queue is populated with messageCnt number of messages
    62  	// each message is an int which is also its priority
    63  	// messages are inserted in increasing order of priority
    64  	// e.g. 1,2,3...10,1,2,3,..10,....messagecnt
    65  	for i := 0; i < messageCnt; i++ {
    66  		priority := (i % maxPriority) + 1
    67  		err := q.Insert(priority)
    68  		assert.NoError(t, err)
    69  	}
    70  
    71  	// create all the workers
    72  	queue.CreateQueueWorkers(ctx, uint64(workerCnt), q, callback)
    73  
    74  	// check that callback was eventually called expected number of times
    75  	assert.Eventually(t, func() bool {
    76  		actualCnt := atomic.LoadInt64(&callbackCnt)
    77  		return actualCnt == int64(messageCnt)
    78  	}, time.Second, 5*time.Millisecond)
    79  }