diff --git a/server/grpc.go b/server/grpc.go index c2e68bb1f..f7fe1192c 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,13 @@ func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQ return err } + 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 { return errors.Wrap(err, "streaming result") @@ -183,7 +192,15 @@ 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) + 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 } @@ -197,6 +214,7 @@ 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) + // TODO: what about resp.CollumnAttrSets? if err != nil { return errToStatusError(err) @@ -215,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) @@ -262,7 +287,14 @@ 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) + 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 3cc808bb2..d24aa0cc3 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") } @@ -953,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() @@ -999,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) { @@ -1383,7 +1409,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 +1455,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() } 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`, })