github.com/grafana/pyroscope@v1.18.0/pkg/validation/exporter/exporter_test.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  // Provenance-includes-location: https://github.com/phlareproject/phlare/blob/master/pkg/util/validation/exporter_test.go
     3  // Provenance-includes-license: Apache-2.0
     4  // Provenance-includes-copyright: The phlare Authors.
     5  
     6  package exporter
     7  
     8  import (
     9  	"bytes"
    10  	"context"
    11  	"testing"
    12  	"time"
    13  
    14  	"github.com/go-kit/log"
    15  	"github.com/grafana/dskit/kv/consul"
    16  	"github.com/grafana/dskit/ring"
    17  	"github.com/grafana/dskit/services"
    18  	"github.com/grafana/dskit/test"
    19  	"github.com/prometheus/client_golang/prometheus"
    20  	"github.com/prometheus/client_golang/prometheus/testutil"
    21  	"github.com/stretchr/testify/assert"
    22  	"github.com/stretchr/testify/require"
    23  
    24  	"github.com/grafana/pyroscope/pkg/validation"
    25  )
    26  
    27  func TestOverridesExporter_withConfig(t *testing.T) {
    28  	tenantLimits := map[string]*validation.Limits{
    29  		"tenant-a": {
    30  			IngestionRateMB:              10,
    31  			IngestionBurstSizeMB:         11,
    32  			MaxGlobalSeriesPerTenant:     12,
    33  			MaxLocalSeriesPerTenant:      13,
    34  			MaxLabelNameLength:           14,
    35  			MaxLabelValueLength:          15,
    36  			MaxLabelNamesPerSeries:       16,
    37  			MaxQueryLookback:             17,
    38  			MaxQueryLength:               18,
    39  			MaxQueryParallelism:          19,
    40  			QuerySplitDuration:           20,
    41  			MaxSessionsPerSeries:         21,
    42  			DistributorAggregationWindow: 22,
    43  			DistributorAggregationPeriod: 23,
    44  			MaxFlameGraphNodesDefault:    24,
    45  			MaxFlameGraphNodesMax:        25,
    46  		},
    47  	}
    48  	ringStore, closer := consul.NewInMemoryClient(ring.GetCodec(), log.NewNopLogger(), nil)
    49  	t.Cleanup(func() { assert.NoError(t, closer.Close()) })
    50  
    51  	cfg1 := Config{RingConfig{}}
    52  	cfg1.Ring.Ring.KVStore.Mock = ringStore
    53  	cfg1.Ring.Ring.InstancePort = 1234
    54  	cfg1.Ring.Ring.HeartbeatPeriod = 15 * time.Second
    55  	cfg1.Ring.Ring.HeartbeatTimeout = 1 * time.Minute
    56  
    57  	// Create an empty ring.
    58  	ctx := context.Background()
    59  	require.NoError(t, ringStore.CAS(ctx, ringKey, func(in interface{}) (out interface{}, retry bool, err error) {
    60  		return ring.NewDesc(), true, nil
    61  	}))
    62  
    63  	// Create an overrides-exporter.
    64  	cfg1.Ring.Ring.InstanceID = "overrides-exporter-1"
    65  	cfg1.Ring.Ring.InstanceAddr = "1.2.3.1"
    66  	exporter, err := NewOverridesExporter(cfg1, &validation.Limits{
    67  		IngestionRateMB:              20,
    68  		IngestionBurstSizeMB:         21,
    69  		MaxGlobalSeriesPerTenant:     22,
    70  		MaxLocalSeriesPerTenant:      23,
    71  		MaxLabelNameLength:           24,
    72  		MaxLabelValueLength:          25,
    73  		MaxLabelNamesPerSeries:       26,
    74  		MaxQueryLookback:             27,
    75  		MaxQueryLength:               28,
    76  		MaxQueryParallelism:          29,
    77  		QuerySplitDuration:           30,
    78  		MaxSessionsPerSeries:         31,
    79  		DistributorAggregationWindow: 32,
    80  		DistributorAggregationPeriod: 33,
    81  		MaxFlameGraphNodesDefault:    34,
    82  		MaxFlameGraphNodesMax:        35,
    83  	}, validation.NewMockTenantLimits(tenantLimits), log.NewNopLogger(), nil)
    84  	require.NoError(t, err)
    85  
    86  	l1 := exporter.ring.lifecycler
    87  	require.NoError(t, err)
    88  	require.NoError(t, services.StartAndAwaitRunning(context.Background(), exporter))
    89  	t.Cleanup(func() { assert.NoError(t, services.StopAndAwaitTerminated(context.Background(), exporter)) })
    90  
    91  	// Wait until it has received the ring update.
    92  	test.Poll(t, time.Second, true, func() interface{} {
    93  		rs, _ := exporter.ring.client.GetAllHealthy(ringOp)
    94  		return rs.Includes(l1.GetInstanceAddr())
    95  	})
    96  
    97  	// Set leader token.
    98  	require.NoError(t, ringStore.CAS(context.Background(), ringKey, func(in interface{}) (out interface{}, retry bool, err error) {
    99  		desc := in.(*ring.Desc)
   100  		instance := desc.Ingesters[l1.GetInstanceID()]
   101  		instance.Tokens = []uint32{leaderToken + 1}
   102  		desc.Ingesters[l1.GetInstanceID()] = instance
   103  		return desc, true, nil
   104  	}))
   105  
   106  	// Wait for update of token.
   107  	test.Poll(t, time.Second, []uint32{leaderToken + 1}, func() interface{} {
   108  		rs, _ := exporter.ring.client.GetAllHealthy(ringOp)
   109  		return rs.Instances[0].Tokens
   110  	})
   111  	limitsMetrics := `
   112  # HELP pyroscope_limits_overrides Resource limit overrides applied to tenants
   113  # TYPE pyroscope_limits_overrides gauge
   114  pyroscope_limits_overrides{limit_name="ingestion_rate_mb",tenant="tenant-a"} 10
   115  pyroscope_limits_overrides{limit_name="ingestion_burst_size_mb",tenant="tenant-a"} 11
   116  pyroscope_limits_overrides{limit_name="max_global_series_per_tenant",tenant="tenant-a"} 12
   117  pyroscope_limits_overrides{limit_name="max_series_per_tenant",tenant="tenant-a"} 13
   118  pyroscope_limits_overrides{limit_name="max_label_name_length",tenant="tenant-a"} 14
   119  pyroscope_limits_overrides{limit_name="max_label_value_length",tenant="tenant-a"} 15
   120  pyroscope_limits_overrides{limit_name="max_label_names_per_series",tenant="tenant-a"} 16
   121  pyroscope_limits_overrides{limit_name="max_query_lookback",tenant="tenant-a"} 17
   122  pyroscope_limits_overrides{limit_name="max_query_length",tenant="tenant-a"} 18
   123  pyroscope_limits_overrides{limit_name="max_query_parallelism",tenant="tenant-a"} 19
   124  pyroscope_limits_overrides{limit_name="split_queries_by_interval",tenant="tenant-a"} 20
   125  pyroscope_limits_overrides{limit_name="max_sessions_per_series",tenant="tenant-a"} 21
   126  pyroscope_limits_overrides{limit_name="distributor_aggregation_window",tenant="tenant-a"} 22
   127  pyroscope_limits_overrides{limit_name="distributor_aggregation_period",tenant="tenant-a"} 23
   128  pyroscope_limits_overrides{limit_name="max_flamegraph_nodes_default",tenant="tenant-a"} 24
   129  pyroscope_limits_overrides{limit_name="max_flamegraph_nodes_max",tenant="tenant-a"} 25
   130  `
   131  
   132  	// Make sure each override matches the values from the supplied `Limit`
   133  	err = testutil.CollectAndCompare(exporter, bytes.NewBufferString(limitsMetrics), "pyroscope_limits_overrides")
   134  	assert.NoError(t, err)
   135  
   136  	limitsMetrics = `
   137  # HELP pyroscope_limits_defaults Resource limit defaults for tenants without overrides
   138  # TYPE pyroscope_limits_defaults gauge
   139  pyroscope_limits_defaults{limit_name="ingestion_rate_mb"} 20
   140  pyroscope_limits_defaults{limit_name="ingestion_burst_size_mb"} 21
   141  pyroscope_limits_defaults{limit_name="max_global_series_per_tenant"} 22
   142  pyroscope_limits_defaults{limit_name="max_series_per_tenant"} 23
   143  pyroscope_limits_defaults{limit_name="max_label_name_length"} 24
   144  pyroscope_limits_defaults{limit_name="max_label_value_length"} 25
   145  pyroscope_limits_defaults{limit_name="max_label_names_per_series"} 26
   146  pyroscope_limits_defaults{limit_name="max_query_lookback"} 27
   147  pyroscope_limits_defaults{limit_name="max_query_length"} 28
   148  pyroscope_limits_defaults{limit_name="max_query_parallelism"} 29
   149  pyroscope_limits_defaults{limit_name="split_queries_by_interval"} 30
   150  pyroscope_limits_defaults{limit_name="max_sessions_per_series"} 31
   151  pyroscope_limits_defaults{limit_name="distributor_aggregation_window"} 32
   152  pyroscope_limits_defaults{limit_name="distributor_aggregation_period"} 33
   153  pyroscope_limits_defaults{limit_name="max_flamegraph_nodes_default"} 34
   154  pyroscope_limits_defaults{limit_name="max_flamegraph_nodes_max"} 35
   155  `
   156  	err = testutil.CollectAndCompare(exporter, bytes.NewBufferString(limitsMetrics), "pyroscope_limits_defaults")
   157  	assert.NoError(t, err)
   158  }
   159  
   160  func TestOverridesExporter_withRing(t *testing.T) {
   161  	tenantLimits := map[string]*validation.Limits{
   162  		"tenant-a": {},
   163  	}
   164  
   165  	ringStore, closer := consul.NewInMemoryClient(ring.GetCodec(), log.NewNopLogger(), nil)
   166  	t.Cleanup(func() { assert.NoError(t, closer.Close()) })
   167  
   168  	cfg1 := Config{RingConfig{}}
   169  	cfg1.Ring.Ring.KVStore.Mock = ringStore
   170  	cfg1.Ring.Ring.InstancePort = 1234
   171  	cfg1.Ring.Ring.HeartbeatPeriod = 15 * time.Second
   172  	cfg1.Ring.Ring.HeartbeatTimeout = 1 * time.Minute
   173  
   174  	// Create an empty ring.
   175  	ctx := context.Background()
   176  	require.NoError(t, ringStore.CAS(ctx, ringKey, func(in interface{}) (out interface{}, retry bool, err error) {
   177  		return ring.NewDesc(), true, nil
   178  	}))
   179  
   180  	// Create an overrides-exporter.
   181  	cfg1.Ring.Ring.InstanceID = "overrides-exporter-1"
   182  	cfg1.Ring.Ring.InstanceAddr = "1.2.3.1"
   183  	e1, err := NewOverridesExporter(cfg1, &validation.Limits{}, validation.NewMockTenantLimits(tenantLimits), log.NewNopLogger(), nil)
   184  	l1 := e1.ring.lifecycler
   185  	require.NoError(t, err)
   186  	require.NoError(t, services.StartAndAwaitRunning(ctx, e1))
   187  	t.Cleanup(func() { assert.NoError(t, services.StopAndAwaitTerminated(ctx, e1)) })
   188  
   189  	// Wait until it has received the ring update.
   190  	test.Poll(t, time.Second, true, func() interface{} {
   191  		rs, _ := e1.ring.client.GetAllHealthy(ringOp)
   192  		return rs.Includes(l1.GetInstanceAddr())
   193  	})
   194  
   195  	// Set leader token.
   196  	require.NoError(t, ringStore.CAS(ctx, ringKey, func(in interface{}) (out interface{}, retry bool, err error) {
   197  		desc := in.(*ring.Desc)
   198  		instance := desc.Ingesters[l1.GetInstanceID()]
   199  		instance.Tokens = []uint32{leaderToken + 1}
   200  		desc.Ingesters[l1.GetInstanceID()] = instance
   201  		return desc, true, nil
   202  	}))
   203  
   204  	// Wait for update of token.
   205  	test.Poll(t, time.Second, []uint32{leaderToken + 1}, func() interface{} {
   206  		rs, _ := e1.ring.client.GetAllHealthy(ringOp)
   207  		return rs.Instances[0].Tokens
   208  	})
   209  
   210  	// This instance is now the only ring member and should export metrics.
   211  	require.True(t, hasOverrideMetrics(e1))
   212  
   213  	// Register a second instance.
   214  	cfg2 := cfg1
   215  	cfg2.Ring.Ring.InstanceID = "overrides-exporter-2"
   216  	cfg2.Ring.Ring.InstanceAddr = "1.2.3.2"
   217  	e2, err := NewOverridesExporter(cfg2, &validation.Limits{}, validation.NewMockTenantLimits(tenantLimits), log.NewNopLogger(), nil)
   218  	require.NoError(t, err)
   219  	require.NoError(t, services.StartAndAwaitRunning(ctx, e2))
   220  	t.Cleanup(func() { assert.NoError(t, services.StopAndAwaitTerminated(ctx, e2)) })
   221  
   222  	// Wait until it has registered itself to the ring and both overrides-exporter instances got the updated ring.
   223  	test.Poll(t, time.Second, true, func() interface{} {
   224  		rs1, _ := e1.ring.client.GetAllHealthy(ringOp)
   225  		rs2, _ := e2.ring.client.GetAllHealthy(ringOp)
   226  		return rs1.Includes(e2.ring.lifecycler.GetInstanceAddr()) && rs2.Includes(e1.ring.lifecycler.GetInstanceAddr())
   227  	})
   228  
   229  	// Only the leader instance (owner of the special token) should export metrics.
   230  	require.True(t, hasOverrideMetrics(e1))
   231  	require.False(t, hasOverrideMetrics(e2))
   232  }
   233  
   234  func hasOverrideMetrics(e1 prometheus.Collector) bool {
   235  	return testutil.CollectAndCount(e1, "pyroscope_limits_overrides") > 0
   236  }