From 3351d2331e15ded87b198090b7744356bc0acdca Mon Sep 17 00:00:00 2001
From: Safeer Jiwan <safeer@tbd.email>
Date: Fri, 31 May 2024 16:12:17 -0700
Subject: [PATCH] chore: intercept and log panics in wrap handlers, then
 re-panic (#1602)

Fixes #1587 by creating a new grpc Interceptor to log panics in each
wrap handler, then re-panic.
---
 internal/rpc/context.go | 50 +++++++++++++++++++++++++++++++++++++++--
 1 file changed, 48 insertions(+), 2 deletions(-)

diff --git a/internal/rpc/context.go b/internal/rpc/context.go
index 40f1c09d94..791dcdd0d0 100644
--- a/internal/rpc/context.go
+++ b/internal/rpc/context.go
@@ -4,6 +4,7 @@ import (
 	"context"
 	"fmt"
 	"net/http"
+	"runtime/debug"
 
 	"connectrpc.com/connect"
 	"connectrpc.com/otelconnect"
@@ -79,7 +80,7 @@ func WithRequestName(ctx context.Context, key model.RequestKey) context.Context
 }
 
 func DefaultClientOptions(level log.Level) []connect.ClientOption {
-	interceptors := []connect.Interceptor{MetadataInterceptor(log.Debug), otelInterceptor()}
+	interceptors := []connect.Interceptor{PanicInterceptor(), MetadataInterceptor(log.Debug), otelInterceptor()}
 	if ftl.Version != "dev" {
 		interceptors = append(interceptors, versionInterceptor{})
 	}
@@ -90,7 +91,7 @@ func DefaultClientOptions(level log.Level) []connect.ClientOption {
 }
 
 func DefaultHandlerOptions() []connect.HandlerOption {
-	interceptors := []connect.Interceptor{MetadataInterceptor(log.Debug), otelInterceptor()}
+	interceptors := []connect.Interceptor{PanicInterceptor(), MetadataInterceptor(log.Debug), otelInterceptor()}
 	if ftl.Version != "dev" {
 		interceptors = append(interceptors, versionInterceptor{})
 	}
@@ -105,6 +106,51 @@ func otelInterceptor() connect.Interceptor {
 	return otel
 }
 
+// PanicInterceptor intercepts panics and logs them.
+func PanicInterceptor() connect.Interceptor {
+	return &panicInterceptor{}
+}
+
+type panicInterceptor struct{}
+
+// Intercept and log any panics, then re-panic. Defer calls to this function to
+// trap panics in the calling function.
+func handlePanic(ctx context.Context) {
+	logger := log.FromContext(ctx)
+	if r := recover(); r != nil {
+		var err error
+		if rerr, ok := r.(error); ok {
+			err = rerr
+		} else {
+			err = fmt.Errorf("%v", r)
+		}
+		stack := string(debug.Stack())
+		logger.Errorf(err, "panic in RPC: %s", stack)
+		panic(err)
+	}
+}
+
+func (*panicInterceptor) WrapStreamingClient(req connect.StreamingClientFunc) connect.StreamingClientFunc {
+	return func(ctx context.Context, s connect.Spec) connect.StreamingClientConn {
+		defer handlePanic(ctx)
+		return req(ctx, s)
+	}
+}
+
+func (*panicInterceptor) WrapStreamingHandler(req connect.StreamingHandlerFunc) connect.StreamingHandlerFunc {
+	return func(ctx context.Context, s connect.StreamingHandlerConn) error {
+		defer handlePanic(ctx)
+		return req(ctx, s)
+	}
+}
+
+func (*panicInterceptor) WrapUnary(uf connect.UnaryFunc) connect.UnaryFunc {
+	return func(ctx context.Context, req connect.AnyRequest) (connect.AnyResponse, error) {
+		defer handlePanic(ctx)
+		return uf(ctx, req)
+	}
+}
+
 // MetadataInterceptor propagates FTL metadata through servers and clients.
 //
 // "errorLevel" is the level at which errors will be logged