Add duration header to all gRPC query results

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)

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()
 }