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

     1  package scheduler
     2  
     3  import (
     4  	"bytes"
     5  	"os"
     6  	"sync"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/grafana/dskit/multierror"
    11  	"github.com/hashicorp/raft"
    12  	"github.com/prometheus/client_golang/prometheus"
    13  	"github.com/prometheus/client_golang/prometheus/testutil"
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  	"go.etcd.io/bbolt"
    17  
    18  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    19  	"github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1/raft_log"
    20  	"github.com/grafana/pyroscope/pkg/metastore/compaction/scheduler/store"
    21  	"github.com/grafana/pyroscope/pkg/test"
    22  )
    23  
    24  func TestCollectorRegistration(t *testing.T) {
    25  	reg := prometheus.NewRegistry()
    26  	config := Config{
    27  		MaxFailures:   5,
    28  		LeaseDuration: 15 * time.Second,
    29  	}
    30  
    31  	for i := 0; i < 2; i++ {
    32  		sc := NewScheduler(config, nil, reg)
    33  		sc.queue.put(&raft_log.CompactionJobState{Name: "a"})
    34  		sc.queue.put(&raft_log.CompactionJobState{
    35  			Name: "b", CompactionLevel: 1, Token: 1,
    36  			Status: metastorev1.CompactionJobStatus_COMPACTION_STATUS_IN_PROGRESS,
    37  		})
    38  		sc.queue.delete("a")
    39  	}
    40  }
    41  
    42  func TestCollectorCollect(t *testing.T) {
    43  	reg := prometheus.NewRegistry()
    44  	config := Config{
    45  		MaxFailures:   5,
    46  		LeaseDuration: 15 * time.Second,
    47  	}
    48  
    49  	sc := NewScheduler(config, nil, reg)
    50  	sc.queue.put(&raft_log.CompactionJobState{
    51  		Name: "a", CompactionLevel: 0, Token: 1,
    52  		Status: metastorev1.CompactionJobStatus_COMPACTION_STATUS_UNSPECIFIED,
    53  	})
    54  	sc.queue.put(&raft_log.CompactionJobState{
    55  		Name: "b", CompactionLevel: 2, Token: 1,
    56  		Status: metastorev1.CompactionJobStatus_COMPACTION_STATUS_UNSPECIFIED,
    57  	})
    58  	sc.queue.delete("a")
    59  
    60  	buf, err := os.ReadFile("testdata/metrics.txt")
    61  	require.NoError(t, err)
    62  	assert.NoError(t, testutil.GatherAndCompare(reg, bytes.NewReader(buf)))
    63  }
    64  
    65  func TestCollectorCollectRace(t *testing.T) {
    66  	reg := prometheus.NewRegistry()
    67  	config := Config{
    68  		MaxFailures:   5,
    69  		LeaseDuration: 15 * time.Second,
    70  	}
    71  
    72  	var wg sync.WaitGroup
    73  	wg.Add(2)
    74  
    75  	db := test.BoltDB(t)
    76  	go func() {
    77  		defer wg.Done()
    78  		s := store.NewJobStore()
    79  		sc := NewScheduler(config, s, reg)
    80  		for i := 0; i < 100; i++ {
    81  			require.NoError(t, db.Update(func(tx *bbolt.Tx) error {
    82  				var merr multierror.MultiError
    83  				merr.Add(s.CreateBuckets(tx))
    84  				merr.Add(sc.UpdateSchedule(tx, &raft_log.CompactionPlanUpdate{
    85  					NewJobs: []*raft_log.NewCompactionJob{{
    86  						State: &raft_log.CompactionJobState{Name: "a", CompactionLevel: 0, Token: 1},
    87  						Plan:  &raft_log.CompactionJobPlan{Name: "a", CompactionLevel: 0},
    88  					}},
    89  				}))
    90  				assigned, err := sc.NewSchedule(tx, &raft.Log{}).AssignJob()
    91  				require.NoError(t, err)
    92  				require.NotNil(t, assigned)
    93  				require.Equal(t, "a", assigned.State.Name)
    94  				merr.Add(sc.UpdateSchedule(tx, &raft_log.CompactionPlanUpdate{
    95  					CompletedJobs: []*raft_log.CompletedCompactionJob{
    96  						{State: &raft_log.CompactionJobState{Name: "a", CompactionLevel: 0, Token: 1}},
    97  					},
    98  				}))
    99  				return merr.Err()
   100  			}))
   101  		}
   102  	}()
   103  
   104  	go func() {
   105  		defer wg.Done()
   106  		for i := 0; i < 1000; i++ {
   107  			_, err := testutil.GatherAndCount(reg)
   108  			require.NoError(t, err)
   109  		}
   110  	}()
   111  
   112  	wg.Wait()
   113  }