feat: before and Afters with context params

app.go

@@ -89,11 +89,12 @@ func (a *App) Run() error {
 	a.mu.Lock()
 	a.instance = instance
 	a.mu.Unlock()
-	eg, ctx := errgroup.WithContext(NewContext(a.ctx, a))
+	sctx := NewContext(a.ctx, a)
+	eg, ctx := errgroup.WithContext(sctx)
 	wg := sync.WaitGroup{}
 
 	for _, fn := range a.opts.beforeStart {
-		if err = fn(); err != nil {
+		if err = fn(sctx); err != nil {
 			return err
 		}
 	}
@@ -108,7 +109,7 @@ func (a *App) Run() error {
 		wg.Add(1)
 		eg.Go(func() error {
 			wg.Done() // here is to ensure server start has begun running before register, so defer is not needed
-			return srv.Start(NewContext(a.opts.ctx, a))
+			return srv.Start(sctx)
 		})
 	}
 	wg.Wait()
@@ -120,7 +121,7 @@ func (a *App) Run() error {
 		}
 	}
 	for _, fn := range a.opts.afterStart {
-		if err = fn(); err != nil {
+		if err = fn(sctx); err != nil {
 			return err
 		}
 	}
@@ -139,15 +140,16 @@ func (a *App) Run() error {
 		return err
 	}
 	for _, fn := range a.opts.afterStop {
-		err = fn()
+		err = fn(sctx)
 	}
 	return err
 }
 
 // Stop gracefully stops the application.
 func (a *App) Stop() (err error) {
+	sctx := NewContext(a.ctx, a)
 	for _, fn := range a.opts.beforeStop {
-		err = fn()
+		err = fn(sctx)
 	}
 
 	a.mu.Lock()

app_test.go

@@ -47,19 +47,19 @@ func TestApp(t *testing.T) {
 		Name("kratos"),
 		Version("v1.0.0"),
 		Server(hs, gs),
-		BeforeStart(func() error {
+		BeforeStart(func(_ context.Context) error {
 			t.Log("BeforeStart...")
 			return nil
 		}),
-		BeforeStop(func() error {
+		BeforeStop(func(_ context.Context) error {
 			t.Log("BeforeStop...")
 			return nil
 		}),
-		AfterStart(func() error {
+		AfterStart(func(_ context.Context) error {
 			t.Log("AfterStart...")
 			return nil
 		}),
-		AfterStop(func() error {
+		AfterStop(func(_ context.Context) error {
 			t.Log("AfterStop...")
 			return nil
 		}),

options.go

@@ -32,10 +32,10 @@ type options struct {
 	servers          []transport.Server
 
 	// Before and After funcs
-	beforeStart []func() error
-	beforeStop  []func() error
-	afterStart  []func() error
-	afterStop   []func() error
+	beforeStart []func(context.Context) error
+	beforeStop  []func(context.Context) error
+	afterStart  []func(context.Context) error
+	afterStop   []func(context.Context) error
 }
 
 // ID with service id.
@@ -101,28 +101,28 @@ func StopTimeout(t time.Duration) Option {
 // Before and Afters
 
 // BeforeStart run funcs before app starts
-func BeforeStart(fn func() error) Option {
+func BeforeStart(fn func(context.Context) error) Option {
 	return func(o *options) {
 		o.beforeStart = append(o.beforeStart, fn)
 	}
 }
 
 // BeforeStop run funcs before app stops
-func BeforeStop(fn func() error) Option {
+func BeforeStop(fn func(context.Context) error) Option {
 	return func(o *options) {
 		o.beforeStop = append(o.beforeStop, fn)
 	}
 }
 
 // AfterStart run funcs after app starts
-func AfterStart(fn func() error) Option {
+func AfterStart(fn func(context.Context) error) Option {
 	return func(o *options) {
 		o.afterStart = append(o.afterStart, fn)
 	}
 }
 
 // AfterStop run funcs after app stops
-func AfterStop(fn func() error) Option {
+func AfterStop(fn func(context.Context) error) Option {
 	return func(o *options) {
 		o.afterStop = append(o.afterStop, fn)
 	}

options_test.go

@@ -155,7 +155,7 @@ func TestStopTimeout(t *testing.T) {
 
 func TestBeforeStart(t *testing.T) {
 	o := &options{}
-	v := func() error {
+	v := func(_ context.Context) error {
 		t.Log("BeforeStart...")
 		return nil
 	}
@@ -164,7 +164,7 @@ func TestBeforeStart(t *testing.T) {
 
 func TestBeforeStop(t *testing.T) {
 	o := &options{}
-	v := func() error {
+	v := func(_ context.Context) error {
 		t.Log("BeforeStop...")
 		return nil
 	}
@@ -173,7 +173,7 @@ func TestBeforeStop(t *testing.T) {
 
 func TestAfterStart(t *testing.T) {
 	o := &options{}
-	v := func() error {
+	v := func(_ context.Context) error {
 		t.Log("AfterStart...")
 		return nil
 	}
@@ -182,7 +182,7 @@ func TestAfterStart(t *testing.T) {
 
 func TestAfterStop(t *testing.T) {
 	o := &options{}
-	v := func() error {
+	v := func(_ context.Context) error {
 		t.Log("AfterStop...")
 		return nil
 	}