dendi-wang/thingsboard
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

New Buffered Rate Limit implementation

Browse Source
This commit is contained in:
Andrew Shvayka 2018-10-24 19:28:11 +03:00
parent c1675db1a1
commit 571f96c4a9
11 changed files with 460 additions and 329 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
application/src/main/resources/thingsboard.yml
View File

@ -140,7 +140,7 @@ cassandra:
buffer_size: "${CASSANDRA_QUERY_BUFFER_SIZE:200000}"
concurrent_limit: "${CASSANDRA_QUERY_CONCURRENT_LIMIT:1000}"
permit_max_wait_time: "${PERMIT_MAX_WAIT_TIME:120000}"
rate_limit_print_interval_ms: "${CASSANDRA_QUERY_RATE_LIMIT_PRINT_MS:30000}"
rate_limit_print_interval_ms: "${CASSANDRA_QUERY_RATE_LIMIT_PRINT_MS:10000}"
# SQL configuration parameters
sql:

5
dao/src/main/java/org/thingsboard/server/dao/nosql/CassandraAbstractDao.java
View File

@ -35,7 +35,6 @@ import org.thingsboard.server.dao.model.type.ComponentTypeCodec;
import org.thingsboard.server.dao.model.type.DeviceCredentialsTypeCodec;
import org.thingsboard.server.dao.model.type.EntityTypeCodec;
import org.thingsboard.server.dao.model.type.JsonCodec;
import org.thingsboard.server.dao.util.BufferedRateLimiter;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;
@ -49,7 +48,7 @@ public abstract class CassandraAbstractDao {
private ConcurrentMap<String, PreparedStatement> preparedStatementMap = new ConcurrentHashMap<>();
@Autowired
private BufferedRateLimiter rateLimiter;
private CassandraBufferedRateExecutor rateLimiter;
private Session session;
@ -115,7 +114,7 @@ public abstract class CassandraAbstractDao {
if (statement.getConsistencyLevel() == null) {
statement.setConsistencyLevel(level);
}
return new RateLimitedResultSetFuture(getSession(), rateLimiter, statement);
return rateLimiter.submit(new CassandraStatementTask(getSession(), statement));
}
private static String statementToString(Statement statement) {

78
dao/src/main/java/org/thingsboard/server/dao/nosql/CassandraBufferedRateExecutor.java Normal file
View File

@ -0,0 +1,78 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.nosql;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.ResultSetFuture;
import com.google.common.util.concurrent.SettableFuture;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.thingsboard.server.dao.nosql.tmp.AbstractBufferedRateExecutor;
import org.thingsboard.server.dao.nosql.tmp.AsyncTaskContext;
import org.thingsboard.server.dao.util.NoSqlAnyDao;
import javax.annotation.PreDestroy;
/**
* Created by ashvayka on 24.10.18.
*/
@Component
@Slf4j
@NoSqlAnyDao
public class CassandraBufferedRateExecutor extends AbstractBufferedRateExecutor<CassandraStatementTask, ResultSetFuture, ResultSet> {
public CassandraBufferedRateExecutor(
@Value("${cassandra.query.buffer_size}") int queueLimit,
@Value("${cassandra.query.concurrent_limit}") int concurrencyLimit,
@Value("${cassandra.query.permit_max_wait_time}") long maxWaitTime,
@Value("${cassandra.query.dispatcher_threads:2}") int dispatcherThreads,
@Value("${cassandra.query.callback_threads:2}") int callbackThreads,
@Value("${cassandra.query.poll_ms:50}") long pollMs) {
super(queueLimit, concurrencyLimit, maxWaitTime, dispatcherThreads, callbackThreads, pollMs);
}
@Scheduled(fixedDelayString = "${cassandra.query.rate_limit_print_interval_ms}")
public void printStats() {
log.info("Permits totalAdded [{}] totalLaunched [{}] totalReleased [{}] totalFailed [{}] totalExpired [{}] totalRejected [{}] currBuffer [{}] ",
totalAdded.getAndSet(0), totalLaunched.getAndSet(0), totalReleased.getAndSet(0),
totalFailed.getAndSet(0), totalExpired.getAndSet(0), totalRejected.getAndSet(0),
concurrencyLevel.get());
}
@PreDestroy
public void stop() {
super.stop();
}
@Override
protected SettableFuture<ResultSet> create() {
return SettableFuture.create();
}
@Override
protected ResultSetFuture wrap(CassandraStatementTask task, SettableFuture<ResultSet> future) {
return new TbResultSetFuture(future);
}
@Override
protected ResultSetFuture execute(AsyncTaskContext<CassandraStatementTask, ResultSet> taskCtx) {
CassandraStatementTask task = taskCtx.getTask();
return task.getSession().executeAsync(task.getStatement());
}
}

32
dao/src/main/java/org/thingsboard/server/dao/nosql/CassandraStatementTask.java Normal file
View File

@ -0,0 +1,32 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.nosql;
import com.datastax.driver.core.Session;
import com.datastax.driver.core.Statement;
import lombok.Data;
import org.thingsboard.server.dao.nosql.tmp.AsyncTask;
/**
* Created by ashvayka on 24.10.18.
*/
@Data
public class CassandraStatementTask implements AsyncTask {
private final Session session;
private final Statement statement;
}

94
dao/src/main/java/org/thingsboard/server/dao/nosql/TbResultSetFuture.java Normal file
View File

@ -0,0 +1,94 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.nosql;
import com.datastax.driver.core.ResultSet;
import com.datastax.driver.core.ResultSetFuture;
import com.google.common.util.concurrent.SettableFuture;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
/**
* Created by ashvayka on 24.10.18.
*/
public class TbResultSetFuture implements ResultSetFuture {
private final SettableFuture<ResultSet> mainFuture;
public TbResultSetFuture(SettableFuture<ResultSet> mainFuture) {
this.mainFuture = mainFuture;
}
@Override
public ResultSet getUninterruptibly() {
return getSafe();
}
@Override
public ResultSet getUninterruptibly(long timeout, TimeUnit unit) throws TimeoutException {
return getSafe(timeout, unit);
}
@Override
public boolean cancel(boolean mayInterruptIfRunning) {
return mainFuture.cancel(mayInterruptIfRunning);
}
@Override
public boolean isCancelled() {
return mainFuture.isCancelled();
}
@Override
public boolean isDone() {
return mainFuture.isDone();
}
@Override
public ResultSet get() throws InterruptedException, ExecutionException {
return mainFuture.get();
}
@Override
public ResultSet get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException {
return mainFuture.get(timeout, unit);
}
@Override
public void addListener(Runnable listener, Executor executor) {
mainFuture.addListener(listener, executor);
}
private ResultSet getSafe() {
try {
return mainFuture.get();
} catch (InterruptedException | ExecutionException e) {
throw new IllegalStateException(e);
}
}
private ResultSet getSafe(long timeout, TimeUnit unit) throws TimeoutException {
try {
return mainFuture.get(timeout, unit);
} catch (InterruptedException | ExecutionException e) {
throw new IllegalStateException(e);
}
}
}

169
dao/src/main/java/org/thingsboard/server/dao/nosql/tmp/AbstractBufferedRateExecutor.java Normal file
View File

@ -0,0 +1,169 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.nosql.tmp;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.SettableFuture;
import lombok.extern.slf4j.Slf4j;
import javax.annotation.Nullable;
import java.util.UUID;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingDeque;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.TimeoutException;
import java.util.concurrent.atomic.AtomicInteger;
/**
* Created by ashvayka on 24.10.18.
*/
@Slf4j
public abstract class AbstractBufferedRateExecutor<T extends AsyncTask, F extends ListenableFuture<V>, V> implements BufferedRateExecutor<T, F> {
private final long maxWaitTime;
private final long pollMs;
private final BlockingQueue<AsyncTaskContext<T, V>> queue;
private final ExecutorService dispatcherExecutor;
private final ExecutorService callbackExecutor;
private final ScheduledExecutorService timeoutExecutor;
private final int concurrencyLimit;
protected final AtomicInteger concurrencyLevel = new AtomicInteger();
protected final AtomicInteger totalAdded = new AtomicInteger();
protected final AtomicInteger totalLaunched = new AtomicInteger();
protected final AtomicInteger totalReleased = new AtomicInteger();
protected final AtomicInteger totalFailed = new AtomicInteger();
protected final AtomicInteger totalExpired = new AtomicInteger();
protected final AtomicInteger totalRejected = new AtomicInteger();
public AbstractBufferedRateExecutor(int queueLimit, int concurrencyLimit, long maxWaitTime, int dispatcherThreads, int callbackThreads, long pollMs) {
this.maxWaitTime = maxWaitTime;
this.pollMs = pollMs;
this.concurrencyLimit = concurrencyLimit;
this.queue = new LinkedBlockingDeque<>(queueLimit);
this.dispatcherExecutor = Executors.newFixedThreadPool(dispatcherThreads);
this.callbackExecutor = Executors.newFixedThreadPool(callbackThreads);
this.timeoutExecutor = Executors.newSingleThreadScheduledExecutor();
for (int i = 0; i < dispatcherThreads; i++) {
dispatcherExecutor.submit(this::dispatch);
}
}
@Override
public F submit(T task) {
SettableFuture<V> settableFuture = create();
F result = wrap(task, settableFuture);
try {
totalAdded.incrementAndGet();
queue.add(new AsyncTaskContext<>(UUID.randomUUID(), task, settableFuture, System.currentTimeMillis()));
} catch (IllegalStateException e) {
totalRejected.incrementAndGet();
settableFuture.setException(e);
}
return result;
}
public void stop() {
if (dispatcherExecutor != null) {
dispatcherExecutor.shutdownNow();
}
if (callbackExecutor != null) {
callbackExecutor.shutdownNow();
}
if (timeoutExecutor != null) {
timeoutExecutor.shutdownNow();
}
}
protected abstract SettableFuture<V> create();
protected abstract F wrap(T task, SettableFuture<V> future);
protected abstract ListenableFuture<V> execute(AsyncTaskContext<T, V> taskCtx);
private void dispatch() {
log.info("Buffered rate executor thread started");
while (!Thread.interrupted()) {
int curLvl = concurrencyLevel.get();
AsyncTaskContext<T, V> taskCtx = null;
try {
if (curLvl <= concurrencyLimit) {
taskCtx = queue.take();
final AsyncTaskContext<T, V> finalTaskCtx = taskCtx;
logTask("Processing", finalTaskCtx);
concurrencyLevel.incrementAndGet();
long timeout = finalTaskCtx.getCreateTime() + maxWaitTime - System.currentTimeMillis();
if (timeout > 0) {
totalLaunched.incrementAndGet();
ListenableFuture<V> result = execute(finalTaskCtx);
result = Futures.withTimeout(result, timeout, TimeUnit.MILLISECONDS, timeoutExecutor);
Futures.addCallback(result, new FutureCallback<V>() {
@Override
public void onSuccess(@Nullable V result) {
logTask("Releasing", finalTaskCtx);
totalReleased.incrementAndGet();
concurrencyLevel.decrementAndGet();
finalTaskCtx.getFuture().set(result);
}
@Override
public void onFailure(Throwable t) {
if (t instanceof TimeoutException) {
logTask("Expired During Execution", finalTaskCtx);
} else {
logTask("Failed", finalTaskCtx);
}
totalFailed.incrementAndGet();
concurrencyLevel.decrementAndGet();
finalTaskCtx.getFuture().setException(t);
log.debug("[{}] Failed to execute task: {}", finalTaskCtx.getId(), finalTaskCtx.getTask(), t);
}
}, callbackExecutor);
} else {
logTask("Expired Before Execution", finalTaskCtx);
totalExpired.incrementAndGet();
concurrencyLevel.decrementAndGet();
taskCtx.getFuture().setException(new TimeoutException());
}
} else {
Thread.sleep(pollMs);
}
} catch (Throwable e) {
if (taskCtx != null) {
log.debug("[{}] Failed to execute task: {}", taskCtx.getId(), taskCtx, e);
totalFailed.incrementAndGet();
concurrencyLevel.decrementAndGet();
} else {
log.debug("Failed to queue task:", e);
}
}
}
log.info("Buffered rate executor thread stopped");
}
private void logTask(String action, AsyncTaskContext<T, V> taskCtx) {
if (log.isTraceEnabled()) {
log.trace("[{}] {} task: {}", taskCtx.getId(), action, taskCtx);
} else {
log.debug("[{}] {} task", taskCtx.getId(), action);
}
}
}

22
dao/src/main/java/org/thingsboard/server/dao/nosql/tmp/AsyncTask.java Normal file
View File

@ -0,0 +1,22 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.nosql.tmp;
/**
* Created by ashvayka on 24.10.18.
*/
public interface AsyncTask {
}

34
dao/src/main/java/org/thingsboard/server/dao/nosql/tmp/AsyncTaskContext.java Normal file
View File

@ -0,0 +1,34 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.nosql.tmp;
import com.google.common.util.concurrent.SettableFuture;
import lombok.Data;
import java.util.UUID;
/**
* Created by ashvayka on 24.10.18.
*/
@Data
public class AsyncTaskContext<T extends AsyncTask, V> {
private final UUID id;
private final T task;
private final SettableFuture<V> future;
private final long createTime;
}

27
dao/src/main/java/org/thingsboard/server/dao/nosql/tmp/BufferedRateExecutor.java Normal file
View File

@ -0,0 +1,27 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.nosql.tmp;
import com.google.common.util.concurrent.ListenableFuture;
/**
* Created by ashvayka on 24.10.18.
*/
public interface BufferedRateExecutor<T extends AsyncTask, F extends ListenableFuture> {
F submit(T task);
}

182
dao/src/main/java/org/thingsboard/server/dao/util/BufferedRateLimiter.java
View File

@ -1,182 +0,0 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.util;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.scheduling.annotation.Scheduled;
import org.springframework.stereotype.Component;
import org.thingsboard.server.dao.exception.BufferLimitException;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.Executors;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
@Component
@Slf4j
@NoSqlAnyDao
public class BufferedRateLimiter implements AsyncRateLimiter {
private final ListeningExecutorService pool = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(10));
private final int permitsLimit;
private final int maxPermitWaitTime;
private final AtomicInteger permits;
private final BlockingQueue<LockedFuture> queue;
private final AtomicInteger maxQueueSize = new AtomicInteger();
private final AtomicInteger maxGrantedPermissions = new AtomicInteger();
private final AtomicInteger totalGranted = new AtomicInteger();
private final AtomicInteger totalReleased = new AtomicInteger();
private final AtomicInteger totalRequested = new AtomicInteger();
public BufferedRateLimiter(@Value("${cassandra.query.buffer_size}") int queueLimit,
@Value("${cassandra.query.concurrent_limit}") int permitsLimit,
@Value("${cassandra.query.permit_max_wait_time}") int maxPermitWaitTime) {
this.permitsLimit = permitsLimit;
this.maxPermitWaitTime = maxPermitWaitTime;
this.permits = new AtomicInteger();
this.queue = new LinkedBlockingQueue<>(queueLimit);
}
@Override
public ListenableFuture<Void> acquireAsync() {
totalRequested.incrementAndGet();
if (queue.isEmpty()) {
if (permits.incrementAndGet() <= permitsLimit) {
if (permits.get() > maxGrantedPermissions.get()) {
maxGrantedPermissions.set(permits.get());
}
totalGranted.incrementAndGet();
return Futures.immediateFuture(null);
}
permits.decrementAndGet();
}
return putInQueue();
}
@Override
public void release() {
permits.decrementAndGet();
totalReleased.incrementAndGet();
reprocessQueue();
}
private void reprocessQueue() {
while (permits.get() < permitsLimit) {
if (permits.incrementAndGet() <= permitsLimit) {
if (permits.get() > maxGrantedPermissions.get()) {
maxGrantedPermissions.set(permits.get());
}
LockedFuture lockedFuture = queue.poll();
if (lockedFuture != null) {
totalGranted.incrementAndGet();
lockedFuture.latch.countDown();
} else {
permits.decrementAndGet();
break;
}
} else {
permits.decrementAndGet();
}
}
}
private LockedFuture createLockedFuture() {
CountDownLatch latch = new CountDownLatch(1);
ListenableFuture<Void> future = pool.submit(() -> {
latch.await();
return null;
});
return new LockedFuture(latch, future, System.currentTimeMillis());
}
private ListenableFuture<Void> putInQueue() {
int size = queue.size();
if (size > maxQueueSize.get()) {
maxQueueSize.set(size);
}
if (queue.remainingCapacity() > 0) {
try {
LockedFuture lockedFuture = createLockedFuture();
if (!queue.offer(lockedFuture, 1, TimeUnit.SECONDS)) {
lockedFuture.cancelFuture();
return Futures.immediateFailedFuture(new BufferLimitException());
}
if(permits.get() < permitsLimit) {
reprocessQueue();
}
if(permits.get() < permitsLimit) {
reprocessQueue();
}
return lockedFuture.future;
} catch (InterruptedException e) {
return Futures.immediateFailedFuture(new BufferLimitException());
}
}
return Futures.immediateFailedFuture(new BufferLimitException());
}
@Scheduled(fixedDelayString = "${cassandra.query.rate_limit_print_interval_ms}")
public void printStats() {
int expiredCount = 0;
for (LockedFuture lockedFuture : queue) {
if (lockedFuture.isExpired()) {
lockedFuture.cancelFuture();
expiredCount++;
}
}
log.info("Permits maxBuffer [{}] maxPermits [{}] expired [{}] currPermits [{}] currBuffer [{}] " +
"totalPermits [{}] totalRequests [{}] totalReleased [{}]",
maxQueueSize.getAndSet(0), maxGrantedPermissions.getAndSet(0), expiredCount,
permits.get(), queue.size(),
totalGranted.getAndSet(0), totalRequested.getAndSet(0), totalReleased.getAndSet(0));
}
private class LockedFuture {
final CountDownLatch latch;
final ListenableFuture<Void> future;
final long createTime;
public LockedFuture(CountDownLatch latch, ListenableFuture<Void> future, long createTime) {
this.latch = latch;
this.future = future;
this.createTime = createTime;
}
void cancelFuture() {
future.cancel(false);
latch.countDown();
}
boolean isExpired() {
return (System.currentTimeMillis() - createTime) > maxPermitWaitTime;
}
}
}

142
dao/src/test/java/org/thingsboard/server/dao/util/BufferedRateLimiterTest.java
View File

@ -1,142 +0,0 @@
/**
* Copyright © 2016-2018 The Thingsboard Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.thingsboard.server.dao.util;
import com.google.common.util.concurrent.FutureCallback;
import com.google.common.util.concurrent.Futures;
import com.google.common.util.concurrent.ListenableFuture;
import com.google.common.util.concurrent.ListeningExecutorService;
import com.google.common.util.concurrent.MoreExecutors;
import org.junit.Test;
import org.thingsboard.server.dao.exception.BufferLimitException;
import javax.annotation.Nullable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.Executors;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertFalse;
import static org.junit.Assert.assertTrue;
import static org.junit.Assert.fail;
public class BufferedRateLimiterTest {
@Test
public void finishedFutureReturnedIfPermitsAreGranted() {
BufferedRateLimiter limiter = new BufferedRateLimiter(10, 10, 100);
ListenableFuture<Void> actual = limiter.acquireAsync();
assertTrue(actual.isDone());
}
@Test
public void notFinishedFutureReturnedIfPermitsAreNotGranted() {
BufferedRateLimiter limiter = new BufferedRateLimiter(10, 1, 100);
ListenableFuture<Void> actual1 = limiter.acquireAsync();
ListenableFuture<Void> actual2 = limiter.acquireAsync();
assertTrue(actual1.isDone());
assertFalse(actual2.isDone());
}
@Test
public void failedFutureReturnedIfQueueIsfull() {
BufferedRateLimiter limiter = new BufferedRateLimiter(1, 1, 100);
ListenableFuture<Void> actual1 = limiter.acquireAsync();
ListenableFuture<Void> actual2 = limiter.acquireAsync();
ListenableFuture<Void> actual3 = limiter.acquireAsync();
assertTrue(actual1.isDone());
assertFalse(actual2.isDone());
assertTrue(actual3.isDone());
try {
actual3.get();
fail();
} catch (Exception e) {
assertTrue(e instanceof ExecutionException);
Throwable actualCause = e.getCause();
assertTrue(actualCause instanceof BufferLimitException);
assertEquals("Rate Limit Buffer is full", actualCause.getMessage());
}
}
@Test
public void releasedPermitTriggerTasksFromQueue() throws InterruptedException {
BufferedRateLimiter limiter = new BufferedRateLimiter(10, 2, 100);
ListenableFuture<Void> actual1 = limiter.acquireAsync();
ListenableFuture<Void> actual2 = limiter.acquireAsync();
ListenableFuture<Void> actual3 = limiter.acquireAsync();
ListenableFuture<Void> actual4 = limiter.acquireAsync();
assertTrue(actual1.isDone());
assertTrue(actual2.isDone());
assertFalse(actual3.isDone());
assertFalse(actual4.isDone());
limiter.release();
TimeUnit.MILLISECONDS.sleep(100L);
assertTrue(actual3.isDone());
assertFalse(actual4.isDone());
limiter.release();
TimeUnit.MILLISECONDS.sleep(100L);
assertTrue(actual4.isDone());
}
@Test
public void permitsReleasedInConcurrentMode() throws InterruptedException {
BufferedRateLimiter limiter = new BufferedRateLimiter(10, 2, 100);
AtomicInteger actualReleased = new AtomicInteger();
AtomicInteger actualRejected = new AtomicInteger();
ListeningExecutorService pool = MoreExecutors.listeningDecorator(Executors.newFixedThreadPool(5));
for (int i = 0; i < 100; i++) {
ListenableFuture<ListenableFuture<Void>> submit = pool.submit(limiter::acquireAsync);
Futures.addCallback(submit, new FutureCallback<ListenableFuture<Void>>() {
@Override
public void onSuccess(@Nullable ListenableFuture<Void> result) {
Futures.addCallback(result, new FutureCallback<Void>() {
@Override
public void onSuccess(@Nullable Void result) {
try {
TimeUnit.MILLISECONDS.sleep(100);
} catch (InterruptedException e) {
e.printStackTrace();
}
limiter.release();
actualReleased.incrementAndGet();
}
@Override
public void onFailure(Throwable t) {
actualRejected.incrementAndGet();
}
});
}
@Override
public void onFailure(Throwable t) {
}
});
}
TimeUnit.SECONDS.sleep(2);
assertTrue("Unexpected released count " + actualReleased.get(),
actualReleased.get() > 10 && actualReleased.get() < 20);
assertTrue("Unexpected rejected count " + actualRejected.get(),
actualRejected.get() > 80 && actualRejected.get() < 90);
}
}