agent/grpc: add metrics to grpc handler

Author: nadiamoe

pkg/agent/protocol/grpc/handler.go

@@ -9,6 +9,7 @@ import (
 	"strings"
 	"time"
 
+	"github.com/grafana/xk6-disruptor/pkg/agent/protocol"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
 	"google.golang.org/grpc/metadata"
@@ -24,10 +25,11 @@ func clientStreamDescForProxy() *grpc.StreamDesc {
 }
 
 // NewHandler returns a StreamHandler that attempts to proxy all requests that are not registered in the server.
-func NewHandler(disruption Disruption, forwardConn *grpc.ClientConn) grpc.StreamHandler {
+func NewHandler(disruption Disruption, forwardConn *grpc.ClientConn, metrics *protocol.MetricMap) grpc.StreamHandler {
 	handler := &handler{
 		disruption:  disruption,
 		forwardConn: forwardConn,
+		metrics:     metrics,
 	}
 
 	// return the handler function
@@ -37,6 +39,7 @@ func NewHandler(disruption Disruption, forwardConn *grpc.ClientConn) grpc.Stream
 type handler struct {
 	disruption  Disruption
 	forwardConn *grpc.ClientConn
+	metrics     *protocol.MetricMap
 }
 
 // contains verifies if a list of strings contains the given string
@@ -52,6 +55,8 @@ func contains(list []string, target string) bool {
 // handles requests from the client. If selected for error injection, returns an error,
 // otherwise, forwards to the server transparently
 func (h *handler) streamHandler(_ interface{}, serverStream grpc.ServerStream) error {
+	h.metrics.Inc(protocol.MetricRequests)
+
 	fullMethodName, ok := grpc.MethodFromServerStream(serverStream)
 	if !ok {
 		return status.Errorf(codes.Internal, "ServerTransportStream not exists in context")
@@ -61,18 +66,23 @@ func (h *handler) streamHandler(_ interface{}, serverStream grpc.ServerStream) e
 	excluded := contains(h.disruption.Excluded, serviceName)
 	if !excluded {
 		if h.disruption.ErrorRate > 0 && rand.Float32() <= h.disruption.ErrorRate {
+			h.metrics.Inc(protocol.MetricRequestsFaulted)
 			return h.injectError(serverStream)
 		}
 
 		// add delay
 		if h.disruption.AverageDelay > 0 {
+			h.metrics.Inc(protocol.MetricRequestsFaulted)
+
 			delay := int64(h.disruption.AverageDelay)
 			if h.disruption.DelayVariation > 0 {
 				variation := int64(h.disruption.DelayVariation)
 				delay = delay + variation - 2*rand.Int63n(variation)
 			}
 			time.Sleep(time.Duration(delay))
 		}
+	} else {
+		h.metrics.Inc(protocol.MetricRequestsExcluded)
 	}
 
 	return h.transparentForward(serverStream)

pkg/agent/protocol/grpc/proxy.go

@@ -44,6 +44,7 @@ type proxy struct {
 	config     ProxyConfig
 	disruption Disruption
 	srv        *grpc.Server
+	metrics    *protocol.MetricMap
 }
 
 // NewProxy return a new Proxy
@@ -79,6 +80,7 @@ func NewProxy(c ProxyConfig, d Disruption) (protocol.Proxy, error) {
 	return &proxy{
 		disruption: d,
 		config:     c,
+		metrics:    &protocol.MetricMap{},
 	}, nil
 }
 
@@ -93,7 +95,7 @@ func (p *proxy) Start() error {
 	if err != nil {
 		return fmt.Errorf("error dialing %s: %w", p.config.UpstreamAddress, err)
 	}
-	handler := NewHandler(p.disruption, conn)
+	handler := NewHandler(p.disruption, conn, p.metrics)
 
 	p.srv = grpc.NewServer(
 		grpc.UnknownServiceHandler(handler),
@@ -123,7 +125,7 @@ func (p *proxy) Stop() error {
 // Metrics returns runtime metrics for the proxy.
 // TODO: Add metrics.
 func (p *proxy) Metrics() map[string]uint {
-	return nil
+	return p.metrics.Map()
 }
 
 // Force stops the proxy without waiting for connections to drain

pkg/agent/protocol/grpc/proxy_test.go

@@ -8,7 +8,9 @@ import (
 	"path/filepath"
 	"testing"
 
+	"github.com/google/go-cmp/cmp"
 	"github.com/google/uuid"
+	"github.com/grafana/xk6-disruptor/pkg/agent/protocol"
 	"github.com/grafana/xk6-disruptor/pkg/testutils/grpc/ping"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/codes"
@@ -187,6 +189,7 @@ func Test_ProxyHandler(t *testing.T) {
 		expectStatus codes.Code
 	}
 
+	// TODO: Add test for excluded endpoints
 	testCases := []TestCase{
 		{
 			title: "default proxy",
@@ -207,7 +210,7 @@ func Test_ProxyHandler(t *testing.T) {
 			expectStatus: codes.OK,
 		},
 		{
-			title: "error injection ",
+			title: "error injection",
 			disruption: Disruption{
 				AverageDelay:   0,
 				DelayVariation: 0,
@@ -331,3 +334,125 @@ func Test_ProxyHandler(t *testing.T) {
 		})
 	}
 }
+
+func Test_ProxyMetrics(t *testing.T) {
+	t.Parallel()
+
+	type TestCase struct {
+		title           string
+		disruption      Disruption
+		expectedMetrics map[string]uint
+	}
+
+	// TODO: Add test for excluded endpoints
+	testCases := []TestCase{
+		{
+			title: "passthrough",
+			disruption: Disruption{
+				AverageDelay:   0,
+				DelayVariation: 0,
+				ErrorRate:      0.0,
+				StatusCode:     0,
+				StatusMessage:  "",
+			},
+			expectedMetrics: map[string]uint{
+				protocol.MetricRequests: 1,
+			},
+		},
+		{
+			title: "error injection",
+			disruption: Disruption{
+				AverageDelay:   0,
+				DelayVariation: 0,
+				ErrorRate:      1.0,
+				StatusCode:     int32(codes.Internal),
+				StatusMessage:  "Internal server error",
+			},
+			expectedMetrics: map[string]uint{
+				protocol.MetricRequests:        1,
+				protocol.MetricRequestsFaulted: 1,
+			},
+		},
+	}
+
+	for _, tc := range testCases {
+		tc := tc
+
+		t.Run(tc.title, func(t *testing.T) {
+			t.Parallel()
+
+			// start test server in a random unix socket
+			serverSocket := filepath.Join(os.TempDir(), uuid.New().String())
+			l, err := net.Listen("unix", serverSocket)
+			if err != nil {
+				t.Errorf("error starting test server in unix:%s: %v", serverSocket, err)
+				return
+			}
+
+			srv := grpc.NewServer()
+			ping.RegisterPingServiceServer(srv, ping.NewPingServer())
+			go func() {
+				if serr := srv.Serve(l); err != nil {
+					t.Logf("error in the server: %v", serr)
+				}
+			}()
+
+			// start proxy in a random unix socket
+			proxySocket := filepath.Join(os.TempDir(), uuid.New().String())
+			config := ProxyConfig{
+				Network:         "unix",
+				ListenAddress:   proxySocket,
+				UpstreamAddress: fmt.Sprintf("unix:%s", serverSocket),
+			}
+
+			proxy, err := NewProxy(config, tc.disruption)
+			if err != nil {
+				t.Errorf("error creating proxy: %v", err)
+				return
+			}
+
+			defer func() {
+				_ = proxy.Stop()
+			}()
+
+			go func() {
+				if perr := proxy.Start(); perr != nil {
+					t.Logf("error starting proxy: %v", perr)
+				}
+			}()
+
+			// connect client to proxy
+			conn, err := grpc.DialContext(
+				context.TODO(),
+				fmt.Sprintf("unix:%s", proxySocket),
+				grpc.WithInsecure(),
+			)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			defer func() {
+				_ = conn.Close()
+			}()
+
+			client := ping.NewPingServiceClient(conn)
+
+			var headers metadata.MD
+			_, _ = client.Ping(
+				context.TODO(),
+				&ping.PingRequest{
+					Error:   0,
+					Message: "ping",
+				},
+				grpc.Header(&headers),
+				grpc.WaitForReady(true),
+			)
+
+			metrics := proxy.Metrics()
+
+			if diff := cmp.Diff(tc.expectedMetrics, metrics); diff != "" {
+				t.Fatalf("expected metrics do not match returned:\n%s", diff)
+			}
+		})
+	}
+}