github.com/lingyao2333/mo-zero@v1.4.1/rest/handler/maxconnshandler.go (about)

     1  package handler
     2  
     3  import (
     4  	"net/http"
     5  
     6  	"github.com/lingyao2333/mo-zero/core/logx"
     7  	"github.com/lingyao2333/mo-zero/core/syncx"
     8  	"github.com/lingyao2333/mo-zero/rest/internal"
     9  )
    10  
    11  // MaxConns returns a middleware that limit the concurrent connections.
    12  func MaxConns(n int) func(http.Handler) http.Handler {
    13  	if n <= 0 {
    14  		return func(next http.Handler) http.Handler {
    15  			return next
    16  		}
    17  	}
    18  
    19  	return func(next http.Handler) http.Handler {
    20  		latch := syncx.NewLimit(n)
    21  
    22  		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    23  			if latch.TryBorrow() {
    24  				defer func() {
    25  					if err := latch.Return(); err != nil {
    26  						logx.Error(err)
    27  					}
    28  				}()
    29  
    30  				next.ServeHTTP(w, r)
    31  			} else {
    32  				internal.Errorf(r, "concurrent connections over %d, rejected with code %d",
    33  					n, http.StatusServiceUnavailable)
    34  				w.WriteHeader(http.StatusServiceUnavailable)
    35  			}
    36  		})
    37  	}
    38  }