From 367425bba155fb8fa98701d0408c5c69b5b35506 Mon Sep 17 00:00:00 2001 From: Cody Soyland Date: Fri, 29 Jan 2021 13:12:59 -0600 Subject: [PATCH 1/4] Add duration header to all gRPC query results --- server/grpc.go | 24 ++++++++++++-- server/grpc_test.go | 76 +++++++++++++++++++++++++++++++++++++++++++-- 2 files changed, 96 insertions(+), 4 deletions(-) diff --git a/server/grpc.go b/server/grpc.go index c2e68bb1f..8ffc49ad3 100644 --- a/server/grpc.go +++ b/server/grpc.go @@ -20,6 +20,7 @@ import ( "fmt" "net" "net/http" + "strconv" "strings" "sync" "time" @@ -34,6 +35,7 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials" + "google.golang.org/grpc/metadata" "google.golang.org/grpc/reflection" "google.golang.org/grpc/status" ) @@ -146,6 +148,10 @@ func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQ return err } + stream.SendHeader(metadata.New(map[string]string{ + "duration": strconv.Itoa(int(duration)), + })) + err = newDurationRowser(results, duration).ToRows(stream.Send) if err != nil { return errors.Wrap(err, "streaming result") @@ -183,7 +189,12 @@ func (h *GRPCHandler) QuerySQLUnary(ctx context.Context, req *pb.QuerySQLRequest if err != nil { return nil, err } - table.Duration = int64(time.Since(start)) + duration := time.Since(start) + table.Duration = int64(duration) + grpc.SendHeader(ctx, metadata.New(map[string]string{ + "duration": strconv.Itoa(int(duration)), + })) + return table, nil } @@ -197,6 +208,11 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ t := time.Now() resp, err := h.api.Query(stream.Context(), &query) durQuery := time.Since(t) + + stream.SendHeader(metadata.New(map[string]string{ + "duration": strconv.Itoa(int(durQuery)), + })) + // TODO: what about resp.CollumnAttrSets? if err != nil { return errToStatusError(err) @@ -262,7 +278,11 @@ func (h *GRPCHandler) QueryPQLUnary(ctx context.Context, req *pb.QueryPQLRequest } durFormat := time.Since(t) - table.Duration = int64(durQuery + durFormat) + duration := durQuery + durFormat + table.Duration = int64(duration) + grpc.SendHeader(ctx, metadata.New(map[string]string{ + "duration": strconv.Itoa(int(duration)), + })) h.stats.Timing(pilosa.MetricGRPCUnaryQueryDurationSeconds, durQuery, 0.1) h.stats.Timing(pilosa.MetricGRPCUnaryFormatDurationSeconds, durFormat, 0.1) diff --git a/server/grpc_test.go b/server/grpc_test.go index 3cc808bb2..89ae9451f 100644 --- a/server/grpc_test.go +++ b/server/grpc_test.go @@ -29,7 +29,9 @@ import ( "github.com/pilosa/pilosa/v2/sql" "github.com/pilosa/pilosa/v2/test" "github.com/pkg/errors" + "google.golang.org/grpc" "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" ) @@ -354,9 +356,11 @@ func TestQueryPQLUnary(t *testing.T) { i := m.MustCreateIndex(t, "i", pilosa.IndexOptions{}) m.MustCreateField(t, i.Name(), "f", pilosa.OptFieldKeys()) - ctx := context.Background() gh := server.NewGRPCHandler(m.API) + stream := &MockServerTransportStream{} + ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream) + resp, err := gh.QueryPQLUnary(ctx, &pb.QueryPQLRequest{ Index: i.Name(), Pql: `Set(0, f="zero")`, @@ -369,6 +373,10 @@ func TestQueryPQLUnary(t *testing.T) { if resp.Duration == 0 { t.Fatal("duration not recorded") } + duration, err := stream.GetDuration() + if duration == 0 || err != nil { + t.Fatal("duration header not recorded") + } _, err = gh.QueryPQLUnary(ctx, &pb.QueryPQLRequest{ Index: i.Name(), @@ -400,6 +408,11 @@ func TestQueryPQL(t *testing.T) { t.Fatal(err) } + duration, err := mock.GetDuration() + if duration == 0 || err != nil { + t.Fatal("duration header not recorded") + } + if len(mock.Results) != 1 { t.Fatal("expecting one result") } @@ -481,7 +494,9 @@ type ( func TestQuerySQL(t *testing.T) { - ctx := context.Background() + stream := &MockServerTransportStream{} + ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream) + gh, tearDownFunc := setUpTestQuerySQLUnary(ctx, t) defer tearDownFunc() @@ -924,6 +939,11 @@ func TestQuerySQL(t *testing.T) { if resp.Duration == 0 { t.Fatal("duration not recorded") } + duration, err := stream.GetDuration() + if duration == 0 || err != nil { + t.Fatal("duration header not recorded") + } + stream.ClearMD() tr := toTableResponse(resp) if err := test.eq(test.exp, tr); err != nil { t.Fatalf("sql: %s, error: %+v", test.sql, err) @@ -942,6 +962,10 @@ func TestQuerySQL(t *testing.T) { if mock.Results[0].Duration == 0 { t.Fatal("duration not recorded") } + duration, err := mock.GetDuration() + if duration == 0 || err != nil { + t.Fatal("duration header not recorded") + } if len(mock.Results) > 1 && mock.Results[1].Duration != 0 { t.Fatal("duration on second result expected to be zero") } @@ -1383,7 +1407,43 @@ func equalUnordered(exp tableResponse, got tableResponse) error { return nil } +type MockServerTransportStream struct { + header metadata.MD +} + +func (stream *MockServerTransportStream) Method() string { + return "" +} + +func (stream *MockServerTransportStream) SetHeader(md metadata.MD) error { + // Should probably merge md with value of stream.header, but this works since we have only one metadata value + stream.header = md + return nil +} + +func (stream *MockServerTransportStream) SendHeader(md metadata.MD) error { + stream.header = md + return nil +} + +func (stream *MockServerTransportStream) SetTrailer(md metadata.MD) error { + return nil +} + +func (stream *MockServerTransportStream) GetDuration() (int, error) { + duration, ok := stream.header["duration"] + if ok { + return strconv.Atoi(duration[0]) + } + return 0, errors.New("duration not recorded") +} + +func (stream *MockServerTransportStream) ClearMD() { + stream.header = metadata.New(map[string]string{}) +} + type mockPilosa_QuerySQLServer struct { + MockServerTransportStream pb.Pilosa_QuerySQLServer Results []*pb.RowResponse } @@ -1393,6 +1453,18 @@ func (m *mockPilosa_QuerySQLServer) Send(result *pb.RowResponse) error { return nil } +func (m *mockPilosa_QuerySQLServer) SendHeader(md metadata.MD) error { + return m.MockServerTransportStream.SendHeader(md) +} + +func (m *mockPilosa_QuerySQLServer) SetHeader(md metadata.MD) error { + return m.MockServerTransportStream.SetHeader(md) +} + +func (m *mockPilosa_QuerySQLServer) SetTrailer(md metadata.MD) { + m.MockServerTransportStream.SetTrailer(md) +} + func (m *mockPilosa_QuerySQLServer) Context() context.Context { return context.Background() } From e2331372d89e5d741e8e5d70d0baef513fe49edc Mon Sep 17 00:00:00 2001 From: Cody Soyland Date: Mon, 1 Feb 2021 09:23:30 -0600 Subject: [PATCH 2/4] Handle errors --- server/grpc.go | 20 ++++++++++++++++---- server/grpc_test.go | 2 +- 2 files changed, 17 insertions(+), 5 deletions(-) diff --git a/server/grpc.go b/server/grpc.go index 8ffc49ad3..f50c9ac80 100644 --- a/server/grpc.go +++ b/server/grpc.go @@ -148,9 +148,12 @@ func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQ return err } - stream.SendHeader(metadata.New(map[string]string{ + err = stream.SendHeader(metadata.New(map[string]string{ "duration": strconv.Itoa(int(duration)), })) + if err != nil { + return errors.Wrap(err, "sending header") + } err = newDurationRowser(results, duration).ToRows(stream.Send) if err != nil { @@ -191,9 +194,12 @@ func (h *GRPCHandler) QuerySQLUnary(ctx context.Context, req *pb.QuerySQLRequest } duration := time.Since(start) table.Duration = int64(duration) - grpc.SendHeader(ctx, metadata.New(map[string]string{ + err = grpc.SendHeader(ctx, metadata.New(map[string]string{ "duration": strconv.Itoa(int(duration)), })) + if err != nil { + return nil, errors.Wrap(err, "sending header") + } return table, nil } @@ -209,9 +215,12 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ resp, err := h.api.Query(stream.Context(), &query) durQuery := time.Since(t) - stream.SendHeader(metadata.New(map[string]string{ + err = stream.SendHeader(metadata.New(map[string]string{ "duration": strconv.Itoa(int(durQuery)), })) + if err != nil { + return errors.Wrap(err, "sending header") + } // TODO: what about resp.CollumnAttrSets? if err != nil { @@ -280,9 +289,12 @@ func (h *GRPCHandler) QueryPQLUnary(ctx context.Context, req *pb.QueryPQLRequest duration := durQuery + durFormat table.Duration = int64(duration) - grpc.SendHeader(ctx, metadata.New(map[string]string{ + err = grpc.SendHeader(ctx, metadata.New(map[string]string{ "duration": strconv.Itoa(int(duration)), })) + if err != nil { + return nil, errors.Wrap(err, "sending header") + } h.stats.Timing(pilosa.MetricGRPCUnaryQueryDurationSeconds, durQuery, 0.1) h.stats.Timing(pilosa.MetricGRPCUnaryFormatDurationSeconds, durFormat, 0.1) diff --git a/server/grpc_test.go b/server/grpc_test.go index 89ae9451f..843fe2fe1 100644 --- a/server/grpc_test.go +++ b/server/grpc_test.go @@ -1462,7 +1462,7 @@ func (m *mockPilosa_QuerySQLServer) SetHeader(md metadata.MD) error { } func (m *mockPilosa_QuerySQLServer) SetTrailer(md metadata.MD) { - m.MockServerTransportStream.SetTrailer(md) + _ = m.MockServerTransportStream.SetTrailer(md) } func (m *mockPilosa_QuerySQLServer) Context() context.Context { From c30e3f0c2d320fe6e5aae81ac120a82f0b918b9d Mon Sep 17 00:00:00 2001 From: Cody Soyland Date: Mon, 1 Feb 2021 09:30:24 -0600 Subject: [PATCH 3/4] Move duration header to fix error handling --- server/grpc.go | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/server/grpc.go b/server/grpc.go index f50c9ac80..f7fe1192c 100644 --- a/server/grpc.go +++ b/server/grpc.go @@ -215,13 +215,6 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ resp, err := h.api.Query(stream.Context(), &query) durQuery := time.Since(t) - err = stream.SendHeader(metadata.New(map[string]string{ - "duration": strconv.Itoa(int(durQuery)), - })) - if err != nil { - return errors.Wrap(err, "sending header") - } - // TODO: what about resp.CollumnAttrSets? if err != nil { return errToStatusError(err) @@ -240,6 +233,13 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ return errors.Wrap(err, "wrapping as type ToRowser") } + err = stream.SendHeader(metadata.New(map[string]string{ + "duration": strconv.Itoa(int(durQuery)), + })) + if err != nil { + return errors.Wrap(err, "sending header") + } + t = time.Now() if err := newDurationRowser(toRowser, durQuery).ToRows(stream.Send); err != nil { return errToStatusError(err) From 0e22ed71ccd3ce533820ce3ff09784000fa9d125 Mon Sep 17 00:00:00 2001 From: Cody Soyland Date: Mon, 1 Feb 2021 10:18:38 -0600 Subject: [PATCH 4/4] Fix instances of context.Background that need mocked context --- server/grpc_test.go | 6 ++++-- server/handler_test.go | 5 ++++- 2 files changed, 8 insertions(+), 3 deletions(-) diff --git a/server/grpc_test.go b/server/grpc_test.go index 843fe2fe1..d24aa0cc3 100644 --- a/server/grpc_test.go +++ b/server/grpc_test.go @@ -977,7 +977,8 @@ func TestQuerySQL(t *testing.T) { func TestQuerySQLUnaryWithError(t *testing.T) { - ctx := context.Background() + stream := &MockServerTransportStream{} + ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream) gh, tearDownFunc := setUpTestQuerySQLUnary(ctx, t) defer tearDownFunc() @@ -1023,7 +1024,8 @@ func TestCRUDIndexes(t *testing.T) { m := test.RunCommand(t) defer m.Close() - ctx := context.Background() + stream := &MockServerTransportStream{} + ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream) gh := server.NewGRPCHandler(m.API) t.Run("CreateIndex", func(t *testing.T) { diff --git a/server/handler_test.go b/server/handler_test.go index 0ffaf3d0e..834399587 100644 --- a/server/handler_test.go +++ b/server/handler_test.go @@ -40,6 +40,7 @@ import ( pb "github.com/pilosa/pilosa/v2/proto" "github.com/pilosa/pilosa/v2/server" "github.com/pilosa/pilosa/v2/test" + "google.golang.org/grpc" ) func TestHandler_PostSchemaCluster(t *testing.T) { @@ -1517,7 +1518,9 @@ func TestQueryHistory(t *testing.T) { test.Do(t, "POST", cmd.URL()+"/index/i0/field/f0", "") gh := server.NewGRPCHandler(cmd.API) - _, err = gh.QuerySQLUnary(context.Background(), &pb.QuerySQLRequest{ + stream := &MockServerTransportStream{} + ctx := grpc.NewContextWithServerTransportStream(context.Background(), stream) + _, err = gh.QuerySQLUnary(ctx, &pb.QuerySQLRequest{ Sql: `select * from i0`, })