From 9a2a97222d5eb7908ec42f890c7a25a4f0cd5785 Mon Sep 17 00:00:00 2001
From: FillZpp <FillZpp.pub@gmail.com>
Date: Mon, 9 Aug 2021 13:46:27 +0800
Subject: [PATCH] Recover panic for controller reconcile and webhook handler

Signed-off-by: FillZpp <FillZpp.pub@gmail.com>
---
 pkg/controller/controller.go               |  4 +++
 pkg/internal/controller/controller.go      | 17 +++++++++--
 pkg/internal/controller/controller_test.go | 33 ++++++++++++++++++++++
 3 files changed, 52 insertions(+), 2 deletions(-)

diff --git a/pkg/controller/controller.go b/pkg/controller/controller.go
index c9e07562a3..88ba786717 100644
--- a/pkg/controller/controller.go
+++ b/pkg/controller/controller.go
@@ -52,6 +52,9 @@ type Options struct {
 	// CacheSyncTimeout refers to the time limit set to wait for syncing caches.
 	// Defaults to 2 minutes if not set.
 	CacheSyncTimeout time.Duration
+
+	// RecoverPanic indicates whether the panic caused by reconcile should be recovered.
+	RecoverPanic bool
 }
 
 // Controller implements a Kubernetes API.  A Controller manages a work queue fed reconcile.Requests
@@ -133,5 +136,6 @@ func NewUnmanaged(name string, mgr manager.Manager, options Options) (Controller
 		SetFields:               mgr.SetFields,
 		Name:                    name,
 		Log:                     options.Log.WithName("controller").WithName(name),
+		RecoverPanic:            options.RecoverPanic,
 	}, nil
 }
diff --git a/pkg/internal/controller/controller.go b/pkg/internal/controller/controller.go
index 224d300b89..87431a438f 100644
--- a/pkg/internal/controller/controller.go
+++ b/pkg/internal/controller/controller.go
@@ -85,6 +85,9 @@ type Controller struct {
 
 	// Log is used to log messages to users during reconciliation, or for example when a watch is started.
 	Log logr.Logger
+
+	// RecoverPanic indicates whether the panic caused by reconcile should be recovered.
+	RecoverPanic bool
 }
 
 // watchDescription contains all the information necessary to start a watch.
@@ -95,7 +98,17 @@ type watchDescription struct {
 }
 
 // Reconcile implements reconcile.Reconciler.
-func (c *Controller) Reconcile(ctx context.Context, req reconcile.Request) (reconcile.Result, error) {
+func (c *Controller) Reconcile(ctx context.Context, req reconcile.Request) (_ reconcile.Result, err error) {
+	if c.RecoverPanic {
+		defer func() {
+			if r := recover(); r != nil {
+				for _, fn := range utilruntime.PanicHandlers {
+					fn(r)
+				}
+				err = fmt.Errorf("panic: %v [recovered]", r)
+			}
+		}()
+	}
 	log := c.Log.WithValues("name", req.Name, "namespace", req.Namespace)
 	ctx = logf.IntoContext(ctx, log)
 	return c.Do.Reconcile(ctx, req)
@@ -295,7 +308,7 @@ func (c *Controller) reconcileHandler(ctx context.Context, obj interface{}) {
 
 	// RunInformersAndControllers the syncHandler, passing it the Namespace/Name string of the
 	// resource to be synced.
-	result, err := c.Do.Reconcile(ctx, req)
+	result, err := c.Reconcile(ctx, req)
 	switch {
 	case err != nil:
 		c.Queue.AddRateLimited(req)
diff --git a/pkg/internal/controller/controller_test.go b/pkg/internal/controller/controller_test.go
index 7545be73a7..9f23fa2abc 100644
--- a/pkg/internal/controller/controller_test.go
+++ b/pkg/internal/controller/controller_test.go
@@ -88,6 +88,39 @@ var _ = Describe("controller", func() {
 			Expect(err).NotTo(HaveOccurred())
 			Expect(result).To(Equal(reconcile.Result{Requeue: true}))
 		})
+
+		It("should not recover panic if RecoverPanic is false by default", func() {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			defer func() {
+				Expect(recover()).ShouldNot(BeNil())
+			}()
+			ctrl.Do = reconcile.Func(func(context.Context, reconcile.Request) (reconcile.Result, error) {
+				var res *reconcile.Result
+				return *res, nil
+			})
+			_, _ = ctrl.Reconcile(ctx,
+				reconcile.Request{NamespacedName: types.NamespacedName{Namespace: "foo", Name: "bar"}})
+		})
+
+		It("should recover panic if RecoverPanic is true", func() {
+			ctx, cancel := context.WithCancel(context.Background())
+			defer cancel()
+
+			defer func() {
+				Expect(recover()).To(BeNil())
+			}()
+			ctrl.RecoverPanic = true
+			ctrl.Do = reconcile.Func(func(context.Context, reconcile.Request) (reconcile.Result, error) {
+				var res *reconcile.Result
+				return *res, nil
+			})
+			_, err := ctrl.Reconcile(ctx,
+				reconcile.Request{NamespacedName: types.NamespacedName{Namespace: "foo", Name: "bar"}})
+			Expect(err).To(HaveOccurred())
+			Expect(err.Error()).To(ContainSubstring("[recovered]"))
+		})
 	})
 
 	Describe("Start", func() {