diff --git a/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java b/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java index 785b5ef4a..332a86869 100644 --- a/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java +++ b/src/main/java/org/whispersystems/textsecuregcm/limits/LeakyBucket.java @@ -42,7 +42,8 @@ public class LeakyBucket { } public boolean add(int amount) { - this.spaceRemaining = getUpdatedSpaceRemaining(); + this.spaceRemaining = getUpdatedSpaceRemaining(); + this.lastUpdateTimeMillis = System.currentTimeMillis(); if (this.spaceRemaining >= amount) { this.spaceRemaining -= amount; diff --git a/src/test/java/org/whispersystems/textsecuregcm/tests/limits/LeakyBucketTest.java b/src/test/java/org/whispersystems/textsecuregcm/tests/limits/LeakyBucketTest.java new file mode 100644 index 000000000..213baec79 --- /dev/null +++ b/src/test/java/org/whispersystems/textsecuregcm/tests/limits/LeakyBucketTest.java @@ -0,0 +1,52 @@ +package org.whispersystems.textsecuregcm.tests.limits; + +import com.fasterxml.jackson.databind.ObjectMapper; +import org.junit.Test; +import org.whispersystems.textsecuregcm.limits.LeakyBucket; + +import java.io.IOException; +import java.util.concurrent.TimeUnit; + +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +public class LeakyBucketTest { + + @Test + public void testFull() { + LeakyBucket leakyBucket = new LeakyBucket(2, 1.0 / 2.0); + + assertTrue(leakyBucket.add(1)); + assertTrue(leakyBucket.add(1)); + assertFalse(leakyBucket.add(1)); + + leakyBucket = new LeakyBucket(2, 1.0 / 2.0); + + assertTrue(leakyBucket.add(2)); + assertFalse(leakyBucket.add(1)); + assertFalse(leakyBucket.add(2)); + } + + @Test + public void testLapseRate() throws IOException { + ObjectMapper mapper = new ObjectMapper(); + String serialized = "{\"bucketSize\":2,\"leakRatePerMillis\":8.333333333333334E-6,\"spaceRemaining\":0,\"lastUpdateTimeMillis\":" + (System.currentTimeMillis() - TimeUnit.MINUTES.toMillis(2)) + "}"; + + LeakyBucket leakyBucket = LeakyBucket.fromSerialized(mapper, serialized); + assertTrue(leakyBucket.add(1)); + + String serializedAgain = leakyBucket.serialize(mapper); + LeakyBucket leakyBucketAgain = LeakyBucket.fromSerialized(mapper, serializedAgain); + + assertFalse(leakyBucketAgain.add(1)); + } + + @Test + public void testLapseShort() throws Exception { + ObjectMapper mapper = new ObjectMapper(); + String serialized = "{\"bucketSize\":2,\"leakRatePerMillis\":8.333333333333334E-6,\"spaceRemaining\":0,\"lastUpdateTimeMillis\":" + (System.currentTimeMillis() - TimeUnit.MINUTES.toMillis(1)) + "}"; + + LeakyBucket leakyBucket = LeakyBucket.fromSerialized(mapper, serialized); + assertFalse(leakyBucket.add(1)); + } +}