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 public class RateLimit {
   7
   8     private class Entry implements Comparable<Entry> {
   9
  10         long firstAccess;
  11
  12         int count = 1;
  13
  14         String feature;
  15
  16         public Entry(long firstAccess, String feature) {
  17             this.firstAccess = firstAccess;
  18             this.feature = feature;
  19         }
  20
  21         public void access() {
  22             count++;
  23         }
  24
  25         @Override
  26         public int compareTo(Entry o) {
  27             return feature.compareTo(o.feature);
  28         }
  29
  30         public boolean isExpired() {
  31             return firstAccess + time < System.currentTimeMillis();
  32         }
  33
  34     }
  35
  36     private final int maxcount;
  37
  38     private final long time;
  39
  40     TreeSet<Entry> set = new TreeSet<Entry>();
  41
  42     HashMap<String, Entry> feat = new HashMap<>();
  43
  44     public RateLimit(int maxcount, long time) {
  45         this.maxcount = maxcount;
  46         this.time = time;
  47     }
  48
  49     public synchronized boolean isLimitExceeded(String feature) {
  50         clean();
  51         Entry e = feat.get(feature);
  52         if (e == null) {
  53             e = new Entry(System.currentTimeMillis(), feature);
  54             set.add(e);
  55             feat.put(feature, e);
  56         } else {
  57             e.access();
  58         }
  59         return e.count > maxcount;
  60     }
  61
  62     private void clean() {
  63         while (set.size() > 0) {
  64             Entry e = set.last();
  65             if (e.isExpired()) {
  66                 set.remove(e);
  67                 feat.remove(e.feature);
  68             } else {
  69                 return;
  70             }
  71         }
  72     }
  73
  74     public synchronized void bypass() {
  75         set.clear();
  76         feat.clear();
  77     }
  78 }