Groovy web console

subscribe to the feed Subscribe
to this
site
Solution (via #groovywebconsole)
tweet this snippet Tweet
this
script

Solution

Published 2 months ago by Guo with tags RateLimiter
Actions  ➤ Edit in console Back to console Show/hide line numbers View recent scripts
/*
     _______.  ______    __       __    __  .___________. __    ______   .__   __.
    /       | /  __  \  |  |     |  |  |  | |           ||  |  /  __  \  |  \ |  |
   |   (----`|  |  |  | |  |     |  |  |  | `---|  |----`|  | |  |  |  | |   \|  |
    \   \    |  |  |  | |  |     |  |  |  |     |  |     |  | |  |  |  | |  . `  |
.----)   |   |  `--'  | |  `----.|  `--'  |     |  |     |  | |  `--'  | |  |\   |
|_______/     \______/  |_______| \______/      |__|     |__|  \______/  |__| \__|

 */


class RateLimiter {
    int limitIntervalInMillSec
    int limitCountPerInterval
    def rateLimitStore = [:]

    synchronized boolean tryAcquire(String requesterId) {
        List requestHistory = retriveRequesterHistory(requesterId)

        def currentTimeInSeconds = System.currentTimeMillis()
        removeExpiredHistory(requestHistory, currentTimeInSeconds)

        if (requestHistory.size() < limitCountPerInterval) {
            requestHistory << currentTimeInSeconds
            true
        } else {
            false
        }
    }

    private void removeExpiredHistory(List requestHistory, currentTimeInSeconds) {
        requestHistory.removeAll {
            currentTimeInSeconds - it > limitIntervalInMillSec
        }
    }

    private List retriveRequesterHistory(String requesterId) {
        List requestHistory = rateLimitStore[requesterId]
        if (requestHistory == null) {
            requestHistory = []
            rateLimitStore[requesterId] = requestHistory
        }
        requestHistory
    }

}


class API {
    RateLimiter rateLimiter

    boolean service1(String customer) {
        if (rateLimiter.tryAcquire(customer)) {
            println "\t${new Date()} service1 called by $customer"
            true
        } else {
            println "\t${new Date()} service1 rejected by $customer"
            false
        }
    }

    boolean service2(String customer) {
        if (rateLimiter.tryAcquire(customer)) {
            println "\t${new Date()} service2 called by $customer"
            true
        } else {
            println "\t${new Date()} service2 rejected by $customer"
            false
        }
    }
}


/*
.___________. _______     _______.___________.    _______.
|           ||   ____|   /       |           |   /       |
`---|  |----`|  |__     |   (----`---|  |----`  |   (----`
    |  |     |   __|     \   \       |  |        \   \
    |  |     |  |____.----)   |      |  |    .----)   |
    |__|     |_______|_______/       |__|    |_______/
 */

def EXPIRED_TIMESTAMP=0
def limitIntervalInMillSec = 1000
def limitCountPerInterval = 5

def rateLimiter = new RateLimiter(limitIntervalInMillSec: limitIntervalInMillSec, limitCountPerInterval: limitCountPerInterval);
{->
    println "RATE LIMITER TEST: should acquire success before reach limit"
    (1..limitCountPerInterval).each {
        assert rateLimiter.tryAcquire("rateLimiterTest")
    }

    println "RATE LIMITER TEST: should acquire fail after reach limit immediately"
    assert !rateLimiter.tryAcquire("rateLimiterTest")


    println "RATE LIMITER TEST: should acquire success again after first request expired"
    rateLimiter.rateLimitStore["rateLimiterTest"][0]=EXPIRED_TIMESTAMP
    assert rateLimiter.tryAcquire("rateLimiterTest")
}()



def api = new API(rateLimiter: rateLimiter);
{->
    println "API TEST: should call service1 success before reach limit "
    (1..limitCountPerInterval).each {
        assert api.service1("apiTest")
    }

    println "API TEST: should fail to call any service after reach limit immediately"
    assert !api.service1("apiTest")
    assert !api.service2("apiTest")


    println "API TEST: should call any success again after requests expired"
    rateLimiter.rateLimitStore["apiTest"][0,1]=[EXPIRED_TIMESTAMP,EXPIRED_TIMESTAMP]
    assert api.service1("apiTest")
    assert api.service2("apiTest")
}()