scriptoriumadmin/featurebase
1
0
Fork 0
mirror of https://github.com/featurebasedb/featurebase.git synced 2026-08-28 10:54:59 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[FB-1479] Adds GRPC interceptor for sentry and chaining mechanism

Browse source 
also adds metadata for sentry performance monitoring for more organized output
This commit is contained in:
Samir Patel 2022-06-15 09:26:58 -05:00 • committed by GitHub
parent b7516eacc5
commit 0d43be934e
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 259 additions and 34 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

14
http_handler.go
View file

@ -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

129
server/grpc.go
View file

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

150
server/grpc_test.go
View file

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