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