feat: make Exit() blocking with early exit support and lifecycle signals

  Allow Exit() to be called before Run() to short-circuit startup and
  unwind callbacks. Exit() now blocks until teardown completes and returns
  the exit cause, enabling the pattern: return lc.Exit(err).

  Add Started(), Stopping(), and Completed() lifecycle signals. Fix
  PanicOnError for async callbacks and race conditions in callback
  registration during concurrent teardown.

Author: IAmRadek

README.md

@@ -35,13 +35,23 @@ Exitplan manages two lifecycle phases:
   `Context()` for workers and other long-running tasks.  
   It is canceled as soon as shutdown begins (via `Exit()`, signal, or startup timeout).
 
-- **Teardown**: after `Exit()` is called. Use `TeardownContext()` in shutdown callbacks.  
+- **Teardown**: calling `Exit(reason or nil)` starts the teardown.  
+  It blocks until all registered callbacks are complete.
+  Use `TeardownContext()` in shutdown callbacks.  
   It is canceled when the global teardown timeout elapses.
 
+> [!NOTE]
+> Calling `Exit()` before `Run()` starts the teardown immediately.
+> If no teardown timeout is set and a callback hangs, Exit() will block indefinitely.
+
 Use
 `Started()` to receive a signal when the application enters the running phase.                                                                                                                                                                                                                                                                                                                                             
 This is useful for readiness probes or coordinating dependent services.
 
+Use `Stopping()` to receive a signal when the application enters the teardown phase.
+
+Use `Completed()` to receive a signal when the teardown phase completes.
+
 ### Startup Timeout
 
 Use `WithStartupTimeout()` to detect stuck initialization:
