Recover querier handler from panic. (#10983)

grafana · Oct 31, 2023 · 4dc9d3b · 4dc9d3b
1 parent bcab696
commit 4dc9d3b
Show file tree

Hide file tree

Showing 4 changed files with 28 additions and 4 deletions.
diff --git a/pkg/loki/modules.go b/pkg/loki/modules.go
@@ -505,10 +505,10 @@ func (t *Loki) initQuerier() (services.Service, error) {
 	}
 
 	internalHandler := queryrangebase.MergeMiddlewares(
+		serverutil.RecoveryMiddleware,
 		queryrange.Instrument{
 			QueryHandlerMetrics: queryrange.NewQueryHandlerMetrics(
 				prometheus.DefaultRegisterer,
-				t.Cfg.MetricsNamespace,
 			),
 		},
 	).Wrap(handler)

diff --git a/pkg/querier/queryrange/instrument.go b/pkg/querier/queryrange/instrument.go
@@ -18,11 +18,11 @@ type QueryHandlerMetrics struct {
 	InflightRequests *prometheus.GaugeVec
 }
 
-func NewQueryHandlerMetrics(registerer prometheus.Registerer, metricsNamespace string) *QueryHandlerMetrics {
+func NewQueryHandlerMetrics(registerer prometheus.Registerer) *QueryHandlerMetrics {
 	return &QueryHandlerMetrics{
 		InflightRequests: promauto.With(registerer).NewGaugeVec(prometheus.GaugeOpts{
-			Namespace: metricsNamespace,
-			Name:      "inflight_requests",
+			Namespace: "loki",
+			Name:      "inflight_requests_grpc",
 			Help:      "Current number of inflight requests.",
 		}, []string{"method", "route"}),
 	}

diff --git a/pkg/util/server/recovery.go b/pkg/util/server/recovery.go
@@ -1,6 +1,7 @@
 package server
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"os"
@@ -11,6 +12,8 @@ import (
 	grpc_recovery "github.com/grpc-ecosystem/go-grpc-middleware/recovery"
 	"github.com/prometheus/client_golang/prometheus"
 	"github.com/prometheus/client_golang/prometheus/promauto"
+
+	"github.com/grafana/loki/pkg/querier/queryrange/queryrangebase"
 )
 
 const maxStacksize = 8 * 1024
@@ -34,6 +37,18 @@ var (
 	})
 	RecoveryGRPCStreamInterceptor = grpc_recovery.StreamServerInterceptor(grpc_recovery.WithRecoveryHandler(onPanic))
 	RecoveryGRPCUnaryInterceptor  = grpc_recovery.UnaryServerInterceptor(grpc_recovery.WithRecoveryHandler(onPanic))
+
+	RecoveryMiddleware queryrangebase.Middleware = queryrangebase.MiddlewareFunc(func(next queryrangebase.Handler) queryrangebase.Handler {
+		return queryrangebase.HandlerFunc(func(ctx context.Context, req queryrangebase.Request) (res queryrangebase.Response, err error) {
+			defer func() {
+				if p := recover(); p != nil {
+					err = onPanic(p)
+				}
+			}()
+			res, err = next.Do(ctx, req)
+			return
+		})
+	})
 )
 
 func onPanic(p interface{}) error {

diff --git a/pkg/util/server/recovery_test.go b/pkg/util/server/recovery_test.go
@@ -9,6 +9,8 @@ import (
 	"github.com/stretchr/testify/require"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/metadata"
+
+	"github.com/grafana/loki/pkg/querier/queryrange/queryrangebase"
 )
 
 func Test_onPanic(t *testing.T) {
@@ -32,6 +34,13 @@ func Test_onPanic(t *testing.T) {
 		panic("foo")
 	}))
 	require.Error(t, err)
+
+	_, err = RecoveryMiddleware.
+		Wrap(queryrangebase.HandlerFunc(func(ctx context.Context, req queryrangebase.Request) (res queryrangebase.Response, err error) {
+			panic("foo")
+		})).
+		Do(context.Background(), nil)
+	require.ErrorContains(t, err, "foo")
 }
 
 type fakeStream struct{}