github.com/yankunsam/loki/v2@v2.6.3-0.20220817130409-389df5235c27/pkg/storage/stores/indexshipper/gatewayclient/gateway_client_test.go (about)

     1  package gatewayclient
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"fmt"
     7  	"net"
     8  	"testing"
     9  
    10  	"github.com/grafana/dskit/flagext"
    11  	"github.com/prometheus/client_golang/prometheus"
    12  	"github.com/stretchr/testify/require"
    13  	"github.com/weaveworks/common/user"
    14  	"google.golang.org/grpc"
    15  
    16  	"github.com/grafana/loki/pkg/logproto"
    17  	"github.com/grafana/loki/pkg/storage/stores/series/index"
    18  	"github.com/grafana/loki/pkg/storage/stores/shipper/indexgateway"
    19  	"github.com/grafana/loki/pkg/storage/stores/shipper/util"
    20  	util_log "github.com/grafana/loki/pkg/util/log"
    21  )
    22  
    23  const (
    24  	// query prefixes
    25  	tableNamePrefix        = "table-name"
    26  	hashValuePrefix        = "hash-value"
    27  	rangeValuePrefixPrefix = "range-value-prefix"
    28  	rangeValueStartPrefix  = "range-value-start"
    29  	valueEqualPrefix       = "value-equal"
    30  
    31  	// response prefixes
    32  	rangeValuePrefix = "range-value"
    33  	valuePrefix      = "value"
    34  
    35  	// the number of index entries for benchmarking will be divided amongst numTables
    36  	//benchMarkNumEntries = 1000000
    37  	//numTables           = 50
    38  )
    39  
    40  type mockIndexGatewayServer struct {
    41  	logproto.IndexGatewayServer
    42  }
    43  
    44  func (m mockIndexGatewayServer) QueryIndex(request *logproto.QueryIndexRequest, server logproto.IndexGateway_QueryIndexServer) error {
    45  	for i, query := range request.Queries {
    46  		resp := logproto.QueryIndexResponse{
    47  			QueryKey: "",
    48  			Rows:     nil,
    49  		}
    50  
    51  		if query.TableName != fmt.Sprintf("%s%d", tableNamePrefix, i) {
    52  			return errors.New("incorrect TableName in query")
    53  		}
    54  		if query.HashValue != fmt.Sprintf("%s%d", hashValuePrefix, i) {
    55  			return errors.New("incorrect HashValue in query")
    56  		}
    57  		if string(query.RangeValuePrefix) != fmt.Sprintf("%s%d", rangeValuePrefixPrefix, i) {
    58  			return errors.New("incorrect RangeValuePrefix in query")
    59  		}
    60  		if string(query.RangeValueStart) != fmt.Sprintf("%s%d", rangeValueStartPrefix, i) {
    61  			return errors.New("incorrect RangeValueStart in query")
    62  		}
    63  		if string(query.ValueEqual) != fmt.Sprintf("%s%d", valueEqualPrefix, i) {
    64  			return errors.New("incorrect ValueEqual in query")
    65  		}
    66  
    67  		for j := 0; j <= i; j++ {
    68  			resp.Rows = append(resp.Rows, &logproto.Row{
    69  				RangeValue: []byte(fmt.Sprintf("%s%d", rangeValuePrefix, j)),
    70  				Value:      []byte(fmt.Sprintf("%s%d", valuePrefix, j)),
    71  			})
    72  		}
    73  
    74  		resp.QueryKey = util.QueryKey(index.Query{
    75  			TableName:        query.TableName,
    76  			HashValue:        query.HashValue,
    77  			RangeValuePrefix: query.RangeValuePrefix,
    78  			RangeValueStart:  query.RangeValueStart,
    79  			ValueEqual:       query.ValueEqual,
    80  		})
    81  
    82  		if err := server.Send(&resp); err != nil {
    83  			return err
    84  		}
    85  	}
    86  
    87  	return nil
    88  }
    89  
    90  func (m mockIndexGatewayServer) GetChunkRef(context.Context, *logproto.GetChunkRefRequest) (*logproto.GetChunkRefResponse, error) {
    91  	return &logproto.GetChunkRefResponse{}, nil
    92  }
    93  
    94  func createTestGrpcServer(t *testing.T) (func(), string) {
    95  	var server mockIndexGatewayServer
    96  	lis, err := net.Listen("tcp", "localhost:0")
    97  	require.NoError(t, err)
    98  	s := grpc.NewServer()
    99  
   100  	logproto.RegisterIndexGatewayServer(s, &server)
   101  	go func() {
   102  		if err := s.Serve(lis); err != nil {
   103  			t.Logf("Failed to serve: %v", err)
   104  		}
   105  	}()
   106  
   107  	return s.GracefulStop, lis.Addr().String()
   108  }
   109  
   110  func TestGatewayClient(t *testing.T) {
   111  	cleanup, storeAddress := createTestGrpcServer(t)
   112  	t.Cleanup(cleanup)
   113  
   114  	var cfg IndexGatewayClientConfig
   115  	cfg.Mode = indexgateway.SimpleMode
   116  	flagext.DefaultValues(&cfg)
   117  	cfg.Address = storeAddress
   118  
   119  	gatewayClient, err := NewGatewayClient(cfg, prometheus.DefaultRegisterer, util_log.Logger)
   120  	require.NoError(t, err)
   121  
   122  	ctx := user.InjectOrgID(context.Background(), "fake")
   123  
   124  	queries := []index.Query{}
   125  	for i := 0; i < 10; i++ {
   126  		queries = append(queries, index.Query{
   127  			TableName:        fmt.Sprintf("%s%d", tableNamePrefix, i),
   128  			HashValue:        fmt.Sprintf("%s%d", hashValuePrefix, i),
   129  			RangeValuePrefix: []byte(fmt.Sprintf("%s%d", rangeValuePrefixPrefix, i)),
   130  			RangeValueStart:  []byte(fmt.Sprintf("%s%d", rangeValueStartPrefix, i)),
   131  			ValueEqual:       []byte(fmt.Sprintf("%s%d", valueEqualPrefix, i)),
   132  		})
   133  	}
   134  
   135  	numCallbacks := 0
   136  	err = gatewayClient.QueryPages(ctx, queries, func(query index.Query, batch index.ReadBatchResult) (shouldContinue bool) {
   137  		itr := batch.Iterator()
   138  
   139  		for j := 0; j <= numCallbacks; j++ {
   140  			require.True(t, itr.Next())
   141  			require.Equal(t, fmt.Sprintf("%s%d", rangeValuePrefix, j), string(itr.RangeValue()))
   142  			require.Equal(t, fmt.Sprintf("%s%d", valuePrefix, j), string(itr.Value()))
   143  		}
   144  
   145  		require.False(t, itr.Next())
   146  		numCallbacks++
   147  		return true
   148  	})
   149  	require.NoError(t, err)
   150  
   151  	require.Equal(t, len(queries), numCallbacks)
   152  }
   153  
   154  /*
   155  ToDo(Sandeep): Comment out benchmark code for now to fix circular dependency
   156  func buildTableName(i int) string {
   157  	return fmt.Sprintf("%s%d", tableNamePrefix, i)
   158  }
   159  
   160  type mockLimits struct{}
   161  
   162  func (m mockLimits) AllByUserID() map[string]*validation.Limits {
   163  	return map[string]*validation.Limits{}
   164  }
   165  
   166  func (m mockLimits) DefaultLimits() *validation.Limits {
   167  	return &validation.Limits{}
   168  }
   169  
   170  func benchmarkIndexQueries(b *testing.B, queries []index.Query) {
   171  	buffer := 1024 * 1024
   172  	listener := bufconn.Listen(buffer)
   173  
   174  	// setup the grpc server
   175  	s := grpc.NewServer(grpc.ChainStreamInterceptor(func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
   176  		return middleware.StreamServerUserHeaderInterceptor(srv, ss, info, handler)
   177  	}))
   178  	conn, _ := grpc.DialContext(context.Background(), "", grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
   179  		return listener.Dial()
   180  	}), grpc.WithTransportCredentials(insecure.NewCredentials()))
   181  	defer func() {
   182  		s.Stop()
   183  		conn.Close()
   184  	}()
   185  
   186  	// setup test data
   187  	dir := b.TempDir()
   188  	bclient, err := local.NewBoltDBIndexClient(local.BoltDBConfig{
   189  		Directory: dir + "/boltdb",
   190  	})
   191  	require.NoError(b, err)
   192  
   193  	for i := 0; i < numTables; i++ {
   194  		// setup directory for table in both cache and object storage
   195  		tableName := buildTableName(i)
   196  		objectStorageDir := filepath.Join(dir, "index", tableName)
   197  		cacheDir := filepath.Join(dir, "cache", tableName)
   198  		require.NoError(b, os.MkdirAll(objectStorageDir, 0o777))
   199  		require.NoError(b, os.MkdirAll(cacheDir, 0o777))
   200  
   201  		// add few rows at a time to the db because doing to many writes in a single transaction puts too much strain on boltdb and makes it slow
   202  		for i := 0; i < benchMarkNumEntries/numTables; i += 10000 {
   203  			end := util_math.Min(i+10000, benchMarkNumEntries/numTables)
   204  			// setup index files in both the cache directory and object storage directory so that we don't spend time syncing files at query time
   205  			testutil.AddRecordsToDB(b, filepath.Join(objectStorageDir, "db1"), bclient, i, end-i, []byte("index"))
   206  			testutil.AddRecordsToDB(b, filepath.Join(cacheDir, "db1"), bclient, i, end-i, []byte("index"))
   207  		}
   208  	}
   209  
   210  	fs, err := local.NewFSObjectClient(local.FSConfig{
   211  		Directory: dir,
   212  	})
   213  	require.NoError(b, err)
   214  	tm, err := downloads.NewTableManager(downloads.Config{
   215  		CacheDir:          dir + "/cache",
   216  		SyncInterval:      15 * time.Minute,
   217  		CacheTTL:          15 * time.Minute,
   218  		QueryReadyNumDays: 30,
   219  		Limits:            mockLimits{},
   220  	}, bclient, storage.NewIndexStorageClient(fs, "index/"), nil, nil)
   221  	require.NoError(b, err)
   222  
   223  	// initialize the index gateway server
   224  	var cfg indexgateway.Config
   225  	flagext.DefaultValues(&cfg)
   226  
   227  	gw, err := indexgateway.NewIndexGateway(cfg, util_log.Logger, prometheus.DefaultRegisterer, nil, tm)
   228  	require.NoError(b, err)
   229  	logproto.RegisterIndexGatewayServer(s, gw)
   230  	go func() {
   231  		if err := s.Serve(listener); err != nil {
   232  			panic(err)
   233  		}
   234  	}()
   235  
   236  	// setup context for querying
   237  	ctx := user.InjectOrgID(context.Background(), "foo")
   238  	ctx, _ = user.InjectIntoGRPCRequest(ctx)
   239  
   240  	// initialize the gateway client
   241  	gatewayClient := GatewayClient{}
   242  	gatewayClient.grpcClient = logproto.NewIndexGatewayClient(conn)
   243  
   244  	// build the response we expect to get from queries
   245  	expected := map[string]int{}
   246  	for i := 0; i < benchMarkNumEntries/numTables; i++ {
   247  		expected[strconv.Itoa(i)] = numTables
   248  	}
   249  
   250  	b.ReportAllocs()
   251  	b.ResetTimer()
   252  	for i := 0; i < b.N; i++ {
   253  		actual := map[string]int{}
   254  		syncMtx := sync.Mutex{}
   255  
   256  		err := gatewayClient.QueryPages(ctx, queries, func(query index.Query, batch index.ReadBatchResult) (shouldContinue bool) {
   257  			itr := batch.Iterator()
   258  			for itr.Next() {
   259  				syncMtx.Lock()
   260  				actual[string(itr.Value())]++
   261  				syncMtx.Unlock()
   262  			}
   263  			return true
   264  		})
   265  		require.NoError(b, err)
   266  		require.Equal(b, expected, actual)
   267  	}
   268  }
   269  
   270  func Benchmark_QueriesMatchingSingleRow(b *testing.B) {
   271  	queries := []index.Query{}
   272  	// do a query per row from each of the tables
   273  	for i := 0; i < benchMarkNumEntries/numTables; i++ {
   274  		for j := 0; j < numTables; j++ {
   275  			queries = append(queries, index.Query{
   276  				TableName:        buildTableName(j),
   277  				RangeValuePrefix: []byte(strconv.Itoa(i)),
   278  				ValueEqual:       []byte(strconv.Itoa(i)),
   279  			})
   280  		}
   281  	}
   282  
   283  	benchmarkIndexQueries(b, queries)
   284  }
   285  
   286  func Benchmark_QueriesMatchingLargeNumOfRows(b *testing.B) {
   287  	var queries []index.Query
   288  	// do a query per table matching all the rows from it
   289  	for j := 0; j < numTables; j++ {
   290  		queries = append(queries, index.Query{
   291  			TableName: buildTableName(j),
   292  		})
   293  	}
   294  	benchmarkIndexQueries(b, queries)
   295  }*/
   296  
   297  func TestDoubleRegistration(t *testing.T) {
   298  	r := prometheus.NewRegistry()
   299  	cleanup, storeAddress := createTestGrpcServer(t)
   300  	t.Cleanup(cleanup)
   301  
   302  	clientCfg := IndexGatewayClientConfig{
   303  		Address: storeAddress,
   304  	}
   305  
   306  	client, err := NewGatewayClient(clientCfg, r, util_log.Logger)
   307  	require.NoError(t, err)
   308  	defer client.Stop()
   309  
   310  	client, err = NewGatewayClient(clientCfg, r, util_log.Logger)
   311  	require.NoError(t, err)
   312  	defer client.Stop()
   313  }