Make TokenBucketGroup a trait to allow easier customization about play-guard HOT 2 CLOSED

  object TokenBuckets {
    val NanosPerSecond: Long = 1000L * 1000 * 1000
  }

  class TokenBuckets(size: Int, rate: Double, minNanosPerUpdate: Long = TokenBuckets.NanosPerSecond)
                    (implicit actorSystem: ActorSystem) {

    require(size > 0)
    require(rate >= 0.0001)
    require(rate < 1e6)

    def now(): Long = System.nanoTime()

    //private val floor = 0
    private val floor = -size // Penalize abusers a little
    private val buckets = new ConcurrentHashMap[String, AtomicInteger]()
    private var lastUpdateAt: Long = now()
    private val tokensPerNano: Double = rate / TokenBuckets.NanosPerSecond
    private val nanosPerToken: Double = 1 / tokensPerNano
    private val nanosPerUpdate: Long = Math.max(nanosPerToken.ceil.toLong, minNanosPerUpdate)

    def consume(ip: String): Int = {
      var int = buckets.get(ip)
      if (int == null) {
        int = new AtomicInteger(size)
        buckets.put(ip, int) // May drop a token if there's a concurrent request, that's ok
      }
      if (int.get() > floor) int.decrementAndGet() else floor
    }

    import scala.concurrent.duration._
    import scala.concurrent.ExecutionContext.Implicits.global

    actorSystem.scheduler.schedule(nanosPerUpdate.nanos, nanosPerUpdate.nanos) {
      val deltaNanos: Long = now() - lastUpdateAt
      val tokensToAdd = (deltaNanos * tokensPerNano).round.toInt
      if (tokensToAdd > 0) { // Should always be true
        buckets.forEach((ip: String, int: AtomicInteger) => {
          val tokensInBucket = int.addAndGet(tokensToAdd)
          if (tokensInBucket >= size) buckets.remove(ip) // May drop a token if there's a concurrent request, that's ok
        })
        lastUpdateAt += (tokensToAdd * nanosPerToken).ceil.toLong
      }
    }
  }