lcfx.go

// Package lcfx provides a helper for go.uber.org/fx to manage long-running,
// asynchronous tasks like web servers or background workers within the fx lifecycle.
package lcfx

import (
	"context"
	"fmt"
	"reflect"
	"runtime"
	"strings"
	"sync"
	"time"

	"go.uber.org/fx"
)

// Module provides the lcfx lifecycle implementation to the fx application.
// It should be included in the fx.New() call.
var Module = fx.Module("lc", fx.Provide(New))

// Lifecycle is an abstraction over the fx.Lifecycle that adds support for
// running hooks asynchronously.
type Lifecycle interface {
	// Append registers a hook with the application's lifecycle. It is identical
	// to fx.Lifecycle.Append.
	Append(fx.Hook)
	// AppendAsync registers a hook to be executed in a non-blocking manner.
	// The OnStart function of the hook is run in a new goroutine, allowing the
	// application to continue starting up.
	AppendAsync(fx.Hook)
}

// lifecycle implements the Lifecycle interface.
// It wraps the standard fx.Lifecycle and fx.Shutdowner to add asynchronous capabilities.
type lifecycle struct {
	fx.Lifecycle
	fx.Shutdowner
	// wg is used to wait for all asynchronous hooks to complete during shutdown.
	wg *sync.WaitGroup
}

// New creates a new lcfx.Lifecycle.
// It registers a stop hook that waits for all async goroutines to finish.
func New(lc fx.Lifecycle, sd fx.Shutdowner) Lifecycle {
	wg := new(sync.WaitGroup)
	// Add a stop hook to the underlying fx.Lifecycle.
	// This ensures that during shutdown, the application will wait for all
	// goroutines started by AppendAsync to finish.
	lc.Append(fx.StopHook(wg.Wait))
	return &lifecycle{
		Lifecycle:  lc,
		Shutdowner: sd,
		wg:         wg,
	}
}

// AppendAsync runs the OnStart hook in a new goroutine and manages its lifecycle.
// This is the core of the lcfx library.
func (l *lifecycle) AppendAsync(hook fx.Hook) {
	if onStart := hook.OnStart; onStart != nil {
		// isWrapped checks if the hook is already wrapped by fx's internal StartStopHook.
		// This is a defensive check to avoid double-wrapping.
		if !isWrapped(hook) {
			hook = fx.StartStopHook(hook.OnStart, hook.OnStop)
		}

		// cancelCtx is used to signal shutdown to the running goroutine.
		cancelCtx, cancel := context.WithCancel(context.Background())

		// Wrap the original OnStart to run it in a goroutine.
		hook.OnStart = func(ctx context.Context) error {
			// The context passed to the user's OnStart is a merged context,
			// which allows passing values from the original fx context while
			// controlling cancellation independently.
			ctx = &mergedContext{
				ctx:       ctx,
				cancelCtx: cancelCtx,
			}
			l.wg.Add(1)
			go func() {
				defer func() {
					// Signal cancellation and decrement the WaitGroup counter when the goroutine exits.
					cancel()
					l.wg.Done()
				}()
				// Execute the original OnStart function.
				if err := onStart(ctx); err != nil {
					// If the async task returns an error, initiate a shutdown of the entire application.
					// The error is ignored as we are already in a shutdown sequence.
					_ = l.Shutdowner.Shutdown(fx.ExitCode(1))
				}
			}()
			return nil // Return immediately, not blocking the application startup.
		}

		// Wrap the original OnStop to ensure cancellation is triggered.
		onStop := hook.OnStop
		hook.OnStop = func(ctx context.Context) error {
			// Trigger cancellation for the goroutine.
			cancel()
			// Execute the original OnStop if it exists.
			if onStop != nil {
				return onStop(ctx)
			}
			return nil
		}
	}

	// Append the modified hook to the actual fx.Lifecycle.
	l.Lifecycle.Append(hook)
}

// mergedContext combines the original fx.Context (for values) with a new
// cancellable context (for shutdown signaling). This allows the async goroutine
// to be cancelled by lcfx while still having access to values from the fx graph.
type mergedContext struct {
	ctx       context.Context // The original context from Fx, used for Value().
	cancelCtx context.Context // The cancellable context for Done(), Err(), and Deadline().
}

func (c *mergedContext) Deadline() (time.Time, bool) {
	return c.cancelCtx.Deadline()
}

func (c *mergedContext) Done() <-chan struct{} {
	return c.cancelCtx.Done()
}

func (c *mergedContext) Err() error {
	return c.cancelCtx.Err()
}

func (c *mergedContext) Value(key any) any {
	return c.ctx.Value(key)
}

// isWrapped uses reflection to check if a hook has already been wrapped by
// fx's internal lifecycle wrapper.
//
// WARNING: This function depends on the internal implementation details of the
// go.uber.org/fx package. It checks for a function name prefix that is not
// part of the public API. This could break if the fx library changes its
// internal structure in a future version.
func isWrapped(hook fx.Hook) bool {
	const prefix = "go.uber.org/fx/internal/lifecycle.Wrap"

	for _, fn := range []any{
		hook.OnStart, hook.OnStop,
	} {
		if fn != nil && strings.HasPrefix(funcName(fn), prefix) {
			return true
		}
	}

	return false
}

// funcName returns the name of a function using reflection.
func funcName(fn any) string {
	v := reflect.ValueOf(fn)
	if v.Kind() != reflect.Func {
		return fmt.Sprint(fn)
	}
	return runtime.FuncForPC(v.Pointer()).Name()
}