add rpc panic recovery middleware

modules/rpc/handler.go

@@ -27,9 +27,13 @@ import (
 	"go.uber.org/fx/modules/rpc/internal/stats"
 	"go.uber.org/fx/service"
 	"go.uber.org/fx/ulog"
+
+	"github.com/pkg/errors"
 	"go.uber.org/yarpc/api/transport"
 )
 
+const _panicResponse = "Server Error"
+
 type contextInboundMiddleware struct {
 	service.Host
 }
@@ -88,3 +92,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)
+		ulog.Logger(ctx).Error("Panic recovered serving request", "error", errors.Errorf("panic in handler: %+v", err))
+		// rethrow panic back to yarpc
+		// before https://github.com/yarpc/yarpc-go/issues/734 fixed, throw a generic error.
+		panic(_panicResponse)
+	}
+}

modules/rpc/handler_test.go

@@ -31,6 +31,7 @@ import (
 	"go.uber.org/yarpc/api/transport"
 
 	"github.com/stretchr/testify/assert"
+	"github.com/uber-go/tally"
 )
 
 type fakeEnveloper struct {
@@ -54,31 +55,31 @@ func TestInboundMiddleware_fxContext(t *testing.T) {
 		Host: service.NopHost(),
 	}
 	stats.SetupRPCMetrics(unary.Host.Metrics())
-	err := unary.Handle(context.Background(), &transport.Request{}, nil, &fakeUnaryHandler{t: t})
+	err := unary.Handle(context.Background(), &transport.Request{}, nil, &fakeUnary{t: t})
 	assert.Equal(t, "handle", err.Error())
 }
 
 func TestOnewayInboundMiddleware_fxContext(t *testing.T) {
 	oneway := contextOnewayInboundMiddleware{
 		Host: service.NopHost(),
 	}
-	err := oneway.HandleOneway(context.Background(), &transport.Request{}, &fakeOnewayHandler{t: t})
+	err := oneway.HandleOneway(context.Background(), &transport.Request{}, &fakeOneway{t: t})
 	assert.Equal(t, "oneway handle", err.Error())
 }
 
 func TestInboundMiddleware_auth(t *testing.T) {
 	unary := authInboundMiddleware{
 		Host: service.NopHost(),
 	}
-	err := unary.Handle(context.Background(), &transport.Request{}, nil, &fakeUnaryHandler{t: t})
+	err := unary.Handle(context.Background(), &transport.Request{}, nil, &fakeUnary{t: t})
 	assert.EqualError(t, err, "handle")
 }
 
 func TestInboundMiddleware_authFailure(t *testing.T) {
 	unary := authInboundMiddleware{
 		Host: service.NopHostAuthFailure(),
 	}
-	err := unary.Handle(context.Background(), &transport.Request{}, nil, &fakeUnaryHandler{t: t})
+	err := unary.Handle(context.Background(), &transport.Request{}, nil, &fakeUnary{t: t})
 	assert.EqualError(t, err, "Error authorizing the service")
 
 }
@@ -87,32 +88,73 @@ func TestOnewayInboundMiddleware_auth(t *testing.T) {
 	oneway := authOnewayInboundMiddleware{
 		Host: service.NopHost(),
 	}
-	err := oneway.HandleOneway(context.Background(), &transport.Request{}, &fakeOnewayHandler{t: t})
+	err := oneway.HandleOneway(context.Background(), &transport.Request{}, &fakeOneway{t: t})
 	assert.EqualError(t, err, "oneway handle")
 }
 
 func TestOnewayInboundMiddleware_authFailure(t *testing.T) {
 	oneway := authOnewayInboundMiddleware{
 		Host: service.NopHostAuthFailure(),
 	}
-	err := oneway.HandleOneway(context.Background(), &transport.Request{}, &fakeOnewayHandler{t: t})
+	err := oneway.HandleOneway(context.Background(), &transport.Request{}, &fakeOneway{t: t})
 	assert.EqualError(t, err, "Error authorizing the service")
 }
 
-type fakeUnaryHandler struct {
+func TestInboundMiddleware_panic(t *testing.T) {
+	host := service.NopHost()
+	testScope := host.Metrics()
+	stats.SetupRPCMetrics(testScope)
+
+	defer testPanicHandler(t, testScope)
+	unary := panicInboundMiddleware{}
+	unary.Handle(context.Background(), &transport.Request{}, nil, &alwaysPanicUnary{})
+}
+
+func TestOnewayInboundMiddleware_panic(t *testing.T) {
+	host := service.NopHost()
+	testScope := host.Metrics()
+	stats.SetupRPCMetrics(testScope)
+
+	defer testPanicHandler(t, testScope)
+	oneway := panicOnewayInboundMiddleware{}
+	oneway.HandleOneway(context.Background(), &transport.Request{}, &alwaysPanicOneway{})
+}
+
+func testPanicHandler(t *testing.T, testScope tally.Scope) {
+	r := recover()
+	assert.EqualValues(t, r, _panicResponse)
+
+	snapshot := testScope.(tally.TestScope).Snapshot()
+	counters := snapshot.Counters()
+	assert.True(t, counters["panic"].Value() > 0)
+}
+
+type fakeUnary struct {
 	t *testing.T
 }
 
-func (f fakeUnaryHandler) Handle(ctx context.Context, _param1 *transport.Request, _param2 transport.ResponseWriter) error {
+func (f fakeUnary) Handle(ctx context.Context, _param1 *transport.Request, _param2 transport.ResponseWriter) error {
 	assert.NotNil(f.t, ctx)
 	return errors.New("handle")
 }
 
-type fakeOnewayHandler struct {
+type fakeOneway struct {
 	t *testing.T
 }
 
-func (f fakeOnewayHandler) HandleOneway(ctx context.Context, p *transport.Request) error {
+func (f fakeOneway) HandleOneway(ctx context.Context, p *transport.Request) error {
 	assert.NotNil(f.t, ctx)
 	return errors.New("oneway handle")
 }
+
+type alwaysPanicUnary struct{}
+
+func (p alwaysPanicUnary) Handle(_ context.Context, _ *transport.Request, _ transport.ResponseWriter) error {
+	panic("panic")
+}
+
+type alwaysPanicOneway struct{}
+
+func (p alwaysPanicOneway) 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},
 		},
 	}