github.com/grafana/pyroscope@v1.18.0/pkg/usagestats/seed_test.go (about)

     1  package usagestats
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"os"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/go-kit/log"
    11  	"github.com/grafana/dskit/flagext"
    12  	"github.com/grafana/dskit/kv"
    13  	"github.com/grafana/dskit/kv/codec"
    14  	"github.com/grafana/dskit/kv/memberlist"
    15  	"github.com/grafana/dskit/services"
    16  	"github.com/stretchr/testify/require"
    17  
    18  	"github.com/grafana/pyroscope/pkg/objstore/client"
    19  	"github.com/grafana/pyroscope/pkg/objstore/providers/filesystem"
    20  )
    21  
    22  type dnsProviderMock struct {
    23  	resolved []string
    24  }
    25  
    26  func (p *dnsProviderMock) Resolve(ctx context.Context, addrs []string) error {
    27  	p.resolved = addrs
    28  	return nil
    29  }
    30  
    31  func (p dnsProviderMock) Addresses() []string {
    32  	return p.resolved
    33  }
    34  
    35  func createMemberlist(t *testing.T, port, memberID int) *memberlist.KV {
    36  	t.Helper()
    37  	var cfg memberlist.KVConfig
    38  	flagext.DefaultValues(&cfg)
    39  	cfg.TCPTransport = memberlist.TCPTransportConfig{
    40  		BindAddrs:            []string{"127.0.0.1"},
    41  		BindPort:             0,
    42  		MaxConcurrentWrites:  3,
    43  		AcquireWriterTimeout: 250 * time.Millisecond,
    44  	}
    45  	cfg.GossipInterval = 100 * time.Millisecond
    46  	cfg.GossipNodes = 3
    47  	cfg.PushPullInterval = 5 * time.Second
    48  	cfg.NodeName = fmt.Sprintf("Member-%d", memberID)
    49  	cfg.Codecs = []codec.Codec{JSONCodec}
    50  
    51  	mkv := memberlist.NewKV(cfg, log.NewNopLogger(), &dnsProviderMock{}, nil)
    52  	require.NoError(t, services.StartAndAwaitRunning(context.Background(), mkv))
    53  	if port != 0 {
    54  		_, err := mkv.JoinMembers([]string{fmt.Sprintf("127.0.0.1:%d", port)})
    55  		require.NoError(t, err, "%s failed to join the cluster: %v", memberID, err)
    56  	}
    57  	t.Cleanup(func() {
    58  		_ = services.StopAndAwaitTerminated(context.TODO(), mkv)
    59  	})
    60  	return mkv
    61  }
    62  
    63  func Test_Memberlist(t *testing.T) {
    64  	stabilityCheckInterval = 100 * time.Millisecond
    65  
    66  	objectClient, err := client.NewBucket(context.Background(), client.Config{
    67  		StorageBackendConfig: client.StorageBackendConfig{
    68  			Backend: client.Filesystem,
    69  			Filesystem: filesystem.Config{
    70  				Directory: t.TempDir(),
    71  			},
    72  		},
    73  	}, "test")
    74  	require.NoError(t, err)
    75  	result := make(chan ClusterSeed, 10)
    76  
    77  	// create a first memberlist to get a valid listening port.
    78  	initMKV := createMemberlist(t, 0, -1)
    79  
    80  	for i := 0; i < 10; i++ {
    81  		go func(i int) {
    82  			leader, err := NewReporter(Config{
    83  				Leader:  true,
    84  				Enabled: true,
    85  			}, kv.Config{
    86  				Store: "memberlist",
    87  				StoreConfig: kv.StoreConfig{
    88  					MemberlistKV: func() (*memberlist.KV, error) {
    89  						return createMemberlist(t, initMKV.GetListeningPort(), i), nil
    90  					},
    91  				},
    92  			}, objectClient, log.NewLogfmtLogger(os.Stdout), nil)
    93  			require.NoError(t, err)
    94  			leader.init(context.Background())
    95  			result <- leader.cluster
    96  		}(i)
    97  	}
    98  
    99  	var UID []string
   100  	for i := 0; i < 10; i++ {
   101  		cluster := <-result
   102  		require.NotNil(t, cluster)
   103  		UID = append(UID, cluster.UID)
   104  	}
   105  	first := UID[0]
   106  	for _, uid := range UID {
   107  		require.Equal(t, first, uid)
   108  	}
   109  }