github.com/nsqio/nsq@v1.3.0/nsqd/in_flight_pqueue_test.go (about)

     1  package nsqd
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"sort"
     7  	"testing"
     8  
     9  	"github.com/nsqio/nsq/internal/test"
    10  )
    11  
    12  func TestPriorityQueue(t *testing.T) {
    13  	c := 100
    14  	pq := newInFlightPqueue(c)
    15  
    16  	for i := 0; i < c+1; i++ {
    17  		pq.Push(&Message{clientID: int64(i), pri: int64(i)})
    18  	}
    19  	test.Equal(t, c+1, len(pq))
    20  	test.Equal(t, c*2, cap(pq))
    21  
    22  	for i := 0; i < c+1; i++ {
    23  		msg := pq.Pop()
    24  		test.Equal(t, int64(i), msg.clientID)
    25  	}
    26  	test.Equal(t, c/4, cap(pq))
    27  }
    28  
    29  func TestUnsortedInsert(t *testing.T) {
    30  	c := 100
    31  	pq := newInFlightPqueue(c)
    32  	ints := make([]int, 0, c)
    33  
    34  	for i := 0; i < c; i++ {
    35  		v := rand.Int()
    36  		ints = append(ints, v)
    37  		pq.Push(&Message{pri: int64(v)})
    38  	}
    39  	test.Equal(t, c, len(pq))
    40  	test.Equal(t, c, cap(pq))
    41  
    42  	sort.Ints(ints)
    43  
    44  	for i := 0; i < c; i++ {
    45  		msg, _ := pq.PeekAndShift(int64(ints[len(ints)-1]))
    46  		test.Equal(t, int64(ints[i]), msg.pri)
    47  	}
    48  }
    49  
    50  func TestRemove(t *testing.T) {
    51  	c := 100
    52  	pq := newInFlightPqueue(c)
    53  
    54  	msgs := make(map[MessageID]*Message)
    55  	for i := 0; i < c; i++ {
    56  		m := &Message{pri: int64(rand.Intn(100000000))}
    57  		copy(m.ID[:], fmt.Sprintf("%016d", m.pri))
    58  		msgs[m.ID] = m
    59  		pq.Push(m)
    60  	}
    61  
    62  	for i := 0; i < 10; i++ {
    63  		idx := rand.Intn((c - 1) - i)
    64  		var fm *Message
    65  		for _, m := range msgs {
    66  			if m.index == idx {
    67  				fm = m
    68  				break
    69  			}
    70  		}
    71  		rm := pq.Remove(idx)
    72  		test.Equal(t, fmt.Sprintf("%s", fm.ID), fmt.Sprintf("%s", rm.ID))
    73  	}
    74  
    75  	lastPriority := pq.Pop().pri
    76  	for i := 0; i < (c - 10 - 1); i++ {
    77  		msg := pq.Pop()
    78  		test.Equal(t, true, lastPriority <= msg.pri)
    79  		lastPriority = msg.pri
    80  	}
    81  }