loki/pkg/queue/treequeue_test.go

package queue

import (
	"testing"

	"github.com/stretchr/testify/require"
)

type dummyRequest struct {
	id int
}

func r(id int) *dummyRequest {
	return &dummyRequest{id}
}

func TestTreeQueue(t *testing.T) {

	t.Run("add sub queues recursively", func(t *testing.T) {
		pathA := QueuePath([]string{"l0", "l1", "l3"})
		pathB := QueuePath([]string{"l0", "l2", "l3"})

		q := newTreeQueue(1, "root")
		require.NotNil(t, q)
		require.Equal(t, "root", q.Name())
		require.Equal(t, 0, q.Len())
		require.Equal(t, 0, q.mapping.Len())

		q.add(pathA)
		require.Equal(t, 1, q.mapping.Len())

		q.add(pathB)
		require.Equal(t, 1, q.mapping.Len())
	})

	t.Run("enqueue/dequeue to/from subqueues", func(t *testing.T) {
		/**
		root: [0]
		  a: [1]
			b: [2]
			  b0: [20]
				b1: [21]
			c: [3]
			  c0: [30]
				  c00: [300]
					c01: [301]
				c1: [31]
				  c10: [310]
					c11: [311]
		**/
		paths := []QueuePath{
			QueuePath([]string{"a"}),
			QueuePath([]string{"b", "b0"}),
			QueuePath([]string{"b", "b1"}),
			QueuePath([]string{"c", "c0", "c00"}),
			QueuePath([]string{"c", "c0", "c01"}),
			QueuePath([]string{"c", "c1", "c10"}),
			QueuePath([]string{"c", "c1", "c11"}),
		}

		q := newTreeQueue(10, "root")
		require.NotNil(t, q)
		for _, p := range paths {
			q.add(p)
		}

		require.Equal(t, 3, q.mapping.Len())

		// no items in any queues
		require.Equal(t, 0, q.Len())

		q.Chan() <- r(0)
		require.Equal(t, 1, q.Len())

		q.mapping.GetByKey("a").Chan() <- r(1)
		require.Equal(t, 2, q.Len())

		q.mapping.GetByKey("b").Chan() <- r(2)
		q.mapping.GetByKey("b").mapping.GetByKey("b0").Chan() <- r(20)
		q.mapping.GetByKey("b").mapping.GetByKey("b1").Chan() <- r(21)
		require.Equal(t, 5, q.Len())

		q.mapping.GetByKey("c").Chan() <- r(3)
		q.mapping.GetByKey("c").mapping.GetByKey("c0").Chan() <- r(30)
		q.mapping.GetByKey("c").mapping.GetByKey("c0").mapping.GetByKey("c00").Chan() <- r(300)
		q.mapping.GetByKey("c").mapping.GetByKey("c0").mapping.GetByKey("c01").Chan() <- r(301)
		q.mapping.GetByKey("c").mapping.GetByKey("c1").Chan() <- r(31)
		q.mapping.GetByKey("c").mapping.GetByKey("c1").mapping.GetByKey("c10").Chan() <- r(310)
		q.mapping.GetByKey("c").mapping.GetByKey("c1").mapping.GetByKey("c11").Chan() <- r(311)
		require.Equal(t, 12, q.Len())
		t.Log(q)

		items := make([]int, 0, q.Len())

		for q.Len() > 0 {
			r := q.Dequeue()
			if r == nil {
				continue
			}
			items = append(items, r.(*dummyRequest).id)
		}
		require.Len(t, items, 12)
		require.Equal(t, []int{0, 1, 2, 3, 20, 30, 21, 31, 300, 310, 301, 311}, items)
	})

	t.Run("dequeue ensure round-robin", func(t *testing.T) {
		/**
		root:
		  a: [100, 101, 102]
			b: [200]
			c: [300, 301]
		**/
		paths := []QueuePath{
			QueuePath([]string{"a"}),
			QueuePath([]string{"b"}),
			QueuePath([]string{"c"}),
		}

		q := newTreeQueue(10, "root")
		require.NotNil(t, q)
		for _, p := range paths {
			q.add(p)
		}

		require.Equal(t, 3, q.mapping.Len())

		// no items in any queues
		require.Equal(t, 0, q.Len())

		q.mapping.GetByKey("a").Chan() <- r(100)
		q.mapping.GetByKey("a").Chan() <- r(101)
		q.mapping.GetByKey("a").Chan() <- r(102)
		q.mapping.GetByKey("b").Chan() <- r(200)
		q.mapping.GetByKey("c").Chan() <- r(300)
		q.mapping.GetByKey("c").Chan() <- r(301)

		t.Log(q)

		items := make([]int, 0, q.Len())

		for q.Len() > 0 {
			r := q.Dequeue()
			if r == nil {
				continue
			}
			items = append(items, r.(*dummyRequest).id)
		}
		require.Len(t, items, 6)
		require.Equal(t, []int{100, 200, 300, 101, 301, 102}, items)
	})

	t.Run("empty sub-queues are removed", func(t *testing.T) {
		q := newTreeQueue(10, "root")
		q.add(QueuePath{"a"})
		q.add(QueuePath{"b"})

		q.mapping.GetByKey("a").Chan() <- r(1)
		q.mapping.GetByKey("b").Chan() <- r(2)

		t.Log(q)

		// drain queue
		r := q.Dequeue()
		for r != nil {
			r = q.Dequeue()
		}

		require.Nil(t, q.mapping.GetByKey("a"))
		require.Nil(t, q.mapping.GetByKey("b"))
	})
}