nsq/nsqd/in_flight_pqueue_test.go

package nsqd

import (
	"fmt"
	"math/rand"
	"sort"
	"testing"

	"github.com/nsqio/nsq/internal/test"
)

func TestPriorityQueue(t *testing.T) {
	c := 100
	pq := newInFlightPqueue(c)

	for i := 0; i < c+1; i++ {
		pq.Push(&Message{clientID: int64(i), pri: int64(i)})
	}
	test.Equal(t, c+1, len(pq))
	test.Equal(t, c*2, cap(pq))

	for i := 0; i < c+1; i++ {
		msg := pq.Pop()
		test.Equal(t, int64(i), msg.clientID)
	}
	test.Equal(t, c/4, cap(pq))
}

func TestUnsortedInsert(t *testing.T) {
	c := 100
	pq := newInFlightPqueue(c)
	ints := make([]int, 0, c)

	for i := 0; i < c; i++ {
		v := rand.Int()
		ints = append(ints, v)
		pq.Push(&Message{pri: int64(v)})
	}
	test.Equal(t, c, len(pq))
	test.Equal(t, c, cap(pq))

	sort.Ints(ints)

	for i := 0; i < c; i++ {
		msg, _ := pq.PeekAndShift(int64(ints[len(ints)-1]))
		test.Equal(t, int64(ints[i]), msg.pri)
	}
}

func TestRemove(t *testing.T) {
	c := 100
	pq := newInFlightPqueue(c)

	msgs := make(map[MessageID]*Message)
	for i := 0; i < c; i++ {
		m := &Message{pri: int64(rand.Intn(100000000))}
		copy(m.ID[:], fmt.Sprintf("%016d", m.pri))
		msgs[m.ID] = m
		pq.Push(m)
	}

	for i := 0; i < 10; i++ {
		idx := rand.Intn((c - 1) - i)
		var fm *Message
		for _, m := range msgs {
			if m.index == idx {
				fm = m
				break
			}
		}
		rm := pq.Remove(idx)
		test.Equal(t, fmt.Sprintf("%s", fm.ID), fmt.Sprintf("%s", rm.ID))
	}

	lastPriority := pq.Pop().pri
	for i := 0; i < (c - 10 - 1); i++ {
		msg := pq.Pop()
		test.Equal(t, true, lastPriority <= msg.pri)
		lastPriority = msg.pri
	}
}