Add Metrics() to HTTP an GRPC proxies (#274)

* agent/protocol: add Metrics API to proxy
* agent/http: add support for metrics
* agent/grpc: add metrics to grpc handler
* agent/grpc: slightly refactor handler to reduce nesting

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,27 +55,35 @@ 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")
 	}
+
 	// full method name has the form /service/method, we want the service
 	serviceName := strings.Split(fullMethodName, "/")[1]
-	excluded := contains(h.disruption.Excluded, serviceName)
-	if !excluded {
-		if h.disruption.ErrorRate > 0 && rand.Float32() <= h.disruption.ErrorRate {
-			return h.injectError(serverStream)
-		}
+	if contains(h.disruption.Excluded, serviceName) {
+		h.metrics.Inc(protocol.MetricRequestsExcluded)
+		return h.transparentForward(serverStream)
+	}
 
-		// add delay
-		if h.disruption.AverageDelay > 0 {
-			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))
+	if rand.Float32() < h.disruption.ErrorRate {
+		h.metrics.Inc(protocol.MetricRequestsDisrupted)
+		return h.injectError(serverStream)
+	}
+
+	// add delay
+	if h.disruption.AverageDelay > 0 {
+		h.metrics.Inc(protocol.MetricRequestsDisrupted)
+
+		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))
 	}
 
 	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),
@@ -120,6 +122,12 @@ func (p *proxy) Stop() error {
 	return nil
 }
 
+// Metrics returns runtime metrics for the proxy.
+// TODO: Add metrics.
+func (p *proxy) Metrics() map[string]uint {
+	return p.metrics.Map()
+}
+
 // Force stops the proxy without waiting for connections to drain
 // In grpc this action is a nop
 func (p *proxy) Force() error {

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.MetricRequestsDisrupted: 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)
+			}
+		})
+	}
+}

pkg/agent/protocol/http/proxy.go

@@ -44,6 +44,7 @@ type proxy struct {
 	config     ProxyConfig
 	disruption Disruption
 	srv        *http.Server
+	metrics    protocol.MetricMap
 }
 
 // NewProxy return a new Proxy for HTTP requests
@@ -85,6 +86,7 @@ type httpHandler struct {
 	upstreamURL url.URL
 	disruption  Disruption
 	client      httpClient
+	metrics     *protocol.MetricMap
 }
 
 // isExcluded checks whether a request should be proxied through without any kind of modification whatsoever.
@@ -146,7 +148,10 @@ func (h *httpHandler) injectError(rw http.ResponseWriter, delay time.Duration) {
 }
 
 func (h *httpHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
+	h.metrics.Inc(protocol.MetricRequests)
+
 	if h.isExcluded(req) {
+		h.metrics.Inc(protocol.MetricRequestsExcluded)
 		//nolint:contextcheck // Unclear which context the linter requires us to propagate here.
 		h.forward(rw, req, 0)
 		return
@@ -159,6 +164,7 @@ func (h *httpHandler) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
 	}
 
 	if h.disruption.ErrorRate > 0 && rand.Float32() <= h.disruption.ErrorRate {
+		h.metrics.Inc(protocol.MetricRequestsDisrupted)
 		h.injectError(rw, delay)
 		return
 	}
@@ -178,6 +184,7 @@ func (p *proxy) Start() error {
 		upstreamURL: *upstreamURL,
 		disruption:  p.disruption,
 		client:      http.DefaultClient,
+		metrics:     &p.metrics,
 	}
 
 	p.srv = &http.Server{
@@ -200,6 +207,11 @@ func (p *proxy) Stop() error {
 	return nil
 }
 
+// Metrics returns runtime metrics for the proxy.
+func (p *proxy) Metrics() map[string]uint {
+	return p.metrics.Map()
+}
+
 // Force stops the proxy without waiting for connections to drain
 func (p *proxy) Force() error {
 	if p.srv != nil {