feat(shutdown): add panic handling

use simplified handling while waiting for
golang/go#53757

Signed-off-by: Artsiom Koltun <[email protected]>

pkg/utils/shutdown.go

@@ -7,6 +7,7 @@ package utils
 import (
  "context"
  "errors"
+ "fmt"
  "log"
  "net"
  "net/http"
@@ -141,7 +142,7 @@ func (s *ShutdownHandler) RunAndWait() error {
  for i := range s.serves {
  fn := s.serves[i]
  s.eg.Go(func() error {
- return fn()
+ return wrapServeFuncPanic(fn)()
  })
  }
 
@@ -163,7 +164,8 @@ func (s *ShutdownHandler) RunAndWait() error {
  for i := len(s.shutdowns) - 1; i >= 0; i-- {
  timeoutCtx, cancel := context.WithTimeout(context.Background(), s.timeoutPerShutdown)
  defer cancel()
- err = errors.Join(err, s.shutdowns[i](timeoutCtx))
+ shutdownFn := wrapShutdownFuncPanic(s.shutdowns[i])
+ err = errors.Join(err, shutdownFn(timeoutCtx))
  }
 
  return err
@@ -172,6 +174,30 @@ func (s *ShutdownHandler) RunAndWait() error {
  return s.eg.Wait()
 }
 
+func wrapServeFuncPanic(fn ServeFunc) ServeFunc {
+ return func() (err error) {
+ defer func() {
+ if r := recover(); r != nil {
+ err = fmt.Errorf("was panic for serve function, recovered value: %v", r)
+ }
+ }()
+ err = fn()
+ return err
+ }
+}
+
+func wrapShutdownFuncPanic(fn ShutdownFunc) ShutdownFunc {
+ return func(ctx context.Context) (err error) {
+ defer func() {
+ if r := recover(); r != nil {
+ err = fmt.Errorf("was panic for shutdown function, recovered value: %v", r)
+ }
+ }()
+ err = fn(ctx)
+ return err
+ }
+}
+
 func runWithCtx(ctx context.Context, fn func() error) error {
  var err error
 

pkg/utils/shutdown_test.go

@@ -7,6 +7,7 @@ package utils
 import (
  "context"
  "errors"
+ "log"
  "os"
  "sync"
  "testing"
@@ -26,19 +27,24 @@ func newServeShutdownPair(
  serveIDs, shutdownIDs *[]int,
  mu *sync.Mutex,
  serveErr, shutdownErr error,
+ servePanic bool, shutdownPanic bool,
 ) *serveShutdownPair {
  shutdownTrigger := make(chan struct{}, 1)
  s := &serveShutdownPair{
  shutdownTrigger: shutdownTrigger,
  serve: func() error {
- if serveErr == nil {
- <-shutdownTrigger
- }
-
  mu.Lock()
  *serveIDs = append(*serveIDs, fnID)
  mu.Unlock()
 
+ if servePanic {
+ log.Panic("Panic!")
+ }
+
+ if serveErr == nil {
+ <-shutdownTrigger
+ }
+
  return serveErr
  },
  shutdown: func(ctx context.Context) error {
@@ -48,38 +54,74 @@ func newServeShutdownPair(
 
  shutdownTrigger <- struct{}{}
 
+ if shutdownPanic {
+ log.Panic("Panic!")
+ }
+
  return shutdownErr
  },
  }
 
  return s
 }
 
+func errString(err error) string {
+ if err == nil {
+ return ""
+ }
+
+ return err.Error()
+}
+
 func TestRunAndWait(t *testing.T) {
  stubErr := errors.New("stub error")
  tests := map[string]struct {
  giveServeErr error
+ giveServePanic bool
  giveShutdownErr error
+ giveShutdownPanic bool
  stoppedByInterrupt bool
- wantErr error
+ wantErr string
  }{
  "all services successfully completed": {
  giveServeErr: nil,
+ giveServePanic: false,
  giveShutdownErr: nil,
+ giveShutdownPanic: false,
  stoppedByInterrupt: true,
- wantErr: nil,
+ wantErr: "",
  },
  "serve failed": {
  giveServeErr: stubErr,
+ giveServePanic: false,
  giveShutdownErr: nil,
  stoppedByInterrupt: false,
- wantErr: stubErr,
+ giveShutdownPanic: false,
+ wantErr: stubErr.Error(),
  },
  "shutdown failed": {
  giveServeErr: nil,
+ giveServePanic: false,
  giveShutdownErr: stubErr,
+ giveShutdownPanic: false,
+ stoppedByInterrupt: true,
+ wantErr: stubErr.Error(),
+ },
+ "serve panic": {
+ giveServeErr: nil,
+ giveServePanic: true,
+ giveShutdownErr: nil,
+ giveShutdownPanic: false,
+ stoppedByInterrupt: false,
+ wantErr: "was panic for serve function, recovered value: Panic!",
+ },
+ "shutdown panic": {
+ giveServeErr: nil,
+ giveServePanic: false,
+ giveShutdownErr: nil,
+ giveShutdownPanic: true,
  stoppedByInterrupt: true,
- wantErr: stubErr,
+ wantErr: "was panic for shutdown function, recovered value: Panic!",
  },
  }
  for testName, tt := range tests {
@@ -89,9 +131,12 @@ func TestRunAndWait(t *testing.T) {
  serveFnIDs := &[]int{}
  shutdownFnIDs := &[]int{}
  mu := sync.Mutex{}
- s0 := newServeShutdownPair(0, serveFnIDs, shutdownFnIDs, &mu, nil, nil)
- s1 := newServeShutdownPair(1, serveFnIDs, shutdownFnIDs, &mu, tt.giveServeErr, tt.giveShutdownErr)
- s2 := newServeShutdownPair(2, serveFnIDs, shutdownFnIDs, &mu, nil, nil)
+ s0 := newServeShutdownPair(0, serveFnIDs, shutdownFnIDs, &mu, nil, nil, false, false)
+ s1 := newServeShutdownPair(1, serveFnIDs, shutdownFnIDs, &mu,
+ tt.giveServeErr, tt.giveShutdownErr,
+ tt.giveServePanic, tt.giveShutdownPanic,
+ )
+ s2 := newServeShutdownPair(2, serveFnIDs, shutdownFnIDs, &mu, nil, nil, false, false)
 
  sh.AddServe(s0.serve, s0.shutdown)
  sh.AddServe(s1.serve, s1.shutdown)
@@ -103,7 +148,7 @@ func TestRunAndWait(t *testing.T) {
 
  err := sh.RunAndWait()
 
- if !errors.Is(err, tt.wantErr) {
+ if errString(err) != tt.wantErr {
  t.Errorf("Expected error: %v, received: %v", tt.wantErr, err)
  }