package filodb.core

import java.util.concurrent.atomic.AtomicInteger
import scala.concurrent.{ExecutionContext, Future}

import filodb.core.messages.{Response, TooManyRequests}

trait FutureUtils {
  val limiter: ConcurrentLimiter

  /**
   * Only executes f if the limiter value is below the limit, otherwise
   * fails fast and returns TooManyRequests
   */
  def rateLimit(f: => Future[Response])(implicit ec: ExecutionContext): Future[Response] = {
    if (limiter.obtain()) {
      val future = f
      future.onComplete {
        case x: Any => limiter.release()   // Must release when future is finished
      }
      future
    } else {
      Future.successful(TooManyRequests)
    }
  }
}

/**
 * A simple concurrent resource limiter using j.u.c.AtomicInteger
 * @param limit the maximum number of outstanding futures at one time
 */
private[core] class ConcurrentLimiter(limit: Int) {
  val counter = new AtomicInteger

  def obtain(): Boolean = {
    val newValue = counter.incrementAndGet()
    if (newValue > limit) {
      counter.decrementAndGet()
      false
    } else {
      true
    }
  }

  def release(): Unit = { counter.decrementAndGet() }

  def get: Int = counter.get()
}
package filodb.core

import java.util.concurrent.atomic.AtomicInteger
import scala.concurrent.Future

import filodb.core.messages.{Response, TooManyRequests}

trait FutureUtils {
  val limiter: ConcurrentLimiter

  /**
   * Only executes f if the limiter value is below the limit, otherwise
   * fails fast and returns TooManyRequests
   */
  def rateLimit(f: => Future[Response]): Future[Response] = {
    if (limiter.obtain()) {
      f
    } else {
      Future.successful(TooManyRequests)
    }
  }
}

/**
 * A simple concurrent resource limiter using j.u.c.AtomicInteger
 * @param limit the maximum number of outstanding futures at one time
 */
private[core] class ConcurrentLimiter(limit: Int) {
  val counter = new AtomicInteger

  def obtain(): Boolean = {
    val newValue = counter.incrementAndGet()
    if (newValue > limit) {
      counter.decrementAndGet()
      false
    } else {
      true
    }
  }

  def release(): Unit = { counter.decrementAndGet() }

  def get: Int = counter.get()
}
--- 
+++ 
@@ -1,7 +1,7 @@
 package filodb.core
 
 import java.util.concurrent.atomic.AtomicInteger
-import scala.concurrent.Future
+import scala.concurrent.{ExecutionContext, Future}
 
 import filodb.core.messages.{Response, TooManyRequests}
 
@@ -12,9 +12,13 @@
    * Only executes f if the limiter value is below the limit, otherwise
    * fails fast and returns TooManyRequests
    */
-  def rateLimit(f: => Future[Response]): Future[Response] = {
+  def rateLimit(f: => Future[Response])(implicit ec: ExecutionContext): Future[Response] = {
     if (limiter.obtain()) {
-      f
+      val future = f
+      future.onComplete {
+        case x: Any => limiter.release()   // Must release when future is finished
+      }
+      future
     } else {
       Future.successful(TooManyRequests)
     }
Fix rateLimiter: release when future is finished
