Andrew Shvayka
9 years ago
21 changed files with 745 additions and 91 deletions
@ -0,0 +1,148 @@ |
|||||
|
/** |
||||
|
* 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.datastax.driver.core.Session; |
||||
|
import com.datastax.driver.core.Statement; |
||||
|
import com.google.common.base.Function; |
||||
|
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.Uninterruptibles; |
||||
|
import org.thingsboard.server.dao.util.AsyncRateLimiter; |
||||
|
|
||||
|
import javax.annotation.Nullable; |
||||
|
import java.util.concurrent.*; |
||||
|
|
||||
|
public class RateLimitedResultSetFuture implements ResultSetFuture { |
||||
|
|
||||
|
private final ListenableFuture<ResultSetFuture> originalFuture; |
||||
|
private final ListenableFuture<Void> rateLimitFuture; |
||||
|
|
||||
|
public RateLimitedResultSetFuture(Session session, AsyncRateLimiter rateLimiter, Statement statement) { |
||||
|
this.rateLimitFuture = rateLimiter.acquireAsync(); |
||||
|
this.originalFuture = Futures.transform(rateLimitFuture, |
||||
|
(Function<Void, ResultSetFuture>) i -> executeAsyncWithRelease(rateLimiter, session, statement)); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public ResultSet getUninterruptibly() { |
||||
|
return safeGet().getUninterruptibly(); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public ResultSet getUninterruptibly(long timeout, TimeUnit unit) throws TimeoutException { |
||||
|
long rateLimitStart = System.nanoTime(); |
||||
|
ResultSetFuture resultSetFuture = null; |
||||
|
try { |
||||
|
resultSetFuture = originalFuture.get(timeout, unit); |
||||
|
} catch (InterruptedException | ExecutionException e) { |
||||
|
throw new IllegalStateException(e); |
||||
|
} |
||||
|
long rateLimitDurationNano = System.nanoTime() - rateLimitStart; |
||||
|
long innerTimeoutNano = unit.toNanos(timeout) - rateLimitDurationNano; |
||||
|
if (innerTimeoutNano > 0) { |
||||
|
return resultSetFuture.getUninterruptibly(innerTimeoutNano, TimeUnit.NANOSECONDS); |
||||
|
} |
||||
|
throw new TimeoutException("Timeout waiting for task."); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public boolean cancel(boolean mayInterruptIfRunning) { |
||||
|
if (originalFuture.isDone()) { |
||||
|
return safeGet().cancel(mayInterruptIfRunning); |
||||
|
} else { |
||||
|
return originalFuture.cancel(mayInterruptIfRunning); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public boolean isCancelled() { |
||||
|
if (originalFuture.isDone()) { |
||||
|
return safeGet().isCancelled(); |
||||
|
} |
||||
|
|
||||
|
return originalFuture.isCancelled(); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public boolean isDone() { |
||||
|
return originalFuture.isDone() && safeGet().isDone(); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public ResultSet get() throws InterruptedException, ExecutionException { |
||||
|
return safeGet().get(); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public ResultSet get(long timeout, TimeUnit unit) throws InterruptedException, ExecutionException, TimeoutException { |
||||
|
long rateLimitStart = System.nanoTime(); |
||||
|
ResultSetFuture resultSetFuture = originalFuture.get(timeout, unit); |
||||
|
long rateLimitDurationNano = System.nanoTime() - rateLimitStart; |
||||
|
long innerTimeoutNano = unit.toNanos(timeout) - rateLimitDurationNano; |
||||
|
if (innerTimeoutNano > 0) { |
||||
|
return resultSetFuture.get(innerTimeoutNano, TimeUnit.NANOSECONDS); |
||||
|
} |
||||
|
throw new TimeoutException("Timeout waiting for task."); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public void addListener(Runnable listener, Executor executor) { |
||||
|
originalFuture.addListener(() -> { |
||||
|
try { |
||||
|
ResultSetFuture resultSetFuture = Uninterruptibles.getUninterruptibly(originalFuture); |
||||
|
resultSetFuture.addListener(listener, executor); |
||||
|
} catch (CancellationException e) { |
||||
|
cancel(false); |
||||
|
return; |
||||
|
} catch (ExecutionException e) { |
||||
|
Futures.immediateFailedFuture(e).addListener(listener, executor); |
||||
|
} |
||||
|
}, executor); |
||||
|
} |
||||
|
|
||||
|
private ResultSetFuture safeGet() { |
||||
|
try { |
||||
|
return originalFuture.get(); |
||||
|
} catch (InterruptedException | ExecutionException e) { |
||||
|
throw new IllegalStateException(e); |
||||
|
} |
||||
|
} |
||||
|
|
||||
|
private ResultSetFuture executeAsyncWithRelease(AsyncRateLimiter rateLimiter, Session session, Statement statement) { |
||||
|
try { |
||||
|
ResultSetFuture resultSetFuture = session.executeAsync(statement); |
||||
|
Futures.addCallback(resultSetFuture, new FutureCallback<ResultSet>() { |
||||
|
@Override |
||||
|
public void onSuccess(@Nullable ResultSet result) { |
||||
|
rateLimiter.release(); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public void onFailure(Throwable t) { |
||||
|
rateLimiter.release(); |
||||
|
} |
||||
|
}); |
||||
|
return resultSetFuture; |
||||
|
} catch (RuntimeException re) { |
||||
|
rateLimiter.release(); |
||||
|
throw re; |
||||
|
} |
||||
|
} |
||||
|
} |
||||
@ -0,0 +1,25 @@ |
|||||
|
/** |
||||
|
* 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.ListenableFuture; |
||||
|
|
||||
|
public interface AsyncRateLimiter { |
||||
|
|
||||
|
ListenableFuture<Void> acquireAsync(); |
||||
|
|
||||
|
void release(); |
||||
|
} |
||||
@ -0,0 +1,160 @@ |
|||||
|
/** |
||||
|
* 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 java.util.concurrent.*; |
||||
|
import java.util.concurrent.atomic.AtomicInteger; |
||||
|
|
||||
|
@Component |
||||
|
@Slf4j |
||||
|
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(); |
||||
|
|
||||
|
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() { |
||||
|
if (queue.isEmpty()) { |
||||
|
if (permits.incrementAndGet() <= permitsLimit) { |
||||
|
if (permits.get() > maxGrantedPermissions.get()) { |
||||
|
maxGrantedPermissions.set(permits.get()); |
||||
|
} |
||||
|
return Futures.immediateFuture(null); |
||||
|
} |
||||
|
permits.decrementAndGet(); |
||||
|
} |
||||
|
|
||||
|
return putInQueue(); |
||||
|
} |
||||
|
|
||||
|
@Override |
||||
|
public void release() { |
||||
|
permits.decrementAndGet(); |
||||
|
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) { |
||||
|
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 IllegalStateException("Rate Limit Buffer is full. Reject")); |
||||
|
} |
||||
|
return lockedFuture.future; |
||||
|
} catch (InterruptedException e) { |
||||
|
return Futures.immediateFailedFuture(new IllegalStateException("Rate Limit Task interrupted. Reject")); |
||||
|
} |
||||
|
} |
||||
|
return Futures.immediateFailedFuture(new IllegalStateException("Rate Limit Buffer is full. Reject")); |
||||
|
} |
||||
|
|
||||
|
@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 is [{}] max concurrent [{}] expired [{}]", maxQueueSize.getAndSet(0), |
||||
|
maxGrantedPermissions.getAndSet(0), expiredCount); |
||||
|
} |
||||
|
|
||||
|
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; |
||||
|
} |
||||
|
|
||||
|
} |
||||
|
|
||||
|
|
||||
|
} |
||||
@ -0,0 +1,156 @@ |
|||||
|
/** |
||||
|
* 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.*; |
||||
|
import com.datastax.driver.core.exceptions.UnsupportedFeatureException; |
||||
|
import com.google.common.util.concurrent.Futures; |
||||
|
import com.google.common.util.concurrent.ListenableFuture; |
||||
|
import org.junit.Test; |
||||
|
import org.junit.runner.RunWith; |
||||
|
import org.mockito.Mock; |
||||
|
import org.mockito.Mockito; |
||||
|
import org.mockito.runners.MockitoJUnitRunner; |
||||
|
import org.mockito.stubbing.Answer; |
||||
|
import org.thingsboard.server.dao.util.AsyncRateLimiter; |
||||
|
|
||||
|
import java.util.concurrent.ExecutionException; |
||||
|
import java.util.concurrent.TimeoutException; |
||||
|
|
||||
|
import static org.junit.Assert.*; |
||||
|
import static org.mockito.Mockito.*; |
||||
|
|
||||
|
@RunWith(MockitoJUnitRunner.class) |
||||
|
public class RateLimitedResultSetFutureTest { |
||||
|
|
||||
|
private RateLimitedResultSetFuture resultSetFuture; |
||||
|
|
||||
|
@Mock |
||||
|
private AsyncRateLimiter rateLimiter; |
||||
|
@Mock |
||||
|
private Session session; |
||||
|
@Mock |
||||
|
private Statement statement; |
||||
|
@Mock |
||||
|
private ResultSetFuture realFuture; |
||||
|
@Mock |
||||
|
private ResultSet rows; |
||||
|
@Mock |
||||
|
private Row row; |
||||
|
|
||||
|
@Test |
||||
|
public void doNotReleasePermissionIfRateLimitFutureFailed() throws InterruptedException { |
||||
|
when(rateLimiter.acquireAsync()).thenReturn(Futures.immediateFailedFuture(new IllegalArgumentException())); |
||||
|
resultSetFuture = new RateLimitedResultSetFuture(session, rateLimiter, statement); |
||||
|
Thread.sleep(1000L); |
||||
|
verify(rateLimiter).acquireAsync(); |
||||
|
try { |
||||
|
assertTrue(resultSetFuture.isDone()); |
||||
|
fail(); |
||||
|
} catch (Exception e) { |
||||
|
assertTrue(e instanceof IllegalStateException); |
||||
|
Throwable actualCause = e.getCause(); |
||||
|
assertTrue(actualCause instanceof ExecutionException); |
||||
|
} |
||||
|
verifyNoMoreInteractions(session, rateLimiter, statement); |
||||
|
|
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
public void getUninterruptiblyDelegateToCassandra() throws InterruptedException, ExecutionException { |
||||
|
when(rateLimiter.acquireAsync()).thenReturn(Futures.immediateFuture(null)); |
||||
|
when(session.executeAsync(statement)).thenReturn(realFuture); |
||||
|
Mockito.doAnswer((Answer<Void>) invocation -> { |
||||
|
Object[] args = invocation.getArguments(); |
||||
|
Runnable task = (Runnable) args[0]; |
||||
|
task.run(); |
||||
|
return null; |
||||
|
}).when(realFuture).addListener(Mockito.any(), Mockito.any()); |
||||
|
|
||||
|
when(realFuture.getUninterruptibly()).thenReturn(rows); |
||||
|
|
||||
|
resultSetFuture = new RateLimitedResultSetFuture(session, rateLimiter, statement); |
||||
|
ResultSet actual = resultSetFuture.getUninterruptibly(); |
||||
|
assertSame(rows, actual); |
||||
|
verify(rateLimiter, times(1)).acquireAsync(); |
||||
|
verify(rateLimiter, times(1)).release(); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
public void addListenerAllowsFutureTransformation() throws InterruptedException, ExecutionException { |
||||
|
when(rateLimiter.acquireAsync()).thenReturn(Futures.immediateFuture(null)); |
||||
|
when(session.executeAsync(statement)).thenReturn(realFuture); |
||||
|
Mockito.doAnswer((Answer<Void>) invocation -> { |
||||
|
Object[] args = invocation.getArguments(); |
||||
|
Runnable task = (Runnable) args[0]; |
||||
|
task.run(); |
||||
|
return null; |
||||
|
}).when(realFuture).addListener(Mockito.any(), Mockito.any()); |
||||
|
|
||||
|
when(realFuture.get()).thenReturn(rows); |
||||
|
when(rows.one()).thenReturn(row); |
||||
|
|
||||
|
resultSetFuture = new RateLimitedResultSetFuture(session, rateLimiter, statement); |
||||
|
|
||||
|
ListenableFuture<Row> transform = Futures.transform(resultSetFuture, ResultSet::one); |
||||
|
Row actualRow = transform.get(); |
||||
|
|
||||
|
assertSame(row, actualRow); |
||||
|
verify(rateLimiter, times(1)).acquireAsync(); |
||||
|
verify(rateLimiter, times(1)).release(); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
public void immidiateCassandraExceptionReturnsPermit() throws InterruptedException, ExecutionException { |
||||
|
when(rateLimiter.acquireAsync()).thenReturn(Futures.immediateFuture(null)); |
||||
|
when(session.executeAsync(statement)).thenThrow(new UnsupportedFeatureException(ProtocolVersion.V3, "hjg")); |
||||
|
resultSetFuture = new RateLimitedResultSetFuture(session, rateLimiter, statement); |
||||
|
ListenableFuture<Row> transform = Futures.transform(resultSetFuture, ResultSet::one); |
||||
|
try { |
||||
|
transform.get(); |
||||
|
fail(); |
||||
|
} catch (Exception e) { |
||||
|
assertTrue(e instanceof ExecutionException); |
||||
|
} |
||||
|
verify(rateLimiter, times(1)).acquireAsync(); |
||||
|
verify(rateLimiter, times(1)).release(); |
||||
|
} |
||||
|
|
||||
|
@Test |
||||
|
public void queryTimeoutReturnsPermit() throws InterruptedException, ExecutionException { |
||||
|
when(rateLimiter.acquireAsync()).thenReturn(Futures.immediateFuture(null)); |
||||
|
when(session.executeAsync(statement)).thenReturn(realFuture); |
||||
|
Mockito.doAnswer((Answer<Void>) invocation -> { |
||||
|
Object[] args = invocation.getArguments(); |
||||
|
Runnable task = (Runnable) args[0]; |
||||
|
task.run(); |
||||
|
return null; |
||||
|
}).when(realFuture).addListener(Mockito.any(), Mockito.any()); |
||||
|
|
||||
|
when(realFuture.get()).thenThrow(new ExecutionException("Fail", new TimeoutException("timeout"))); |
||||
|
resultSetFuture = new RateLimitedResultSetFuture(session, rateLimiter, statement); |
||||
|
ListenableFuture<Row> transform = Futures.transform(resultSetFuture, ResultSet::one); |
||||
|
try { |
||||
|
transform.get(); |
||||
|
fail(); |
||||
|
} catch (Exception e) { |
||||
|
assertTrue(e instanceof ExecutionException); |
||||
|
} |
||||
|
verify(rateLimiter, times(1)).acquireAsync(); |
||||
|
verify(rateLimiter, times(1)).release(); |
||||
|
} |
||||
|
|
||||
|
} |
||||
@ -0,0 +1,134 @@ |
|||||
|
/** |
||||
|
* 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.*; |
||||
|
import org.junit.Test; |
||||
|
|
||||
|
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.*; |
||||
|
|
||||
|
|
||||
|
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 IllegalStateException); |
||||
|
assertEquals("Rate Limit Buffer is full. Reject", 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); |
||||
|
|
||||
|
} |
||||
|
|
||||
|
|
||||
|
} |
||||
Loading…
Reference in new issue