src/club/wpia/gigi/util/RateLimit.java

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