Add metrics to AsyncAPI message handler in dequeuer (#2231)

Author: Miguel Varela Ramos

cmd/dequeuer/main.go

@@ -141,6 +141,11 @@ func main() {
 		exit(log, err, "unable to initialize metrics client")
 	}
 
+	adminHandler := http.NewServeMux()
+	adminHandler.Handle("/healthz", dequeuer.HealthcheckHandler(func() bool {
+		return probe.AreProbesHealthy(probes)
+	}))
+
 	var dequeuerConfig dequeuer.SQSDequeuerConfig
 	var messageHandler dequeuer.MessageHandler
 
@@ -177,24 +182,22 @@ func main() {
 			TargetURL:  targetURL,
 		}
 
-		messageHandler = dequeuer.NewAsyncMessageHandler(config, awsClient, log)
+		asyncStatsReporter := dequeuer.NewAsyncPrometheusStatsReporter()
+		messageHandler = dequeuer.NewAsyncMessageHandler(config, awsClient, asyncStatsReporter, log)
 		dequeuerConfig = dequeuer.SQSDequeuerConfig{
 			Region:           clusterConfig.Region,
 			QueueURL:         queueURL,
 			StopIfNoMessages: false,
 		}
 
+		// report prometheus metrics for async api kinds
+		adminHandler.Handle("/metrics", asyncStatsReporter)
 	default:
 		exit(log, err, fmt.Sprintf("kind %s is not supported", apiKind))
 	}
 
 	errCh := make(chan error)
 
-	adminHandler := http.NewServeMux()
-	adminHandler.Handle("/healthz", dequeuer.HealthcheckHandler(func() bool {
-		return probe.AreProbesHealthy(probes)
-	}))
-
 	go func() {
 		server := &http.Server{
 			Addr:    ":" + strconv.Itoa(adminPort),
@@ -219,13 +222,17 @@ func main() {
 		})
 	}()
 
