fix(GODT-2500): Recover in deferred function.

async/logging.go

@@ -0,0 +1,14 @@
+package async
+
+import (
+	"context"
+
+	"github.com/ProtonMail/gluon/logging"
+)
+
+func GoAnnotated(ctx context.Context, panicHandler PanicHandler, fn func(context.Context), labelMap ...logging.Labels) {
+	go func() {
+		defer HandlePanic(panicHandler)
+		logging.DoAnnotated(ctx, fn, labelMap...)
+	}()
+}

async/panic_handler.go

@@ -1,15 +1,25 @@
 package async
 
 type PanicHandler interface {
-	HandlePanic()
+	HandlePanic(interface{})
 }
 
 type NoopPanicHandler struct{}
 
-func (n NoopPanicHandler) HandlePanic() {}
+func (n NoopPanicHandler) HandlePanic(r interface{}) {}
 
 func HandlePanic(panicHandler PanicHandler) {
-	if panicHandler != nil {
-		panicHandler.HandlePanic()
+	if panicHandler == nil {
+		return
 	}
+
+	if _, ok := panicHandler.(NoopPanicHandler); ok {
+		return
+	}
+
+	if _, ok := panicHandler.(*NoopPanicHandler); ok {
+		return
+	}
+
+	panicHandler.HandlePanic(recover())
 }

async/panic_handler_test.go

@@ -0,0 +1,48 @@
+package async
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+type recoverHandler struct{}
+
+func (h recoverHandler) HandlePanic(r interface{}) {
+	fmt.Println("recoverHandler", r)
+}
+
+func TestPanicHandler(t *testing.T) {
+	require.NotPanics(t, func() {
+		defer HandlePanic(recoverHandler{})
+		panic("there")
+	})
+
+	require.PanicsWithValue(t, "where", func() {
+		defer HandlePanic(NoopPanicHandler{})
+		panic("where")
+	})
+
+	require.PanicsWithValue(t, "everywhere", func() {
+		defer HandlePanic(nil)
+		panic("everywhere")
+	})
+
+	require.NotPanics(t, func() {
+		defer HandlePanic(recoverHandler{})
+		panic(nil)
+	})
+
+	require.NotPanics(t, func() {
+		defer HandlePanic(recoverHandler{})
+	})
+
+	require.NotPanics(t, func() {
+		defer HandlePanic(NoopPanicHandler{})
+	})
+
+	require.NotPanics(t, func() {
+		defer HandlePanic(&NoopPanicHandler{})
+	})
+}

async/queued_channel.go

@@ -3,8 +3,6 @@ package async
 import (
 	"context"
 	"sync"
-
-	"github.com/ProtonMail/gluon/logging"
 )
 
 // QueuedChannel represents a channel on which queued items can be published without having to worry if the reader
@@ -30,7 +28,7 @@ func NewQueuedChannel[T any](chanBufferSize, queueCapacity int, panicHandler Pan
 	queue.closed.store(false)
 
 	// Start the queue consumer.
-	logging.GoAnnotated(context.Background(), panicHandler, func(ctx context.Context) {
+	GoAnnotated(context.Background(), panicHandler, func(ctx context.Context) {
 		defer close(queue.ch)
 
 		for {

internal/backend/update_injector.go

@@ -32,7 +32,7 @@ func newUpdateInjector(connector connector.Connector, userID string, panicHandle
 
 	injector.forwardWG.Add(1)
 
-	logging.GoAnnotated(context.Background(), panicHandler, func(ctx context.Context) {
+	async.GoAnnotated(context.Background(), panicHandler, func(ctx context.Context) {
 		injector.forward(ctx, connector.GetUpdates())
 	}, logging.Labels{
 		"Action": "Forwarding updates",

internal/backend/user.go

@@ -125,7 +125,7 @@ func newUser(
 	user.updateWG.Add(1)
 
 	// nolint:contextcheck
-	logging.GoAnnotated(context.Background(), panicHandler, func(ctx context.Context) {
+	async.GoAnnotated(context.Background(), panicHandler, func(ctx context.Context) {
 		defer user.updateWG.Done()
 
 		updateCh := user.updateInjector.GetUpdates()

internal/session/command.go

@@ -5,6 +5,7 @@ import (
 	"context"
 	"errors"
 
+	"github.com/ProtonMail/gluon/async"
 	"github.com/ProtonMail/gluon/imap/command"
 	"github.com/ProtonMail/gluon/internal/response"
 	"github.com/ProtonMail/gluon/logging"
@@ -21,7 +22,7 @@ type commandResult struct {
 func (s *Session) startCommandReader(ctx context.Context) <-chan commandResult {
 	cmdCh := make(chan commandResult)
 
-	logging.GoAnnotated(ctx, s.panicHandler, func(ctx context.Context) {
+	async.GoAnnotated(ctx, s.panicHandler, func(ctx context.Context) {
 		defer close(cmdCh)
 
 		tlsHeaders := [][]byte{

internal/session/handle_idle.go

@@ -4,6 +4,7 @@ import (
 	"context"
 	"time"
 
+	"github.com/ProtonMail/gluon/async"
 	"github.com/ProtonMail/gluon/imap/command"
 	"github.com/ProtonMail/gluon/internal/response"
 	"github.com/ProtonMail/gluon/logging"
@@ -22,7 +23,7 @@ func (s *Session) handleIdle(ctx context.Context, tag string, _ *command.Idle, c
 	}
 
 	return s.state.Idle(ctx, func(pending []response.Response, resCh chan response.Response) error {
-		logging.GoAnnotated(ctx, s.panicHandler, func(ctx context.Context) {
+		async.GoAnnotated(ctx, s.panicHandler, func(ctx context.Context) {
 			if s.idleBulkTime != 0 {
 				sendResponsesInBulks(s, resCh, s.idleBulkTime)
 			} else {

logging/logging.go

@@ -20,24 +20,6 @@ const (
 	LineKey = "line"
 )
 
-type PanicHandler interface {
-	HandlePanic()
-}
-
-func GoAnnotated(ctx context.Context, panicHandler PanicHandler, fn func(context.Context), labelMap ...Labels) {
-	pprofDo(ctx, toLabelSet(labelMap...), func(ctx context.Context) {
-		go func() {
-			defer func() {
-				if panicHandler != nil {
-					panicHandler.HandlePanic()
-				}
-			}()
-
-			fn(ctx)
-		}()
-	})
-}
-
 func DoAnnotated(ctx context.Context, fn func(context.Context), labelMap ...Labels) {
 	pprofDo(ctx, toLabelSet(labelMap...), fn)
 }

server.go

@@ -334,7 +334,7 @@ func (s *Server) getNextID() int {
 func (s *Server) newEventCh(ctx context.Context) chan events.Event {
 	eventCh := make(chan events.Event)
 
-	logging.GoAnnotated(ctx, s.panicHandler, func(ctx context.Context) {
+	async.GoAnnotated(ctx, s.panicHandler, func(ctx context.Context) {
 		for event := range eventCh {
 			s.publish(event)
 		}
@@ -364,11 +364,7 @@ func newConnCh(l net.Listener, panicHandler async.PanicHandler) <-chan net.Conn
 	connCh := make(chan net.Conn)
 
 	go func() {
-		defer func() {
-			if panicHandler != nil {
-				panicHandler.HandlePanic()
-			}
-		}()
+		defer async.HandlePanic(panicHandler)
 
 		defer close(connCh)
 

tests/full_state_test.go

@@ -112,7 +112,7 @@ func TestReceptionOnIdle(t *testing.T) {
 		wg.Add(2)
 
 		// idling.
-		logging.GoAnnotated(context.Background(), async.NoopPanicHandler{}, func(ctx context.Context) {
+		async.GoAnnotated(context.Background(), async.NoopPanicHandler{}, func(ctx context.Context) {
 			defer wg.Done()
 			done <- c.Idle(stop, nil)
 		}, logging.Labels{
@@ -121,7 +121,7 @@ func TestReceptionOnIdle(t *testing.T) {
 		})
 
 		// receiving messages from another client.
-		logging.GoAnnotated(context.Background(), async.NoopPanicHandler{}, func(ctx context.Context) {
+		async.GoAnnotated(context.Background(), async.NoopPanicHandler{}, func(ctx context.Context) {
 			defer wg.Done()
 
 			cli := sess.newClient()