Bob Balaban's Blog

     
    alt

    Bob Balaban

     

    Token bucket!

    Bob Balaban  June 8 2014 06:41:25 AM
    Greetings, geeks!

    I came across this problem on a project I'm doing: given a RESTful API on a web application server, how do you "throttle" calls so that a given authenticated user can only make a certain number of calls (N) per time interval (W), where both N and W are configurable. The solution I came up with is basically a “Token Bucket” with a sliding time window. What surprised me when I got something working was how easy it really was to implement.

    The context in which this bit of code operates is:
     - You know who the caller is (they're already authenticated)
     - The application can maintain a "user context" that you can retrieve on each HTTP "call" or invocation

    The algorithm is as follows:
    Given a per-user TreeSet (a built-in Java collection class that implements a sorted set) containing long values representing java.util.Date time values (milliseconds). The actual set will contain Long values, not long, but that’s transparent to the algorithm. The set contains a list of times the user has invoked the API, within the last interval W

      - When an invocation arrives at the server, set NOW = long value of “now”, from java.util.Date.

     - For the current user’s list of previous call times, do:
         - Inspect the oldest entry, if older than (NOW – W), remove it
         - Until all old entries are removed, or the list is empty.

     - If the number of entries remaining in the set + 1 is > N, return error
     - Else, add NOW to the set and continue

    The application server context is per-user, and you cache a TreeSet instance for each authenticated client in the context. Here's a link to the file:
    TimeBasedAccessList.java

    As it's only 44 lines, here's all the code...

    import java.util.Date;
    import java.util.TreeSet;

    public class TimeBasedAccessList {

          private volatile TreeSet timeSet = new TreeSet();
          private int nCalls;                        // count of allowed calls
          private long timeWindow;        // in millisec
         
          public TimeBasedAccessList(int N, long M) {
                  nCalls = N;
                  timeWindow = M;
          }

          public synchronized boolean isCallOK() {
                  long now = new Date().getTime();
                  long earliest = now - this.timeWindow;
                 
                  // remove entries older than "earliest"
                  // can't use "for" loop with removals inside
                  while (!this.timeSet.isEmpty()) {
                          Long L = this.timeSet.first();         // "lowest" valued element is earliest timestamp
                          if (L.longValue() <= earliest) {
                                  this.timeSet.remove(L);
                          }
                          else {
                                  // no more entries out of window
                                  break;
                          }
                  } // end for
                 
                  // see how many entries there are
                  if (this.timeSet.size() >= this.nCalls) {
                          return false;
                  }
                 
                  // add new entry
                  this.timeSet.add(new Long(now));
                  return true;
          } // end isCallOK
         
    } // end class
    Not much, is it? There's meant to be an instance of the class for each authenticated user. The constructor captures the size of the time window and the number of allowed calls in that time. All the interesting stuff happens in the isCallOK() method. It's possible on a web server that two calls from the same user could come in at the same time, and be dispatched on separate threads. Therefore the isCallOK() method needs to be synchronized to be thread-safe.

    In my particular project, I'm using a Tomcat app server, and the consumer of this class is a Tomcat filter. Filters are cool, you designate one to the server in your application (.WAR file), and Tomcat invokes your filter code before any servlets or JSPs that you have configured. The filter can examine the request data, modify things, do server-wide logging, whatever. In my case, I'm using this particular filter as a gateway: if the user is making too many calls, the filter just returns an HTTP error code (I picked 503 - Unavailable, which usually means "I'm busy, try again later"). If not, then the filter passes control on to either another filter, or to the designated endpoint for the caller's request..

    A few comments on the details of the code...

    The TreeSet object (built-in class in the Java Collections, um, collection) can only store Objects, not scalars, like "long", so we have to convert between "long" and "Long" (the Object version). The only slightly tricky bit of this was looping over the TreeSet and removing entries outside the specified time window. Originally I coded it as a "for" loop: for (Long L : this.timeSet), etc. The problem was that it threw an exception. The reason is documented, though a little obscure. When you do a "for" loop over a collection object, the JVM actually builds the looping functionality using an Iterator instance that it gets from the collection object. That part's fine, but when you use an Iterator (either explicitly by coding it, or invisibly, as with a for loop), you're supposed to make any changes to the contents of the collection (such as removing items) via the Iterator, not directly to the collection object itself. That's because if you don't do it through the Iterator, the Iterator's knowledge of what's in the collection gets out of sync with the collection's actual contents.

    I could have fixed the loop by explicitly getting the Iterator instance and using it to, um, iterate, as well as to remove old entries. But instead I chose to code it by iteratively calling TreeSet.first(), which returns the smallest valued (and in this case, therefore the earliest) entry, or null if the set is empty. The loop exits when the first entry is found that is inside the specified time window, or there are no entries left.

    So, the TreeSet class is the hero here, because it automatically sorts the entries for me. I looked around for a Queue class where I could add new entries to one end, and remove old entries from the other, because that would eliminate the need for sorting overhead. I didn't find one, though, and I was too lazy to write one, so I went with TreeSet. The efficiency/speed of this code is still very good, testing whether the next server invocation for a given user is ok to proceed (calling the isCallOK() method) only consumes a few milliseconds.

    So, there you have it. Geek ya later!

    Follow me on Twitter @LooseleafLLC
    This article ©Copyright 2014 by Looseleaf Software LLC, all rights reserved. You may link to this page, but may not copy without prior approval.

    Comments

    1Vladimir  05/02/2017 4:50:02 AM  Token bucket!

    There is functional bug in this imlementation. If two events happen in same millisecond, then last overwrites first, and as result last event stored in both "nCalls" and "timeSet", and first stored only in "nCalls", so token consumed by first event will never returned back to bucket.

    2Bob Balaban  05/06/2017 3:35:21 PM  Token bucket!

    Thanks for your comment, Vladimir. I don't think you are correct, becase elements added to the list never overwrite, so there is never a conflict.

    Besides that, the method is synchronized, so there are no multi-threading conflicts. I haven't tested the situation you describe, but I think it will still work correctly. If not, I think the fix would be to store nano-time instead of millisecond-time.