github.com/sohaha/zlsgo@v1.7.13-0.20240501141223-10dd1a906f76/znet/timeout/timeout.go (about)

     1  package timeout
     2  
     3  import (
     4  	"context"
     5  	"net/http"
     6  	"time"
     7  
     8  	"github.com/sohaha/zlsgo/znet"
     9  )
    10  
    11  func New(waitingTime time.Duration, custom ...znet.HandlerFunc) znet.HandlerFunc {
    12  	return func(c *znet.Context) {
    13  		ctx, cancel := context.WithTimeout(c.Request.Context(), waitingTime)
    14  		defer cancel()
    15  		done := make(chan struct{}, 1)
    16  		go func() {
    17  			c.Next()
    18  			done <- struct{}{}
    19  		}()
    20  		for {
    21  			select {
    22  			case <-done:
    23  				return
    24  			case <-ctx.Done():
    25  				if len(custom) > 0 {
    26  					custom[0](c)
    27  				} else {
    28  					c.Abort(http.StatusGatewayTimeout)
    29  				}
    30  				return
    31  			}
    32  		}
    33  	}
    34  }