github.com/gofiber/fiber/v2@v2.47.0/middleware/timeout/timeout.go (about)

     1  package timeout
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"log"
     7  	"sync"
     8  	"time"
     9  
    10  	"github.com/gofiber/fiber/v2"
    11  )
    12  
    13  var once sync.Once
    14  
    15  // New wraps a handler and aborts the process of the handler if the timeout is reached.
    16  //
    17  // Deprecated: This implementation contains data race issues. Use NewWithContext instead.
    18  // Find documentation and sample usage on https://docs.gofiber.io/api/middleware/timeout
    19  func New(handler fiber.Handler, timeout time.Duration) fiber.Handler {
    20  	once.Do(func() {
    21  		log.Printf("[Warning] - [TIMEOUT] timeout contains data race issues, not ready for production!")
    22  	})
    23  
    24  	if timeout <= 0 {
    25  		return handler
    26  	}
    27  
    28  	// logic is from fasthttp.TimeoutWithCodeHandler https://github.com/valyala/fasthttp/blob/master/server.go#L418
    29  	return func(ctx *fiber.Ctx) error {
    30  		ch := make(chan struct{}, 1)
    31  
    32  		go func() {
    33  			defer func() {
    34  				if err := recover(); err != nil {
    35  					log.Printf("[Warning] - [TIMEOUT] recover error %v", err)
    36  				}
    37  			}()
    38  			if err := handler(ctx); err != nil {
    39  				log.Printf("[Warning] - [TIMEOUT] handler error %v", err)
    40  			}
    41  			ch <- struct{}{}
    42  		}()
    43  
    44  		select {
    45  		case <-ch:
    46  		case <-time.After(timeout):
    47  			return fiber.ErrRequestTimeout
    48  		}
    49  
    50  		return nil
    51  	}
    52  }
    53  
    54  // NewWithContext implementation of timeout middleware. Set custom errors(context.DeadlineExceeded vs) for get fiber.ErrRequestTimeout response.
    55  func NewWithContext(h fiber.Handler, t time.Duration, tErrs ...error) fiber.Handler {
    56  	return func(ctx *fiber.Ctx) error {
    57  		timeoutContext, cancel := context.WithTimeout(ctx.UserContext(), t)
    58  		defer cancel()
    59  		ctx.SetUserContext(timeoutContext)
    60  		if err := h(ctx); err != nil {
    61  			if errors.Is(err, context.DeadlineExceeded) {
    62  				return fiber.ErrRequestTimeout
    63  			}
    64  			for i := range tErrs {
    65  				if errors.Is(err, tErrs[i]) {
    66  					return fiber.ErrRequestTimeout
    67  				}
    68  			}
    69  			return err
    70  		}
    71  		return nil
    72  	}
    73  }