Skip to content

Commit

Permalink
Add more unit tests
Browse files Browse the repository at this point in the history
  • Loading branch information
brenoepics committed Jan 16, 2024
1 parent 68ab4cc commit 1dd5c6c
Show file tree
Hide file tree
Showing 7 changed files with 625 additions and 260 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -14,11 +14,6 @@
*/
public class RateLimitBucket<T> {

// The key is the subscription key, as global rate-limits are shared across the same account.
private static final Map<String, Long> globalRateLimitResetTimestamp = new ConcurrentHashMap<>();

private final AzureApiImpl<T> api;

private final ConcurrentLinkedQueue<RestRequest<T>> requestQueue = new ConcurrentLinkedQueue<>();

private final RestEndpoint endpoint;
Expand All @@ -30,25 +25,14 @@ public class RateLimitBucket<T> {
/**
* Creates a RateLimitBucket for the given endpoint / parameter combination.
*
* @param api The api/shard to use.
* @param endpoint The REST endpoint the rate-limit is tracked for.
* @param majorUrlParameter The url parameter this bucket is specific for. Maybe null.
*/
public RateLimitBucket(AzureApi api, RestEndpoint endpoint, String majorUrlParameter) {
this.api = (AzureApiImpl<T>) api;
public RateLimitBucket(RestEndpoint endpoint, String majorUrlParameter) {
this.endpoint = endpoint;
this.majorUrlParameter = majorUrlParameter;
}

/**
* Sets a global rate-limit.
*
* @param api A azure api instance.
* @param resetTimestamp The reset timestamp of the global rate-limit.
*/
public static void setGlobalRateLimitResetTimestamp(AzureApi api, long resetTimestamp) {
globalRateLimitResetTimestamp.put(api.getSubscriptionKey(), resetTimestamp);
}

/**
* Adds the given request to the bucket's queue.
Expand Down Expand Up @@ -97,8 +81,7 @@ public void setRateLimitResetTimestamp(long rateLimitResetTimestamp) {
* @return The time in seconds how long you have to wait till there's space in the bucket again.
*/
public int getTimeTillSpaceGetsAvailable() {
long globalRLResetTimestamp =
RateLimitBucket.globalRateLimitResetTimestamp.getOrDefault(api.getSubscriptionKey(), 0L);
long globalRLResetTimestamp = 0L;
long timestamp = System.currentTimeMillis();
if (rateLimitRemaining > 0 && (globalRLResetTimestamp - timestamp) <= 0) {
return 0;
Expand Down
Loading

0 comments on commit 1dd5c6c

Please sign in to comment.