github.com/grafana/pyroscope@v1.18.0/pkg/querybackend/client/client_test.go (about)

     1  package querybackendclient
     2  
     3  import (
     4  	"context"
     5  	"flag"
     6  	"fmt"
     7  	"testing"
     8  	"time"
     9  
    10  	"github.com/go-kit/log"
    11  	"github.com/grafana/dskit/grpcclient"
    12  	"github.com/stretchr/testify/require"
    13  	"golang.org/x/sync/errgroup"
    14  	"google.golang.org/grpc"
    15  	"google.golang.org/grpc/resolver"
    16  
    17  	metastorev1 "github.com/grafana/pyroscope/api/gen/proto/go/metastore/v1"
    18  	queryv1 "github.com/grafana/pyroscope/api/gen/proto/go/query/v1"
    19  	"github.com/grafana/pyroscope/pkg/querybackend"
    20  	"github.com/grafana/pyroscope/pkg/querybackend/queryplan"
    21  	"github.com/grafana/pyroscope/pkg/test"
    22  )
    23  
    24  const (
    25  	nServers            = 12
    26  	nServerResponseTime = 200 * time.Millisecond
    27  
    28  	nBlocksInQuery     = 4000
    29  	nConcurrentQueries = 5
    30  )
    31  
    32  type QueryHandler struct {
    33  }
    34  
    35  func (q QueryHandler) Invoke(ctx context.Context, request *queryv1.InvokeRequest) (*queryv1.InvokeResponse, error) {
    36  	time.Sleep(nServerResponseTime)
    37  	return &queryv1.InvokeResponse{}, nil
    38  }
    39  
    40  type multiResolverBuilder struct {
    41  	targets []string
    42  }
    43  
    44  func (b *multiResolverBuilder) Build(target resolver.Target, cc resolver.ClientConn, opts resolver.BuildOptions) (resolver.Resolver, error) {
    45  	r := &multiResolver{
    46  		cc:      cc,
    47  		address: b.targets,
    48  	}
    49  	r.updateState()
    50  	return r, nil
    51  }
    52  
    53  func (b *multiResolverBuilder) Scheme() string {
    54  	return "multi"
    55  }
    56  
    57  // Resolves all DNS queries to a given set of IPs
    58  //
    59  // Ignores the name being resolved.
    60  type multiResolver struct {
    61  	cc      resolver.ClientConn
    62  	address []string
    63  }
    64  
    65  func (r *multiResolver) updateState() {
    66  	addresses := make([]resolver.Address, len(r.address))
    67  	for i, addr := range r.address {
    68  		addresses[i] = resolver.Address{Addr: addr}
    69  	}
    70  	_ = r.cc.UpdateState(resolver.State{Addresses: addresses})
    71  }
    72  
    73  func (r *multiResolver) ResolveNow(resolver.ResolveNowOptions) {}
    74  
    75  func (r *multiResolver) Close() {}
    76  
    77  // Test_Concurrency tests the concurrent invocation of queries against multiple backend servers.
    78  //
    79  // This test sets up a simulated environment with `nServers` gRPC servers, each acting as a
    80  // query backend. It uses `bufconn.Listener` for in-memory gRPC communication to avoid
    81  // actual network I/O.
    82  func Test_Concurrency(t *testing.T) {
    83  	addresses := make([]string, 0, nServers)
    84  	for i := 0; i < nServers; i++ {
    85  		address := fmt.Sprintf("localhost:%d", 10004+i)
    86  		addresses = append(addresses, address)
    87  	}
    88  
    89  	listeners, dialOpt := test.CreateInMemoryListeners(addresses)
    90  
    91  	grpcClientCfg := grpcclient.Config{}
    92  	grpcClientCfg.RegisterFlags(flag.NewFlagSet("", flag.PanicOnError))
    93  
    94  	resolver.Register(&multiResolverBuilder{targets: addresses})
    95  	backendAddress := "multi:///"
    96  
    97  	cl, err := New(backendAddress, grpcClientCfg, 30*time.Second, dialOpt)
    98  	require.NoError(t, err)
    99  
   100  	for i := 0; i < nServers; i++ {
   101  		gclInterceptor, err := querybackend.CreateConcurrencyInterceptor(log.NewNopLogger())
   102  		require.NoError(t, err)
   103  
   104  		b, err := querybackend.New(querybackend.Config{
   105  			Address:          backendAddress,
   106  			GRPCClientConfig: grpcClientCfg,
   107  		}, test.NewTestingLogger(t), nil, cl, QueryHandler{})
   108  		require.NoError(t, err)
   109  
   110  		grpcOptions := []grpc.ServerOption{
   111  			grpc.ChainUnaryInterceptor(gclInterceptor),
   112  		}
   113  		serv := grpc.NewServer(grpcOptions...)
   114  		require.NoError(t, err)
   115  
   116  		queryv1.RegisterQueryBackendServiceServer(serv, b)
   117  
   118  		go func() {
   119  			require.NoError(t, serv.Serve(listeners[addresses[i]]))
   120  		}()
   121  	}
   122  
   123  	blocks := make([]*metastorev1.BlockMeta, 0, nBlocksInQuery)
   124  	for i := 0; i < nBlocksInQuery; i++ {
   125  		blocks = append(blocks, &metastorev1.BlockMeta{
   126  			Id: fmt.Sprintf("block-%d", i),
   127  		})
   128  	}
   129  
   130  	g, ctx := errgroup.WithContext(context.Background())
   131  	for i := 0; i < nConcurrentQueries; i++ {
   132  		g.Go(func() error {
   133  			resp, err := cl.Invoke(ctx, &queryv1.InvokeRequest{
   134  				QueryPlan: queryplan.Build(blocks, 4, 20),
   135  			})
   136  			require.NoError(t, err)
   137  			require.NotNil(t, resp)
   138  			return err
   139  		})
   140  	}
   141  	err = g.Wait()
   142  	require.NoError(t, err)
   143  }