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 }