refactor: improve OnStart hook context, move logic to method, and test DI access

Author: HeerakKashyap

examples/http-server/main.go

@@ -28,11 +28,11 @@ func main() {
 	a.GET("/trace", TraceHandler)
 	a.GET("/mysql", MysqlHandler)
 
-	// Register an OnStart Hook
 	a.OnStart(func(ctx *gofr.Context) error {
-		// Try accessing a DI-managed service, e.g., DB or HTTPService
+
 		fmt.Println("OnStart hook executed!")
-		fmt.Printf("DB: %#v\n", ctx.Container.SQL)
+		fmt.Printf("SQL: %#v\n", ctx.Container.SQL)
+		fmt.Printf("Redis: %#v\n", ctx.Container.Redis)
 		return nil
 	})
 

pkg/gofr/gofr.go

@@ -49,6 +49,26 @@ type App struct {
 	onStartHooks        []func(ctx *Context) error
 }
 
+func (a *App) newContextForHooks(ctx context.Context) *Context {
+	logger := logging.NewContextLogger(ctx, a.container.Logger)
+	return &Context{
+		Context:       ctx,
+		Container:     a.container,
+		Request:       noopRequest{},
+		ContextLogger: *logger,
+	}
+}
+
+func (a *App) runOnStartHooks(ctx context.Context) {
+	gofrCtx := a.newContextForHooks(ctx)
+	for _, hook := range a.onStartHooks {
+		if err := hook(gofrCtx); err != nil {
+			a.Logger().Errorf("OnStart hook failed: %v", err)
+			os.Exit(1)
+		}
+	}
+}
+
 // Shutdown stops the service(s) and close the application.
 // It shuts down the HTTP, gRPC, Metrics servers and closes the container's active connections to datasources.
 func (a *App) Shutdown(ctx context.Context) error {

pkg/gofr/run.go

@@ -8,9 +8,6 @@ import (
 	"sync"
 	"syscall"
 
-	"go.opentelemetry.io/otel"
-
-	"gofr.dev/pkg/gofr/logging"
 )
 
 // Run starts the application. If it is an HTTP server, it will start the server.
@@ -19,30 +16,13 @@ func (a *App) Run() {
 		a.cmd.Run(a.container)
 	}
 
-	// Create a context that is canceled on receiving termination signals
+	
 	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGINT, syscall.SIGTERM)
 	defer stop()
 
-	// --- Add this block for OnStart hooks ---
-	onStartCtx, span := otel.GetTracerProvider().Tracer("gofr-onstart").Start(context.Background(), "OnStart")
-	defer span.End()
-
-	logger := logging.NewContextLogger(onStartCtx, a.container.Logger)
-
-	gofrCtx := &Context{
-		Context:       onStartCtx,
-		Container:     a.container,
-		Request:       noopRequest{}, // Make sure noopRequest{} is available
-		ContextLogger: *logger,
-	}
-
-	for _, hook := range a.onStartHooks {
-		if err := hook(gofrCtx); err != nil {
-			a.Logger().Errorf("OnStart hook failed: %v", err)
-			os.Exit(1)
-		}
-	}
-	// --- End OnStart block ---
+	onStartCtx := context.WithoutCancel(ctx)
+    a.runOnStartHooks(onStartCtx)
+	
 
 	timeout, err := getShutdownTimeoutFromConfig(a.Config)
 	if err != nil {