github.com/grafana/pyroscope@v1.18.0/pkg/metastore/compaction/scheduler/scheduler_queue_test.go (about)

     1  package scheduler
     2  
     3  import (
     4  	"container/heap"
     5  	"math/rand"
     6  	"testing"
     7  
     8  	"github.com/stretchr/testify/assert"
     9  
    10  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    11  	"github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1/raft_log"
    12  )
    13  
    14  func TestJobQueue_order(t *testing.T) {
    15  	items := []*raft_log.CompactionJobState{
    16  		{
    17  			Name:            "job-6",
    18  			CompactionLevel: 0,
    19  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_UNSPECIFIED,
    20  			LeaseExpiresAt:  5,
    21  			Failures:        0,
    22  		},
    23  		{
    24  			Name:            "job-0",
    25  			CompactionLevel: 1,
    26  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_UNSPECIFIED,
    27  			LeaseExpiresAt:  2,
    28  			Failures:        0,
    29  		},
    30  		{
    31  			Name:            "job-1",
    32  			CompactionLevel: 1,
    33  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_UNSPECIFIED,
    34  			LeaseExpiresAt:  2,
    35  			Failures:        0,
    36  		},
    37  		{
    38  			Name:            "job-2",
    39  			CompactionLevel: 0,
    40  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_IN_PROGRESS,
    41  			LeaseExpiresAt:  2,
    42  			Failures:        0,
    43  		},
    44  		{
    45  			Name:            "job-5",
    46  			CompactionLevel: 1,
    47  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_UNSPECIFIED,
    48  			LeaseExpiresAt:  3,
    49  			Failures:        0,
    50  		},
    51  		{
    52  			Name:            "job-3",
    53  			CompactionLevel: 0,
    54  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_IN_PROGRESS,
    55  			LeaseExpiresAt:  2,
    56  			Failures:        5,
    57  		},
    58  		{
    59  			Name:            "job-4",
    60  			CompactionLevel: 0,
    61  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_IN_PROGRESS,
    62  			LeaseExpiresAt:  1,
    63  			Failures:        5,
    64  		},
    65  	}
    66  
    67  	test := func(items []*raft_log.CompactionJobState) {
    68  		q := newJobQueue()
    69  		for _, item := range items {
    70  			q.put(item)
    71  		}
    72  
    73  		j3 := q.delete("job-1")
    74  		j1 := q.delete("job-3")
    75  		jx := q.delete("job-x")
    76  		assert.Nil(t, jx)
    77  
    78  		q.put(j1)
    79  		q.put(j3)
    80  		q.put(j3)
    81  		q.put(j1)
    82  
    83  		q.put(&raft_log.CompactionJobState{
    84  			Name:            "job-4",
    85  			CompactionLevel: 0,
    86  			Status:          metastorev1.CompactionJobStatus_COMPACTION_STATUS_IN_PROGRESS,
    87  			LeaseExpiresAt:  3, // Should be after job-3.
    88  			Failures:        5,
    89  		})
    90  
    91  		expected := []string{"job-6", "job-2", "job-3", "job-4", "job-0", "job-1", "job-5"}
    92  		dequeued := make([]string, 0, len(items))
    93  		for range items {
    94  			x := jobQueuePop(q)
    95  			assert.NotNil(t, x)
    96  			dequeued = append(dequeued, x.Name)
    97  		}
    98  		assert.Equal(t, expected, dequeued)
    99  		assert.Nil(t, jobQueuePop(q))
   100  	}
   101  
   102  	rnd := rand.New(rand.NewSource(123))
   103  	for i := 0; i < 25; i++ {
   104  		rnd.Shuffle(len(items), func(i, j int) {
   105  			items[i], items[j] = items[j], items[i]
   106  		})
   107  		test(items)
   108  	}
   109  }
   110  
   111  func TestJobQueue_delete(t *testing.T) {
   112  	q := newJobQueue()
   113  	items := []*raft_log.CompactionJobState{
   114  		{Name: "job-1"},
   115  		{Name: "job-2"},
   116  		{Name: "job-3"},
   117  		{Name: "job-4"},
   118  	}
   119  
   120  	for _, item := range items {
   121  		q.put(item)
   122  	}
   123  
   124  	for _, item := range items {
   125  		q.delete(item.Name)
   126  	}
   127  
   128  	assert.Nil(t, jobQueuePop(q))
   129  }
   130  
   131  func TestJobQueue_empty(t *testing.T) {
   132  	q := newJobQueue()
   133  	q.delete("job-1")
   134  	assert.Nil(t, jobQueuePop(q))
   135  	q.put(&raft_log.CompactionJobState{Name: "job-1"})
   136  	q.delete("job-1")
   137  	assert.Nil(t, jobQueuePop(q))
   138  }
   139  
   140  // The function is for testing purposes only.
   141  func jobQueuePop(q *schedulerQueue) *raft_log.CompactionJobState {
   142  	for i := range q.levels {
   143  		level := q.level(uint32(i))
   144  		if level.jobs.Len() > 0 {
   145  			x := heap.Pop(level.jobs).(*jobEntry).CompactionJobState
   146  			delete(q.jobs, x.Name)
   147  			return x
   148  		}
   149  	}
   150  	return nil
   151  }