diff --git a/http_handler.go b/http_handler.go index 3364d8c39..f0a9d8d48 100644 --- a/http_handler.go +++ b/http_handler.go @@ -435,11 +435,19 @@ func (h *Handler) monitorPerformance(next http.Handler) http.Handler { prefixes["query-history"] = struct{}{} pathParts := strings.Split(r.URL.Path, "/") - if len(pathParts) > 1 { + // checks < 5 to exclude import endpoints for now from sentry transactions as + // there could potentially be many of them. Pricing is per transaction. + if len(pathParts) > 1 && len(pathParts) < 5 { if _, ok := prefixes[pathParts[1]]; ok { path := scrubPath(pathParts) - txName := fmt.Sprintf("URL: %s, Method: %s", path, r.Method) - span := monitor.StartSpan(r.Context(), "http", txName) + span := monitor.StartSpan(r.Context(), "HTTP", path) + qreq := r.Context().Value(contextKeyQueryRequest) + req, ok := qreq.(*QueryRequest) + if ok && req != nil { + span.SetTag("PQL Query", req.Query) + span.SetTag("SQL Query", req.SQLQuery) + span.SetTag("Index", mux.Vars(r)["index"]) + } next.ServeHTTP(w, r) span.Finish() return diff --git a/server/grpc.go b/server/grpc.go index ed2954b87..4c100bb80 100644 --- a/server/grpc.go +++ b/server/grpc.go @@ -17,6 +17,7 @@ import ( "github.com/molecula/featurebase/v3/authn" "github.com/molecula/featurebase/v3/authz" "github.com/molecula/featurebase/v3/logger" + "github.com/molecula/featurebase/v3/monitor" "github.com/molecula/featurebase/v3/pql" pb "github.com/molecula/featurebase/v3/proto" vdsm_pb "github.com/molecula/featurebase/v3/proto/vdsm" @@ -191,9 +192,13 @@ func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQ LogQuery(ctx, "QuerySQL", req, h.queryLogger) } + span := monitor.StartSpan(ctx, "GRPC", "/pilosa.Pilosa/QuerySQL") + + span.SetTag("SQL Query", req.Sql) start := time.Now() results, err := h.execSQL(ctx, req.Sql) duration := time.Since(start) + monitor.Finish(span) if err != nil { return err } @@ -302,9 +307,13 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ } LogQuery(ctx, "QueryPQL", req, h.queryLogger) } + span := monitor.StartSpan(ctx, "GRPC", "/pilosa.Pilosa/QueryPQL") + span.SetTag("PQL Query", req.Pql) + span.SetTag("Index", req.Index) t := time.Now() resp, err := h.api.Query(stream.Context(), &query) durQuery := time.Since(t) + monitor.Finish(span) if err != nil { return errToStatusError(err) @@ -1605,40 +1614,47 @@ func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) { creds := credentials.NewTLS(server.tlsConfig) gopts = append(gopts, grpc.Creds(creds)) } - //if auth enabled - if server.auth != nil { - gopts = append(gopts, grpc.UnaryInterceptor( - func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { - ctx, err := Valid(ctx, server.auth) - if err != nil { - return nil, err - } - LogQuery(ctx, info.FullMethod, req, server.logger) - // reset the molecula-chip cookie just in case the token was refreshed - md, ok := metadata.FromIncomingContext(ctx) - if uinfo, yeah := ctx.Value("userinfo").(*authn.UserInfo); ok && yeah { - server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken) - } - return handler(ctx, req) - }, - )) - gopts = append(gopts, grpc.StreamInterceptor( - func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { - ctx, err := Valid(ss.Context(), server.auth) - if err != nil { - return err - } - // reset the molecula-chip cookie just in case the token was refreshed - md, ok := metadata.FromIncomingContext(ctx) - if uinfo, yeah := ctx.Value("userinfo").(*authn.UserInfo); ok && yeah { - server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken) - } - return handler(srv, &wrappedStream{ss, ctx}) - }, - )) + // gRPC doesn’t allow multiple interceptors so they have to be manually chained. + var unaryInterceptors []grpc.UnaryServerInterceptor + var streamInterceptors []grpc.StreamServerInterceptor + + if server.auth != nil { + unaryInterceptors = append(unaryInterceptors, func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + ctx, err := Valid(ctx, server.auth) + if err != nil { + return nil, err + } + LogQuery(ctx, info.FullMethod, req, server.logger) + + // reset the molecula-chip cookie just in case the token was refreshed + md, ok := metadata.FromIncomingContext(ctx) + if uinfo, yeah := ctx.Value("userinfo").(*authn.UserInfo); ok && yeah { + server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken) + } + return handler(ctx, req) + }) + streamInterceptors = append(streamInterceptors, func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + ctx, err := Valid(ss.Context(), server.auth) + if err != nil { + return err + } + // reset the molecula-chip cookie just in case the token was refreshed + md, ok := metadata.FromIncomingContext(ctx) + if uinfo, yeah := ctx.Value("userinfo").(*authn.UserInfo); ok && yeah { + server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken) + } + return handler(srv, &wrappedStream{ss, ctx}) + }) } + if monitor.IsOn() { + unaryInterceptors = append(unaryInterceptors, monitorUnaryInterceptor) + } + + gopts = append(gopts, grpc.UnaryInterceptor(ChainUnaryInterceptor(unaryInterceptors...))) + gopts = append(gopts, grpc.StreamInterceptor(ChainStreamInterceptors(streamInterceptors...))) + // create grpc server server.grpcServer = grpc.NewServer(gopts...) grpcHandler := NewGRPCHandler(server.api).WithLogger(server.logger).WithStats(server.stats).WithQueryLogger(server.queryLogger) @@ -1755,3 +1771,54 @@ func getTokensFromMetadata(md metadata.MD) (string, string) { } return strings.TrimPrefix(access[0], "Bearer "), refresh[0] } + +func monitorUnaryInterceptor(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + if !monitor.IsOn() { + return handler(ctx, req) + } + + span := monitor.StartSpan(ctx, "GRPC", info.FullMethod) + switch r := req.(type) { + case *pb.QueryPQLRequest: + span.SetTag("PQL Query", r.Pql) + span.SetTag("Index", r.Index) + case *pb.QuerySQLRequest: + span.SetTag("SQL Query", r.Sql) + } + resp, err := handler(ctx, req) + monitor.Finish(span) + return resp, err +} + +// Chains together multiple unary interceptors. +func ChainUnaryInterceptor(interceptors ...grpc.UnaryServerInterceptor) grpc.UnaryServerInterceptor { + chain := func(interceptor grpc.UnaryServerInterceptor, unaryHandler grpc.UnaryHandler, info *grpc.UnaryServerInfo) grpc.UnaryHandler { + return func(ctx context.Context, req interface{}) (interface{}, error) { + return interceptor(ctx, req, info, unaryHandler) + } + } + + return func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp interface{}, err error) { + chained := handler + for i := len(interceptors) - 1; i >= 0; i-- { + chained = chain(interceptors[i], chained, info) + } + return chained(ctx, req) + } +} + +// Chains together multiple stream interceptors. +func ChainStreamInterceptors(interceptors ...grpc.StreamServerInterceptor) grpc.StreamServerInterceptor { + chain := func(interceptor grpc.StreamServerInterceptor, streamHandler grpc.StreamHandler, info *grpc.StreamServerInfo) grpc.StreamHandler { + return func(srv interface{}, stream grpc.ServerStream) error { + return interceptor(srv, stream, info, streamHandler) + } + } + return func(srv interface{}, stream grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) (err error) { + chained := handler + for i := len(interceptors) - 1; i >= 0; i-- { + chained = chain(interceptors[i], chained, info) + } + return chained(srv, stream) + } +} diff --git a/server/grpc_test.go b/server/grpc_test.go index d29c8df2e..57344bd47 100644 --- a/server/grpc_test.go +++ b/server/grpc_test.go @@ -25,6 +25,7 @@ import ( "github.com/molecula/featurebase/v3/server" "github.com/molecula/featurebase/v3/sql" "github.com/molecula/featurebase/v3/test" + "github.com/molecula/featurebase/v3/vprint" "github.com/pkg/errors" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -1859,3 +1860,152 @@ func writeTestFile(t *testing.T, filename, content string) string { defer f.Close() return fname } + +func Test_ChainUnaryInterceptor(t *testing.T) { + + salt := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + r := req.(*pb.QueryPQLRequest) + r.Pql = fmt.Sprintf("%s, add salt", r.Pql) + return handler(ctx, r) + + } + pepper := func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) { + r := req.(*pb.QueryPQLRequest) + r.Pql = fmt.Sprintf("%s, add pepper", r.Pql) + return handler(ctx, r) + + } + + interceptors0 := []grpc.UnaryServerInterceptor{} + interceptors1 := []grpc.UnaryServerInterceptor{salt} + interceptors2 := []grpc.UnaryServerInterceptor{salt, pepper} + // interceptors5 := []grpc.UnaryServerInterceptor{interceptor, interceptor, interceptor, interceptor, interceptor} + + type args struct { + interceptors []grpc.UnaryServerInterceptor + } + tests := []struct { + name string + interceptors []grpc.UnaryServerInterceptor + want string + }{ + {"0", interceptors0, "Soup"}, + {"1", interceptors1, "Soup, add salt"}, + {"2", interceptors2, "Soup, add salt, add pepper"}, + } + for _, tt := range tests { + ctx := context.Background() + req := &pb.QueryPQLRequest{ + Pql: "Soup", + } + info := &grpc.UnaryServerInfo{} + handler := func(ctx context.Context, req interface{}) (interface{}, error) { + r := req.(*pb.QueryPQLRequest) + return r.Pql, nil + } + t.Run(tt.name, func(t *testing.T) { + r := req + chained := server.ChainUnaryInterceptor(tt.interceptors...) + resp, err := chained(ctx, r, info, handler) + vprint.VV("resp: %+v", resp) + + if err != nil { + t.Errorf("ChainUnaryInterceptor() error = %v", err) + } + if resp != tt.want { + t.Errorf("ChainUnaryInterceptor() = %v, want %v", resp, tt.want) + } + + }) + } +} + +// Testing object that implements the ServerStream interface +type MockStream struct { + context context.Context +} + +func (ms MockStream) SetHeader(md metadata.MD) error { + ms.context = context.WithValue(context.Background(), "metadata", md) + return nil +} + +func (ms MockStream) SendHeader(metadata.MD) error { + return nil +} + +func (ms MockStream) SetTrailer(metadata.MD) {} + +func (ms MockStream) Context() context.Context { + if ms.context == nil { + md := metadata.New(map[string]string{}) + ms.context = context.WithValue(context.Background(), "metadata", md) + } + return ms.context +} + +func (ms MockStream) SendMsg(m interface{}) error { + return nil +} + +func (ms MockStream) RecvMsg(m interface{}) error { + return nil +} + +func fromIncomingContext(ctx context.Context) metadata.MD { + return ctx.Value("metadata").(metadata.MD) +} + +func Test_ChainStreamInterceptor(t *testing.T) { + + salt := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + md := fromIncomingContext(ss.Context()) + md.Append("ingredient", "with salt") + ss.SetHeader(md) + return handler(srv, ss) + } + pepper := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { + md := fromIncomingContext(ss.Context()) + md.Append("ingredient", "and pepper") + ss.SetHeader(md) + return handler(srv, ss) + } + + interceptors0 := []grpc.StreamServerInterceptor{} + interceptors1 := []grpc.StreamServerInterceptor{salt} + interceptors2 := []grpc.StreamServerInterceptor{salt, pepper} + + type args struct { + interceptors []grpc.StreamServerInterceptor + } + tests := []struct { + name string + interceptors []grpc.StreamServerInterceptor + want []string + }{ + {"0", interceptors0, []string{"Soup"}}, + {"1", interceptors1, []string{"Soup", "with salt"}}, + {"2", interceptors2, []string{"Soup", "with salt", "and pepper"}}, + } + for _, tt := range tests { + result := make([]string, 0) + srv := "asdf" + md := metadata.New(map[string]string{}) + ss := MockStream{context: context.WithValue(context.Background(), "metadata", md)} + info := &grpc.StreamServerInfo{} + handler := func(srv interface{}, stream grpc.ServerStream) error { + md := fromIncomingContext(stream.Context()) + vals := md.Get("ingredient") + result = append(result, "Soup") + result = append(result, vals...) + return nil + } + t.Run(tt.name, func(t *testing.T) { + chained := server.ChainStreamInterceptors(tt.interceptors...) + chained(srv, ss, info, handler) + if !reflect.DeepEqual(result, tt.want) { + t.Errorf("ChainStreamInterceptor() = %v, want %v", result, tt.want) + } + }) + } +}