@@ -140,8 +150,12 @@ func main() {
 
 	// Signal readiness when Run() starts
 	go func() {
-		<-ex.Started()
-		fmt.Println("Application is now running and ready")
+		select {
+		case <-ex.Started():
+			fmt.Println("Application is now running")
+		case <-ex.Stopping():
+			fmt.Println("Application is shutting down before it was ready")
+		}
 		// e.g., signal readiness probe, notify dependent services
 	}()
 
@@ -214,6 +228,49 @@ func main() {
 
 ```
 
+### Early Exit During Setup
+
+If initialization fails, use `Exit()` to short-circuit and unwind callbacks without calling `Run()`:
+
+```go
+package main
+
+import (
+	"fmt"
+	"syscall"
+	"time"
+
+	"github.com/struct0x/exitplan"
+)
+
+func run() error {
+	ex := exitplan.New(
+		exitplan.WithSignal(syscall.SIGINT, syscall.SIGTERM),
+		exitplan.WithTeardownTimeout(5*time.Second),
+	)
+
+	db, err := connectDB()
+	if err != nil {
+		return ex.Exit(err) // teardown runs, then returns err
+	}
+	ex.OnExit(func() { db.Close() })
+
+	cache, err := connectCache()
+	if err != nil {
+		return ex.Exit(err) // db.Close() runs, then returns err
+	}
+	ex.OnExit(func() { cache.Close() })
+
+	return ex.Run()
+}
+
+func main() {
+	if err := run(); err != nil {
+		fmt.Println(err)
+	}
+}
+```
+
 ## License
 
 This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.

exitplan.go

@@ -4,15 +4,17 @@ It provides a way to register functions that will be called when the application
 
 The application is considered to be running after calling Exitplan.Run() and before calling Exitplan.Exit().
 Use Exitplan.Started() to receive a signal when the application enters the running phase.
+Use Exitplan.Stopping() to receive a signal when the application enters the teardown phase.
 Use Exitplan.Context() to get a context bound to the application lifetime.
+Use Exitplan.Completed() to receive a signal when teardown is completed.
 
 The application is considered to be tearing down after calling Exitplan.Exit().
 You can use Exitplan.TeardownContext() to get a context that can be used to control the teardown phase.
 It is canceled when the teardown timeout is reached.
 
 Ordering of shutdown callbacks
 Exitplan executes registered exit callbacks in LIFO order (last registered, first executed).
-Async only offloads the execution to a goroutine, but Exitplan still waits for all callbacks up to the
+Async only offloads the execution to a goroutine, but the Exitplan still waits for all callbacks up to the
 teardown timeout.
 */
 package exitplan
@@ -43,8 +45,7 @@ type Exitplan struct {
 	runningCtx    context.Context
 	runningCancel context.CancelCauseFunc
 
-	startingCtx     context.Context
-	startingCancel  context.CancelFunc
+	started         chan struct{}
 	startingTimeout time.Duration
 
 	teardownCtx     context.Context
@@ -53,6 +54,10 @@ type Exitplan struct {
 
 	callbacks []*callback
 
+	panicValue any
+	panicOnce  sync.Once
+	ranViaRun  bool
+
 	errorHandler func(error)
 	callbacksMu  *sync.Mutex
 }
@@ -64,6 +69,7 @@ func New(opts ...opt) *Exitplan {
 	teardownCtx, teardownCancel := context.WithCancel(context.Background())
 	l := &Exitplan{
 		die:         make(chan struct{}),
+		started:     make(chan struct{}),
 		callbacksMu: &sync.Mutex{},
 		callbacks:   make([]*callback, 0),
 
@@ -85,26 +91,57 @@ func New(opts ...opt) *Exitplan {
 func (l *Exitplan) start() {
 	l.phase.Store(int32(phaseStarting))
 
-	ctx, cancel := context.WithCancel(context.Background())
 	if l.startingTimeout > 0 {
-		ctx, cancel = context.WithTimeout(ctx, l.startingTimeout)
+		ctx, cancel := context.WithTimeout(context.Background(), l.startingTimeout)
 		go func() {
+			defer cancel()
 			<-ctx.Done()
 
-			if l.phase.Load() == int32(phaseStarting) {
+			if l.phase.CompareAndSwap(int32(phaseStarting), int32(phaseTeardown)) {
 				l.runningCancel(ErrStartupTimeout)
 				close(l.die)
 			}
 		}()
 	}
 
-	l.startingCtx = ctx
-	l.startingCancel = cancel
+	go func() {
+		<-l.die
+
+		go func() {
+			if l.teardownTimeout > 0 {
+				<-time.After(l.teardownTimeout)
+				l.teardownCancel()
+			}
+		}()
+
+		defer l.teardownCancel()
+		defer func() {
+			if r := recover(); r != nil {
+				l.capturePanic(r)
+			}
+		}()
+
+		l.exit()
+	}()
 }
 
-// Started returns a channel that is closed after Run is called.
+// Started returns a channel that is closed after Run is called and the application enters the running phase.
+// It will never fire if Exit is called before Run, or if the startup timeout expires.
 func (l *Exitplan) Started() <-chan struct{} {
-	return l.startingCtx.Done()
+	return l.started
+}
+
+// Completed returns a channel closed after either all callback or teardown context is completed.
+// This signal means that now the application can exit.
+func (l *Exitplan) Completed() <-chan struct{} {
+	return l.teardownCtx.Done()
+}
+
+// Stopping returns a channel closed when the teardown phase starts.
+// This is equivalent to calling Exitplan.Context().Done().
+// It's fired before the application enters the teardown phase by calling Exit.
+func (l *Exitplan) Stopping() <-chan struct{} {
+	return l.runningCtx.Done()
 }
 
 // Context returns a main context. It will be canceled when the application is about to exit.
@@ -121,7 +158,7 @@ func (l *Exitplan) TeardownContext() context.Context {
 }
 
 // OnExit registers a callback that will be called when the application is about to exit.
-// Has no effect after calling Exitplan.Run().
+// Has no effect after calling Exitplan.Run() or Exitplan.Exit().
 // Use exitplan.Async option to execute the callback in a separate goroutine.
 // exitplan.PanicOnError has no effect on this function.
 // See also: OnExitWithError, OnExitWithContext, OnExitWithContextError.
@@ -135,7 +172,7 @@ func (l *Exitplan) OnExit(callback func(), exitOpts ...exitCallbackOpt) {
 }
 
 // OnExitWithError registers a callback that will be called when the application is about to exit.
-// Has no effect after calling Exitplan.Run().
+// Has no effect after calling Exitplan.Run() or Exitplan.Exit().
 // The callback can return an error that will be passed to the error handler.
 // Use exitplan.Async option to execute the callback in a separate goroutine.
 // Use exitplan.PanicOnError to panic with the error returned by the callback.
@@ -146,7 +183,7 @@ func (l *Exitplan) OnExitWithError(callback func() error, exitOpts ...exitCallba
 }
 
 // OnExitWithContext registers a callback that will be called when the application is about to exit.
-// Has no effect after calling Exitplan.Run().
+// Has no effect after calling Exitplan.Run() or Exitplan.Exit().
 // The callback will receive a context that will be canceled after the teardown timeout.
 // Use exitplan.Async option to execute the callback in a separate goroutine.
 // exitplan.PanicOnError has no effect on this function.
@@ -158,7 +195,7 @@ func (l *Exitplan) OnExitWithContext(callback func(context.Context), exitOpts ..
 }
 
 // OnExitWithContextError registers a callback that will be called when the application is about to exit.
-// Has no effect after calling Exitplan.Run().
+// Has no effect after calling Exitplan.Run() or Exitplan.Exit().
 // The callback will receive a context that will be canceled after the teardown timeout.
 // The callback can return an error that will be passed to the error handler.
 // Use exitplan.Async option to execute the callback in a separate goroutine.
@@ -168,10 +205,6 @@ func (l *Exitplan) OnExitWithContextError(callback func(context.Context) error,
 }
 
 func (l *Exitplan) addCallback(cb func(context.Context) error, exitOpts ...exitCallbackOpt) {
-	if l.phase.Load() != int32(phaseStarting) {
-		return
-	}
-
 	c := &callback{
 		name: callerLocation(3),
 		fn:   cb,
@@ -183,63 +216,84 @@ func (l *Exitplan) addCallback(cb func(context.Context) error, exitOpts ...exitC
 
 	l.callbacksMu.Lock()
 	defer l.callbacksMu.Unlock()
+
+	if l.phase.Load() != int32(phaseStarting) {
+		return
+	}
+
 	l.callbacks = append(l.callbacks, c)
 }
 
-// Exit stops the application. It will cause the application to exit with the specified reason.
+// Exit stops the application and waits for all registered callbacks to complete.
+// It returns the provided reason, enabling the pattern: return lc.Exit(err).
 // If the reason is nil, ErrGracefulShutdown will be used.
-// Multiple calls to Exit are safe, but only the first one will have an effect.
-func (l *Exitplan) Exit(reason error) {
+// Exit can be called before or after Run.
+// Multiple calls to Exit are safe; all will block until teardown completes.
+func (l *Exitplan) Exit(reason error) error {
 	if reason == nil {
 		reason = ErrGracefulShutdown
 	}
 
-	if !l.phase.CompareAndSwap(int32(phaseRunning), int32(phaseTeardown)) {
-		return
+	if l.phase.CompareAndSwap(int32(phaseRunning), int32(phaseTeardown)) ||
+		l.phase.CompareAndSwap(int32(phaseStarting), int32(phaseTeardown)) {
+		l.runningCancel(reason)
+		close(l.die)
 	}
 
-	l.runningCancel(reason)
-	close(l.die)
+	<-l.teardownCtx.Done()
+
+	if !l.ranViaRun && l.panicValue != nil {
+		panic(l.panicValue)
+	}
+
+	return context.Cause(l.runningCtx)
 }
 
 // Run starts the application. It will block until the application is stopped by calling Exit.
 // It will also block until all the registered callbacks are executed.
 // If the teardown timeout is set, it will be used to cancel the context passed to the callbacks.
+// If Exit was called before Run, Run will wait for teardown to complete and return the exit cause.
 // Returns the error that caused the application to stop.
 func (l *Exitplan) Run() (exitCause error) {
-	if !l.phase.CompareAndSwap(int32(phaseStarting), int32(phaseRunning)) {
-		panic("Exitplan.Run() called after Exitplan.Exit()")
-	}
+	l.ranViaRun = true
 
-	l.startingCancel()
-
-	<-l.die
+	// If Exit() was already called, skip the starting→running transition.
+	if l.phase.CompareAndSwap(int32(phaseStarting), int32(phaseRunning)) {
+		close(l.started)
+	}
 
-	go func() {
-		if l.teardownTimeout > 0 {
-			<-time.After(l.teardownTimeout)
-			l.teardownCancel()
-		}
-	}()
+	<-l.teardownCtx.Done()
 
-	l.exit()
+	if l.panicValue != nil {
+		panic(l.panicValue)
+	}
 
 	return context.Cause(l.runningCtx)
 }
 
 func (l *Exitplan) exit() {
 	ctx := l.TeardownContext()
 
+	l.callbacksMu.Lock()
+	callbacks := make([]*callback, len(l.callbacks))
+	copy(callbacks, l.callbacks)
+	l.callbacksMu.Unlock()
+
 	wg := sync.WaitGroup{}
 
-	for _, cb := range l.callbacks {
+	for _, cb := range callbacks {
 		if cb.executeBehaviour != executeAsync {
 			continue
 		}
 
 		wg.Add(1)
 		go func(cb *callback) {
 			defer wg.Done()
+			defer func() {
+				if r := recover(); r != nil {
+					l.capturePanic(r)
+				}
+			}()
 
 			execCtx := ctx
 			if cb.timeout > 0 {
@@ -257,8 +311,8 @@ func (l *Exitplan) exit() {
 		}(cb)
 	}
 
-	for i := len(l.callbacks) - 1; i >= 0; i-- {
-		cb := l.callbacks[i]
+	for i := len(callbacks) - 1; i >= 0; i-- {
+		cb := callbacks[i]
 		if cb.executeBehaviour != executeSync {
 			continue
 		}
@@ -303,6 +357,12 @@ func (l *Exitplan) handleExitError(errBehaviour exitBehaviour, err error) {
 	}
 }
 
+func (l *Exitplan) capturePanic(r any) {
+	l.panicOnce.Do(func() {
+		l.panicValue = r
+	})
+}
+
 func callbackWithContext(ctx context.Context, callback func() error) error {
 	errChan := make(chan error)
 	go func() {