src/org/cacert/gigi/util/RateLimit.java

   1 package org.cacert.gigi.util;
   2
   3 import java.util.HashMap;
   4 import java.util.TreeSet;
   5
   6 import org.cacert.gigi.GigiApiException;
   7
   8 public class RateLimit {
   9
  10     public static final class RateLimitException extends GigiApiException {
  11
  12         public RateLimitException() {
  13             super("Rate limit exceeded.");
  14         }
  15     }
  16
  17     private class Entry implements Comparable<Entry> {
  18
  19         long firstAccess;
  20
  21         int count = 1;
  22
  23         String feature;
  24
  25         public Entry(long firstAccess, String feature) {
  26             this.firstAccess = firstAccess;
  27             this.feature = feature;
  28         }
  29
  30         public void access() {
  31             count++;
  32         }
  33
  34         @Override
  35         public int compareTo(Entry o) {
  36             return feature.compareTo(o.feature);
  37         }
  38
  39         public boolean isExpired() {
  40             return firstAccess + time < System.currentTimeMillis();
  41         }
  42
  43     }
  44
  45     private final int maxcount;
  46
  47     private final long time;
  48
  49     TreeSet<Entry> set = new TreeSet<Entry>();
  50
  51     HashMap<String, Entry> feat = new HashMap<>();
  52
  53     public RateLimit(int maxcount, long time) {
  54         this.maxcount = maxcount;
  55         this.time = time;
  56     }
  57
  58     public synchronized boolean isLimitExceeded(String feature) {
  59         clean();
  60         Entry e = feat.get(feature);
  61         if (e == null) {
  62             e = new Entry(System.currentTimeMillis(), feature);
  63             set.add(e);
  64             feat.put(feature, e);
  65         } else {
  66             e.access();
  67         }
  68         return e.count > maxcount;
  69     }
  70
  71     private void clean() {
  72         while (set.size() > 0) {
  73             Entry e = set.last();
  74             if (e.isExpired()) {
  75                 set.remove(e);
  76                 feat.remove(e.feature);
  77             } else {
  78                 return;
  79             }
  80         }
  81     }
  82
  83     public synchronized void bypass() {
  84         set.clear();
  85         feat.clear();
  86     }
  87 }