/**
* @see io.apiman.gateway.engine.components.IRateLimiterComponent#accept(java.lang.String,
* io.apiman.gateway.engine.rates.RateBucketPeriod, int,
* io.apiman.gateway.engine.async.IAsyncResultHandler)
*/
@Override
public void accept(
String bucketId,
RateBucketPeriod period,
int limit,
IAsyncResultHandler<RateLimitResponse> handler) {
RateLimiterBucket bucket = null;
synchronized (mutex) {
bucket = (RateLimiterBucket) getCache().get(bucketId);
if (bucket == null) {
bucket = new RateLimiterBucket();
getCache().put(bucketId, bucket);
}
bucket.resetIfNecessary(period);
RateLimitResponse response = new RateLimitResponse();
if (bucket.getCount() >= limit) {
response.setAccepted(false);
} else {
bucket.setCount(bucket.getCount() + 1);
bucket.setLast(System.currentTimeMillis());
response.setAccepted(true);
}
int reset = (int) (bucket.getResetMillis(period) / 1000L);
response.setReset(reset);
response.setRemaining(limit - bucket.getCount());
handler.handle(AsyncResultImpl.<RateLimitResponse>create(response));
getCache().put(bucketId, bucket);
}
}
/**
* Update the bucket in ES and then return the rate limit response to the original handler. If the
* update fails because we have a stale version, then try the whole thing again (because we
* conflicted with another request).
*
* @param id
* @param bucket
* @param rlr
* @param version
* @param limit
* @param period
* @param bucketId
* @param increment
* @param handler
*/
protected void updateBucketAndReturn(
final String id,
final RateLimiterBucket bucket,
final RateLimitResponse rlr,
final long version,
final String bucketId,
final RateBucketPeriod period,
final long limit,
final long increment,
final IAsyncResultHandler<RateLimitResponse> handler) {
Index index =
new Index.Builder(bucket)
.refresh(false)
.index(getIndexName())
.setParameter(Parameters.OP_TYPE, "index") // $NON-NLS-1$
.setParameter(Parameters.VERSION, String.valueOf(version))
.type("rateBucket")
.id(id)
.build(); //$NON-NLS-1$
try {
getClient().execute(index);
handler.handle(AsyncResultImpl.create(rlr));
} catch (Throwable e) {
// FIXME need to fix this now that we've switched to jest!
// if (ESUtils.rootCause(e) instanceof VersionConflictEngineException) {
// // If we got a version conflict, then it means some other request
// // managed to update the ES document since we retrieved it. Therefore
// // everything we've done is out of date, so we should do it all
// // over again.
// accept(bucketId, period, limit, increment, handler);
// } else {
handler.handle(AsyncResultImpl.<RateLimitResponse>create(e));
// }
}
}
/**
* @see io.apiman.gateway.engine.components.IRateLimiterComponent#accept(java.lang.String,
* io.apiman.gateway.engine.rates.RateBucketPeriod, long, long,
* io.apiman.gateway.engine.async.IAsyncResultHandler)
*/
@Override
public void accept(
final String bucketId,
final RateBucketPeriod period,
final long limit,
final long increment,
final IAsyncResultHandler<RateLimitResponse> handler) {
final String id = id(bucketId);
try {
Get get = new Get.Builder(getIndexName(), id).type("rateBucket").build(); // $NON-NLS-1$
JestResult result = getClient().execute(get);
RateLimiterBucket bucket;
long version;
if (result.isSucceeded()) {
// use the existing bucket
version = result.getJsonObject().get("_version").getAsLong(); // $NON-NLS-1$
bucket = result.getSourceAsObject(RateLimiterBucket.class);
} else {
// make a new bucket
version = 0;
bucket = new RateLimiterBucket();
}
bucket.resetIfNecessary(period);
final RateLimitResponse rlr = new RateLimitResponse();
if (bucket.getCount() > limit) {
rlr.setAccepted(false);
} else {
rlr.setAccepted(bucket.getCount() < limit);
bucket.setCount(bucket.getCount() + increment);
bucket.setLast(System.currentTimeMillis());
}
int reset = (int) (bucket.getResetMillis(period) / 1000L);
rlr.setReset(reset);
rlr.setRemaining(limit - bucket.getCount());
updateBucketAndReturn(id, bucket, rlr, version, bucketId, period, limit, increment, handler);
} catch (Throwable e) {
handler.handle(AsyncResultImpl.create(e, RateLimitResponse.class));
}
}
@Override
public void end() {
responseHandler.handle(AsyncResultImpl.<IApiConnectionResponse>create(this));
}