add a rpc panic recovery middleware

modules/rpc/handler.go

@@ -30,6 +30,8 @@ import (
 	"go.uber.org/yarpc/api/transport"
 )
 
+const panicRespnose = "Server Error"
+
 type contextInboundMiddleware struct {
 	service.Host
 }
@@ -88,3 +90,27 @@ func authorize(ctx context.Context, host service.Host) (context.Context, error)
 	}
 	return ctx, nil
 }
+
+type panicInboundMiddleware struct{}
+
+func (p panicInboundMiddleware) Handle(ctx context.Context, req *transport.Request, resw transport.ResponseWriter, handler transport.UnaryHandler) error {
+	defer panicRecovery(ctx)
+	return handler.Handle(ctx, req, resw)
+}
+
+type panicOnewayInboundMiddleware struct{}
+
+func (p panicOnewayInboundMiddleware) HandleOneway(ctx context.Context, req *transport.Request, handler transport.OnewayHandler) error {
+	defer panicRecovery(ctx)
+	return handler.HandleOneway(ctx, req)
+}
+
+func panicRecovery(ctx context.Context) {
+	if err := recover(); err != nil {
+		stats.RPCPanicCounter.Inc(1)
+		fx.Logger(ctx).Error("Panic recovered serving request", "error", err)
+		// rethrow panic back to yarpc
+		// before https://github.com/yarpc/yarpc-go/issues/734 fixed, throw a generic error.
+		panic(panicRespnose)
+	}
+}

modules/rpc/handler_test.go

@@ -116,3 +116,32 @@ func (f fakeOnewayHandler) HandleOneway(ctx context.Context, p *transport.Reques
 	assert.NotNil(f.t, ctx)
 	return errors.New("oneway handle")
 }
+
+func TestInboundMiddleware_panic(t *testing.T) {
+	defer testPanicHandler(t)
+	unary := panicInboundMiddleware{}
+	unary.Handle(context.Background(), &transport.Request{}, nil, &panicUnaryHandler{})
+}
+
+func TestOnewayInboundMiddleware_panic(t *testing.T) {
+	defer testPanicHandler(t)
+	oneway := panicOnewayInboundMiddleware{}
+	oneway.HandleOneway(context.Background(), &transport.Request{}, &panicOnewayHandler{})
+}
+
+func testPanicHandler(t *testing.T) {
+	r := recover()
+	assert.EqualValues(t, r, panicRespnose)
+}
+
+type panicUnaryHandler struct{}
+
+func (p panicUnaryHandler) Handle(_ context.Context, _ *transport.Request, _ transport.ResponseWriter) error {
+	panic("panic")
+}
+
+type panicOnewayHandler struct{}
+
+func (p panicOnewayHandler) HandleOneway(_ context.Context, _ *transport.Request) error {
+	panic("panic")
+}

modules/rpc/internal/stats/metrics.go

@@ -46,11 +46,14 @@ var (
 	RPCAuthFailCounter tally.Counter
 	// RPCHandleTimer is a turnaround time for rpc handler
 	RPCHandleTimer tally.Scope
+	// RPCPanicCounter counts panics occurred for rpc handler
+	RPCPanicCounter tally.Counter
 )
 
 // SetupRPCMetrics allocates counters for necessary setup
 func SetupRPCMetrics(scope tally.Scope) {
 	rpcTagsScope := scope.Tagged(RPCTags)
 	RPCAuthFailCounter = rpcTagsScope.Tagged(map[string]string{TagMiddleware: "auth"}).Counter("fail")
 	RPCHandleTimer = rpcTagsScope.Tagged(RPCTags)
+	RPCPanicCounter = rpcTagsScope.Counter("panic")
 }

modules/rpc/yarpc.go

@@ -126,10 +126,12 @@ func (c *dispatcherController) addDefaultMiddleware(host service.Host) {
 	cfg := yarpcConfig{
 		inboundMiddleware: []middleware.UnaryInbound{
 			contextInboundMiddleware{host},
+			panicInboundMiddleware{},
 			authInboundMiddleware{host},
 		},
 		onewayInboundMiddleware: []middleware.OnewayInbound{
 			contextOnewayInboundMiddleware{host},
+			panicOnewayInboundMiddleware{},
 			authOnewayInboundMiddleware{host},
 		},
 	}