github.com/grafana/pyroscope@v1.18.0/pkg/scheduler/schedulerdiscovery/ring_test.go (about)

     1  // SPDX-License-Identifier: AGPL-3.0-only
     2  
     3  package schedulerdiscovery
     4  
     5  import (
     6  	"net"
     7  	"strconv"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/go-kit/log"
    12  	"github.com/grafana/dskit/flagext"
    13  	"github.com/grafana/dskit/ring"
    14  	"github.com/stretchr/testify/assert"
    15  	"github.com/stretchr/testify/require"
    16  )
    17  
    18  func TestRingConfig_DefaultConfigToBasicLifecyclerConfig(t *testing.T) {
    19  	cfg := Config{}
    20  	flagext.DefaultValues(&cfg)
    21  	cfg.SchedulerRing.InstanceAddr = "127.0.0.1"
    22  	cfg.SchedulerRing.InstancePort = 9095
    23  
    24  	expected := ring.BasicLifecyclerConfig{
    25  		ID:                              cfg.SchedulerRing.InstanceID,
    26  		Addr:                            net.JoinHostPort(cfg.SchedulerRing.InstanceAddr, strconv.Itoa(cfg.SchedulerRing.InstancePort)),
    27  		HeartbeatPeriod:                 cfg.SchedulerRing.HeartbeatPeriod,
    28  		HeartbeatTimeout:                cfg.SchedulerRing.HeartbeatTimeout,
    29  		TokensObservePeriod:             0,
    30  		NumTokens:                       1,
    31  		KeepInstanceInTheRingOnShutdown: false,
    32  	}
    33  
    34  	actual, err := toBasicLifecyclerConfig(cfg.SchedulerRing, log.NewNopLogger())
    35  	require.NoError(t, err)
    36  	assert.Equal(t, expected, actual)
    37  }
    38  
    39  func TestRingConfig_CustomConfigToBasicLifecyclerConfig(t *testing.T) {
    40  	// Customize the query-scheduler ring config
    41  	cfg := Config{}
    42  	flagext.DefaultValues(&cfg)
    43  	cfg.SchedulerRing.HeartbeatPeriod = 1 * time.Second
    44  	cfg.SchedulerRing.HeartbeatTimeout = 10 * time.Second
    45  	cfg.SchedulerRing.InstanceID = "test"
    46  	cfg.SchedulerRing.InstancePort = 10
    47  	cfg.SchedulerRing.InstanceAddr = "1.2.3.4"
    48  	cfg.SchedulerRing.ListenPort = 10
    49  
    50  	// The lifecycler config should be generated based upon the query-scheduler
    51  	// ring config
    52  	expected := ring.BasicLifecyclerConfig{
    53  		ID:                              "test",
    54  		Addr:                            net.JoinHostPort(cfg.SchedulerRing.InstanceAddr, strconv.Itoa(cfg.SchedulerRing.InstancePort)),
    55  		HeartbeatPeriod:                 1 * time.Second,
    56  		HeartbeatTimeout:                10 * time.Second,
    57  		TokensObservePeriod:             0,
    58  		NumTokens:                       1,
    59  		KeepInstanceInTheRingOnShutdown: false,
    60  	}
    61  
    62  	actual, err := toBasicLifecyclerConfig(cfg.SchedulerRing, log.NewNopLogger())
    63  	require.NoError(t, err)
    64  	assert.Equal(t, expected, actual)
    65  }
    66  
    67  func TestRingConfig_AddressFamilies(t *testing.T) {
    68  	cfg := Config{}
    69  	flagext.DefaultValues(&cfg)
    70  
    71  	t.Run("IPv4", func(t *testing.T) {
    72  		cfg.SchedulerRing.InstanceAddr = "1.2.3.4"
    73  		cfg.SchedulerRing.InstancePort = 10
    74  		actual, err := toBasicLifecyclerConfig(cfg.SchedulerRing, log.NewNopLogger())
    75  		require.NoError(t, err)
    76  		assert.Equal(t, "1.2.3.4:10", actual.Addr)
    77  	})
    78  
    79  	t.Run("IPv6", func(t *testing.T) {
    80  		cfg.SchedulerRing.InstanceAddr = "::1"
    81  		cfg.SchedulerRing.InstancePort = 10
    82  		cfg.SchedulerRing.EnableIPv6 = true
    83  		actual, err := toBasicLifecyclerConfig(cfg.SchedulerRing, log.NewNopLogger())
    84  		require.NoError(t, err)
    85  		assert.Equal(t, "[::1]:10", actual.Addr)
    86  	})
    87  }