diff --git a/injector.go b/injector.go index 5fa3fe8..bca4e99 100644 --- a/injector.go +++ b/injector.go @@ -128,6 +128,14 @@ type InjectorOpts struct { // This hook can be used for logging, metrics collection, or error handling. HookAfterShutdown []func(scope *Scope, serviceName string, err error) + // HookBeforeHealthCheck is called before a service health check. + // This hook can be used for cleanup preparation or logging. + HookBeforeHealthCheck []func(scope *Scope, serviceName string) + + // HookAfterHealthCheck is called after a service is health checked. + // This hook can be used for logging, metrics collection, or error handling. + HookAfterHealthCheck []func(scope *Scope, serviceName string) + // Logf is the logging function used by the DI container for internal logging. // If not provided, no logging will occur. This function should handle the format // string and arguments similar to fmt.Printf. @@ -151,7 +159,7 @@ type InjectorOpts struct { // StructTagKey specifies the tag key used for struct field injection. // Default: "do" (see DefaultStructTagKey constant). // This allows customization of the struct tag format for injection. - StructTagKey string + StructTagKey string } func (o *InjectorOpts) copy() *InjectorOpts { @@ -162,6 +170,8 @@ func (o *InjectorOpts) copy() *InjectorOpts { HookAfterInvocation: append([]func(*Scope, string, error){}, o.HookAfterInvocation...), HookBeforeShutdown: append([]func(*Scope, string){}, o.HookBeforeShutdown...), HookAfterShutdown: append([]func(*Scope, string, error){}, o.HookAfterShutdown...), + HookBeforeHealthCheck: append([]func(*Scope, string){}, o.HookBeforeHealthCheck...), + HookAfterHealthCheck: append([]func(*Scope, string){}, o.HookAfterHealthCheck...), Logf: o.Logf, HealthCheckParallelism: o.HealthCheckParallelism, HealthCheckGlobalTimeout: o.HealthCheckGlobalTimeout, @@ -205,3 +215,15 @@ func (o *InjectorOpts) onAfterShutdown(scope *Scope, serviceName string, err err fn(scope, serviceName, err) } } + +func (o *InjectorOpts) onBeforeHealthCheck(scope *Scope, serviceName string) { + for _, fn := range o.HookBeforeHealthCheck { + fn(scope, serviceName) + } +} + +func (o *InjectorOpts) onAfterHealthCheck(scope *Scope, serviceName string) { + for _, fn := range o.HookAfterHealthCheck { + fn(scope, serviceName) + } +} diff --git a/root_scope.go b/root_scope.go index 268105c..ec5cfbb 100644 --- a/root_scope.go +++ b/root_scope.go @@ -85,6 +85,12 @@ func NewWithOpts(opts *InjectorOpts, packages ...func(Injector)) *RootScope { if opts.HookAfterShutdown == nil { opts.HookAfterShutdown = []func(*Scope, string, error){} } + if opts.HookBeforeHealthCheck == nil { + opts.HookBeforeHealthCheck = []func(*Scope, string){} + } + if opts.HookAfterHealthCheck == nil { + opts.HookAfterHealthCheck = []func(*Scope, string){} + } root := &RootScope{ self: newScope(DefaultRootScopeName, nil, nil), @@ -289,6 +295,20 @@ func (s *RootScope) AddAfterShutdownHook(hook func(*Scope, string, error)) { s.opts.HookAfterShutdown = append(s.opts.HookAfterShutdown, hook) } +// AddBeforeHealthCheckHook adds a hook that will be called before a service health check. +// +// Play: https://go.dev/play/p/wtKubQHkFLK +func (s *RootScope) AddBeforeHealthCheckHook(hook func(*Scope, string)) { + s.opts.HookBeforeHealthCheck = append(s.opts.HookBeforeHealthCheck, hook) +} + +// AddAfterHealthCheckHook adds a hook that will be called after a service is health checked. +// +// Play: https://go.dev/play/p/wtKubQHkFLK +func (s *RootScope) AddAfterHealthCheckHook(hook func(*Scope, string)) { + s.opts.HookAfterHealthCheck = append(s.opts.HookAfterHealthCheck, hook) +} + // Clone clones injector with provided services but not with invoked instances. // // Play: https://go.dev/play/p/DqIlXhZ8c4t diff --git a/scope.go b/scope.go index f5a547e..6f000d0 100644 --- a/scope.go +++ b/scope.go @@ -740,10 +740,15 @@ func (s *Scope) serviceHealthCheck(ctx context.Context, name string) error { // A timeout error is not triggered when the service is not a healthchecker. // If the healthchecker does not support context.Timeout, the error will be triggered by raceWithTimeout(). - return raceWithTimeout( + + s.RootScope().opts.onBeforeHealthCheck(s, name) + err := raceWithTimeout( ctx, service.healthcheck, ) + s.RootScope().opts.onAfterHealthCheck(s, name) + + return err } // Should never happen.