-	for _, probe := range probes {
-		stopper := probe.StartProbing()
-		defer func() {
-			stopper <- struct{}{}
-		}()
+	var stopChs []chan struct{}
+	for _, p := range probes {
+		stopChs = append(stopChs, p.StartProbing())
 	}
 
+	defer func() {
+		for _, stopCh := range stopChs {
+			stopCh <- struct{}{}
+		}
+	}()
+
 	select {
 	case err = <-errCh:
 		exit(log, err, "error during message dequeueing or error from admin server")

go.mod

@@ -30,7 +30,7 @@ require (
 	github.com/morikuni/aec v1.0.0 // indirect
 	github.com/onsi/ginkgo v1.14.1
 	github.com/onsi/gomega v1.10.3
-	github.com/opencontainers/go-digest v1.0.0
+	github.com/opencontainers/go-digest v1.0.0 // indirect
 	github.com/ory/dockertest/v3 v3.6.5
 	github.com/patrickmn/go-cache v2.1.0+incompatible
 	github.com/pkg/errors v0.9.1

go.sum

@@ -547,7 +547,6 @@ github.com/onsi/gomega v0.0.0-20170829124025-dcabb60a477c/go.mod h1:C1qb7wdrVGGV
 github.com/onsi/gomega v1.7.0/go.mod h1:ex+gbHU/CVuBBDIJjb2X0qEXbFg53c61hWP/1CpauHY=
 github.com/onsi/gomega v1.7.1/go.mod h1:XdKZgCCFLUoM/7CFJVPcG8C1xQ1AJ0vpAezJrB7JYyY=
 github.com/onsi/gomega v1.10.1/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
-github.com/onsi/gomega v1.10.2 h1:aY/nuoWlKJud2J6U0E3NWsjlg+0GtwXxgEqthRdzlcs=
 github.com/onsi/gomega v1.10.2/go.mod h1:iN09h71vgCQne3DLsj+A5owkum+a2tYe+TOCB1ybHNo=
 github.com/onsi/gomega v1.10.3 h1:gph6h/qe9GSUw1NhH1gp+qb+h8rXD8Cy60Z32Qw3ELA=
 github.com/onsi/gomega v1.10.3/go.mod h1:V9xEwhxec5O8UDM77eCW8vLymOMltsqPVYWrpDsH8xc=

pkg/dequeuer/async_handler.go

@@ -22,6 +22,7 @@ import (
 	"io"
 	"net/http"
 	"strings"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/s3"
@@ -39,11 +40,12 @@ const (
 )
 
 type AsyncMessageHandler struct {
-	config      AsyncMessageHandlerConfig
-	aws         *awslib.Client
-	log         *zap.SugaredLogger
-	storagePath string
-	httpClient  *http.Client
+	config       AsyncMessageHandlerConfig
+	aws          *awslib.Client
+	log          *zap.SugaredLogger
+	storagePath  string
+	httpClient   *http.Client
+	eventHandler RequestEventHandler
 }
 
 type AsyncMessageHandlerConfig struct {
@@ -58,13 +60,14 @@ type userPayload struct {
 	ContentType string
 }
 
-func NewAsyncMessageHandler(config AsyncMessageHandlerConfig, awsClient *awslib.Client, logger *zap.SugaredLogger) *AsyncMessageHandler {
+func NewAsyncMessageHandler(config AsyncMessageHandlerConfig, awsClient *awslib.Client, eventHandler RequestEventHandler, logger *zap.SugaredLogger) *AsyncMessageHandler {
 	return &AsyncMessageHandler{
-		config:      config,
-		aws:         awsClient,
-		log:         logger,
-		storagePath: async.StoragePath(config.ClusterUID, config.APIName),
-		httpClient:  &http.Client{},
+		config:       config,
+		aws:          awsClient,
+		log:          logger,
+		storagePath:  async.StoragePath(config.ClusterUID, config.APIName),
+		httpClient:   &http.Client{},
+		eventHandler: eventHandler,
 	}
 }
 
@@ -175,11 +178,18 @@ func (h *AsyncMessageHandler) submitRequest(payload *userPayload, requestID stri
 
 	req.Header.Set("Content-Type", payload.ContentType)
 	req.Header.Set(CortexRequestIDHeader, requestID)
+
+	startTime := time.Now()
 	response, err := h.httpClient.Do(req)
 	if err != nil {
 		return nil, ErrorUserContainerNotReachable(err)
 	}
 
+	requestEvent := RequestEvent{
+		StatusCode: response.StatusCode,
+		Duration:   time.Since(startTime),
+	}
+
 	defer func() {
 		_ = response.Body.Close()
 	}()
@@ -197,6 +207,8 @@ func (h *AsyncMessageHandler) submitRequest(payload *userPayload, requestID stri
 		return nil, ErrorUserContainerResponseNotJSONDecodable()
 	}
 
+	h.eventHandler.HandleEvent(requestEvent)
+
 	return result, nil
 }
 

pkg/dequeuer/async_handler_test.go

@@ -49,12 +49,17 @@ func TestAsyncMessageHandler_Handle(t *testing.T) {
 		_, _ = w.Write([]byte("{}"))
 	}))
 
+	var requestEventsCount int
+	eventHandler := NewRequestEventHandlerFunc(func(event RequestEvent) {
+		requestEventsCount++
+	})
+
 	asyncHandler := NewAsyncMessageHandler(AsyncMessageHandlerConfig{
 		ClusterUID: "cortex-test",
 		Bucket:     _testBucket,
 		APIName:    "async-test",
 		TargetURL:  server.URL,
-	}, awsClient, log)
+	}, awsClient, eventHandler, log)
 
 	_, err := awsClient.S3().CreateBucket(&s3.CreateBucketInput{
 		Bucket: aws.String(_testBucket),
@@ -75,6 +80,7 @@ func TestAsyncMessageHandler_Handle(t *testing.T) {
 		fmt.Sprintf("%s/%s/status/%s", asyncHandler.storagePath, requestID, async.StatusCompleted),
 	)
 	require.NoError(t, err)
+	require.Equal(t, 1, requestEventsCount)
 }
 
 func TestAsyncMessageHandler_Handle_Errors(t *testing.T) {
@@ -105,12 +111,14 @@ func TestAsyncMessageHandler_Handle_Errors(t *testing.T) {
 	log := newLogger(t)
 	awsClient := testAWSClient(t)
 
+	eventHandler := NewRequestEventHandlerFunc(func(event RequestEvent) {})
+
 	asyncHandler := NewAsyncMessageHandler(AsyncMessageHandlerConfig{
 		ClusterUID: "cortex-test",
 		Bucket:     _testBucket,
 		APIName:    "async-test",
 		TargetURL:  "http://fake.cortex.dev",
-	}, awsClient, log)
+	}, awsClient, eventHandler, log)
 
 	for _, tt := range cases {
 		t.Run(tt.name, func(t *testing.T) {

pkg/dequeuer/async_stats.go

@@ -0,0 +1,65 @@
+/*
+Copyright 2021 Cortex Labs, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package dequeuer
+
+import (
+	"net/http"
+	"strconv"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/prometheus/client_golang/prometheus/promhttp"
+)
+
+type AsyncStatsReporter struct {
+	handler      http.Handler
+	latencies    *prometheus.HistogramVec
+	requestCount *prometheus.CounterVec
+}
+
+func NewAsyncPrometheusStatsReporter() *AsyncStatsReporter {
+	latenciesHist := promauto.NewHistogramVec(prometheus.HistogramOpts{
+		Name: "cortex_async_latency",
+		Help: "Histogram of the latencies for an AsyncAPI kind in seconds",
+	}, []string{"status_code"})
+
+	requestCounter := promauto.NewCounterVec(prometheus.CounterOpts{
+		Name: "cortex_async_request_count",
+		Help: "Request count for an AsyncAPI",
+	}, []string{"status_code"})
+
+	handler := promhttp.Handler()
+
+	return &AsyncStatsReporter{
+		handler:      handler,
+		latencies:    latenciesHist,
+		requestCount: requestCounter,
+	}
+}
+
+func (r *AsyncStatsReporter) HandleEvent(event RequestEvent) {
+	labels := map[string]string{
+		"status_code": strconv.Itoa(event.StatusCode),
+	}
+
+	r.latencies.With(labels).Observe(event.Duration.Seconds())
+	r.requestCount.With(labels).Add(1)
+}
+
+func (r *AsyncStatsReporter) ServeHTTP(w http.ResponseWriter, req *http.Request) {
+	r.handler.ServeHTTP(w, req)
+}

pkg/dequeuer/async_stats_test.go

@@ -0,0 +1,100 @@
+/*
+Copyright 2021 Cortex Labs, Inc.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package dequeuer
+
+import (
+	"net/http/httptest"
+	"strings"
+	"testing"
+	"time"
+
+	"github.com/prometheus/client_golang/prometheus"
+	"github.com/prometheus/client_golang/prometheus/promauto"
+	"github.com/prometheus/client_golang/prometheus/testutil"
+	"github.com/stretchr/testify/require"
+)
+
+func TestNewAsyncPrometheusStatsReporter(t *testing.T) {
+	t.Parallel()
+
+	statsReporter := NewAsyncPrometheusStatsReporter()
+
+	statsReporter.HandleEvent(
+		RequestEvent{
+			StatusCode: 200,
+			Duration:   100 * time.Millisecond,
+		},
+	)
+
+	w := httptest.NewRecorder()
+	r := httptest.NewRequest("GET", "/metrics", nil)
+	statsReporter.ServeHTTP(w, r)
+
+	result := w.Body.String()
+	require.Contains(t, result, "cortex_async_latency")
+	require.Contains(t, result, "cortex_async_request_count")
+}
+
+func TestAsyncStatsReporter_HandleEvent(t *testing.T) {
+	t.Parallel()
+
+	reg := prometheus.NewRegistry()
+
+	latenciesHist := promauto.With(reg).NewHistogramVec(prometheus.HistogramOpts{
+		Name: "cortex_async_latency",
+		Help: "Histogram of the latencies for an AsyncAPI kind in seconds",
+	}, []string{"status_code"})
+
+	requestCounter := promauto.With(reg).NewCounterVec(prometheus.CounterOpts{
+		Name: "cortex_async_request_count",
+		Help: "Request count for an AsyncAPI",
+	}, []string{"status_code"})
+
+	statsReporter := AsyncStatsReporter{
+		latencies:    latenciesHist,
+		requestCount: requestCounter,
+	}
+
+	statsReporter.HandleEvent(
+		RequestEvent{
+			StatusCode: 200,
+			Duration:   100 * time.Millisecond,
+		},
+	)
+
+	expectedHist := `
+# HELP cortex_async_latency Histogram of the latencies for an AsyncAPI kind in seconds
+# TYPE cortex_async_latency histogram
+cortex_async_latency_bucket{status_code="200",le="0.005"} 0
+cortex_async_latency_bucket{status_code="200",le="0.01"} 0
+cortex_async_latency_bucket{status_code="200",le="0.025"} 0
+cortex_async_latency_bucket{status_code="200",le="0.05"} 0
+cortex_async_latency_bucket{status_code="200",le="0.1"} 1
+cortex_async_latency_bucket{status_code="200",le="0.25"} 1
+cortex_async_latency_bucket{status_code="200",le="0.5"} 1
+cortex_async_latency_bucket{status_code="200",le="1"} 1
+cortex_async_latency_bucket{status_code="200",le="2.5"} 1
+cortex_async_latency_bucket{status_code="200",le="5"} 1
+cortex_async_latency_bucket{status_code="200",le="10"} 1
+cortex_async_latency_bucket{status_code="200",le="+Inf"} 1
+cortex_async_latency_sum{status_code="200"} 0.1
+cortex_async_latency_count{status_code="200"} 1
+`
+
+	require.Equal(t, float64(1), testutil.ToFloat64(statsReporter.requestCount))
+	require.NoError(t, testutil.CollectAndCompare(statsReporter.latencies, strings.NewReader(expectedHist)))
+}

pkg/dequeuer/batch_handler.go

@@ -128,7 +128,9 @@ func (h *BatchMessageHandler) submitRequest(messageBody string, isOnJobComplete
 	if err != nil {
 		return ErrorUserContainerNotReachable(err)
 	}
-	defer response.Body.Close()
+	defer func() {
+		_ = response.Body.Close()
+	}()
 
 	if response.StatusCode == http.StatusNotFound && isOnJobComplete {
 		return nil