parent
75aec0a8d4
commit
0928e4c035
|
@ -17,14 +17,13 @@
|
|||
package org.whispersystems.textsecuregcm;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import org.whispersystems.textsecuregcm.configuration.RedisConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.FederationConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.GraphiteConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.MemcacheConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.NexmoConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.PushConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.RateLimitsConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.RedPhoneConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.RedisConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.S3Configuration;
|
||||
import org.whispersystems.textsecuregcm.configuration.TwilioConfiguration;
|
||||
import org.whispersystems.textsecuregcm.configuration.WebsocketConfiguration;
|
||||
|
|
|
@ -17,26 +17,28 @@
|
|||
package org.whispersystems.textsecuregcm.limits;
|
||||
|
||||
import com.fasterxml.jackson.annotation.JsonProperty;
|
||||
import com.fasterxml.jackson.core.JsonProcessingException;
|
||||
import com.fasterxml.jackson.databind.ObjectMapper;
|
||||
|
||||
import java.io.IOException;
|
||||
|
||||
public class LeakyBucket {
|
||||
|
||||
@JsonProperty
|
||||
private final int bucketSize;
|
||||
|
||||
@JsonProperty
|
||||
private final double leakRatePerMillis;
|
||||
|
||||
@JsonProperty
|
||||
private int spaceRemaining;
|
||||
|
||||
@JsonProperty
|
||||
private long lastUpdateTimeMillis;
|
||||
|
||||
public LeakyBucket(int bucketSize, double leakRatePerMillis) {
|
||||
this(bucketSize, leakRatePerMillis, bucketSize, System.currentTimeMillis());
|
||||
}
|
||||
|
||||
private LeakyBucket(int bucketSize, double leakRatePerMillis, int spaceRemaining, long lastUpdateTimeMillis) {
|
||||
this.bucketSize = bucketSize;
|
||||
this.leakRatePerMillis = leakRatePerMillis;
|
||||
this.spaceRemaining = bucketSize;
|
||||
this.lastUpdateTimeMillis = System.currentTimeMillis();
|
||||
this.spaceRemaining = spaceRemaining;
|
||||
this.lastUpdateTimeMillis = lastUpdateTimeMillis;
|
||||
}
|
||||
|
||||
public boolean add(int amount) {
|
||||
|
@ -56,4 +58,40 @@ public class LeakyBucket {
|
|||
return Math.min(this.bucketSize,
|
||||
(int)Math.floor(this.spaceRemaining + (elapsedTime * this.leakRatePerMillis)));
|
||||
}
|
||||
|
||||
public String serialize(ObjectMapper mapper) throws JsonProcessingException {
|
||||
return mapper.writeValueAsString(new LeakyBucketEntity(bucketSize, leakRatePerMillis, spaceRemaining, lastUpdateTimeMillis));
|
||||
}
|
||||
|
||||
public static LeakyBucket fromSerialized(ObjectMapper mapper, String serialized) throws IOException {
|
||||
LeakyBucketEntity entity = mapper.readValue(serialized, LeakyBucketEntity.class);
|
||||
|
||||
return new LeakyBucket(entity.bucketSize, entity.leakRatePerMillis,
|
||||
entity.spaceRemaining, entity.lastUpdateTimeMillis);
|
||||
}
|
||||
|
||||
private static class LeakyBucketEntity {
|
||||
@JsonProperty
|
||||
private int bucketSize;
|
||||
|
||||
@JsonProperty
|
||||
private double leakRatePerMillis;
|
||||
|
||||
@JsonProperty
|
||||
private int spaceRemaining;
|
||||
|
||||
@JsonProperty
|
||||
private long lastUpdateTimeMillis;
|
||||
|
||||
public LeakyBucketEntity() {}
|
||||
|
||||
private LeakyBucketEntity(int bucketSize, double leakRatePerMillis,
|
||||
int spaceRemaining, long lastUpdateTimeMillis)
|
||||
{
|
||||
this.bucketSize = bucketSize;
|
||||
this.leakRatePerMillis = leakRatePerMillis;
|
||||
this.spaceRemaining = spaceRemaining;
|
||||
this.lastUpdateTimeMillis = lastUpdateTimeMillis;
|
||||
}
|
||||
}
|
||||
}
|
|
@ -73,7 +73,7 @@ public class RateLimiter {
|
|||
|
||||
private void setBucket(String key, LeakyBucket bucket) {
|
||||
try (Jedis jedis = cacheClient.getResource()) {
|
||||
String serialized = mapper.writeValueAsString(bucket);
|
||||
String serialized = bucket.serialize(mapper);
|
||||
jedis.setex(getBucketName(key), (int) Math.ceil((bucketSize / leakRatePerMillis) / 1000), serialized);
|
||||
} catch (JsonProcessingException e) {
|
||||
throw new IllegalArgumentException(e);
|
||||
|
@ -85,7 +85,7 @@ public class RateLimiter {
|
|||
String serialized = jedis.get(getBucketName(key));
|
||||
|
||||
if (serialized != null) {
|
||||
return mapper.readValue(serialized, LeakyBucket.class);
|
||||
return LeakyBucket.fromSerialized(mapper, serialized);
|
||||
}
|
||||
} catch (IOException e) {
|
||||
logger.warn("Deserialization error", e);
|
||||
|
|
Loading…
Reference in New Issue