Search code examples
gomux

Return error when a handler reaches it max client limit


I have written a small wrapper function which use counting semaphore concept to limit number of connections to a particular handler (as this handler is resource consuming). Below is the code which achieve the same.

func LimitNumClients(f http.HandlerFunc, maxClients int) http.HandlerFunc {
    // Counting semaphore using a buffered channel
    sema := make(chan struct{}, maxClients)

    return func(w http.ResponseWriter, req *http.Request) {
        sema <- struct{}{}
        defer func() { <-sema }()
        f(w, req)
    }
}

And then wrapper it up in the handler as below

Route{
        "Test",
        "GET",
        /test,
        LimitNumClients(testhandler, 5),
    },

Now I want to reply back with 501 error when the client limit is reached for any new connection. How to achieve the same.


Solution

  • You may use a non-blocking send operation. If that succeeds, continue as normal, if sending on sema would block, then send back an error and return from the limiter without calling the handler:

    return func(w http.ResponseWriter, req *http.Request) {
        select {
        case sema <- struct{}{}:
        default:
            http.Error(w, "rate limit reached", 501)
            return
        }
    
        defer func() { <-sema }()
        f(w, req)
    }
    

    Also note that to signal a "rate limit reached" error, the returned status code should be HTTP 429 Too Many Requests, see RFC 6585.

    So instead return this:

    http.Error(w, "rate limit reached", http.StatusTooManyRequests)