fix: update to avoid race condition

middleware/loadshedding/loadshedding.go

@@ -11,34 +11,34 @@ import (
 // If a request exceeds the specified timeout, a custom load-shedding handler is executed.
 func New(timeout time.Duration, loadSheddingHandler fiber.Handler, exclude func(fiber.Ctx) bool) fiber.Handler {
 	return func(c fiber.Ctx) error {
-		// Skip load-shedding for excluded requests
+		// Skip load-shedding logic for requests matching the exclusion criteria
 		if exclude != nil && exclude(c) {
 			return c.Next()
 		}
 
-		// Create a context with the specified timeout
+		// Create a context with a timeout for the current request
 		ctx, cancel := context.WithTimeout(c.Context(), timeout)
 		defer cancel()
 
-		// Channel to signal request completion
-		done := make(chan error, 1)
+		// Set the new context with a timeout
+		c.SetContext(ctx)
 
-		// Capture the current handler execution
-		handler := func() error {
-			return c.Next()
-		}
+		// Process the request and capture any error
+		err := c.Next()
+
+		// Create a channel to signal when request processing completes
+		done := make(chan error, 1)
 
-		// Process the handler in a separate goroutine
+		// Send the result of the request processing to the channel
 		go func() {
-			done <- handler()
+			done <- err
 		}()
 
+		// Handle either request completion or timeout
 		select {
-		case <-ctx.Done():
-			// Timeout occurred; invoke the load-shedding handler
+		case <-ctx.Done(): // Triggered if the timeout expires
 			return loadSheddingHandler(c)
-		case err := <-done:
-			// Request completed successfully; return any handler error
+		case err := <-done: // Triggered if request processing completes
 			return err
 		}
 	}

middleware/loadshedding/loadshedding_test.go

@@ -53,20 +53,24 @@ func Test_LoadSheddingTimeout(t *testing.T) {
 	t.Parallel()
 	app := fiber.New()
 
-	// Middleware without exclusions
+	// Middleware with a 1-second timeout
 	app.Use(loadshedding.New(
-		1*time.Second, // Set timeout for the middleware
+		1*time.Second, // Middleware timeout
 		loadSheddingHandler,
 		nil,
 	))
-	app.Get("/", timeoutHandler)
+	app.Get("/", func(c fiber.Ctx) error {
+		// Simulate long-running request without creating goroutines
+		time.Sleep(2 * time.Second)
+		return c.SendString("This should not appear")
+	})
 
-	// Create a custom HTTP client with a sufficient timeout
+	// Create a custom request
+	req := httptest.NewRequest(fiber.MethodGet, "/", nil)
 
 	// Test timeout behavior
-	req := httptest.NewRequest(fiber.MethodGet, "/", nil)
 	resp, err := app.Test(req, fiber.TestConfig{
-		Timeout: 3 * time.Second,
+		Timeout: 3 * time.Second, // Ensure the test timeout exceeds middleware timeout
 	})
 	require.NoError(t, err)
 	require.Equal(t, fiber.StatusServiceUnavailable, resp.StatusCode)