Merge pull request skinny-framework#257 from seratch/memory-leak

Fix memory leak at SkinnyFilterActivation

framework/src/main/scala/org/scalatra/SkinnyScalatraBase.scala

@@ -1,10 +1,12 @@
 package org.scalatra
 
 import org.scalatra.ScalatraBase._
+import java.util.concurrent.atomic.AtomicInteger
 import javax.servlet.Filter
 import javax.servlet.http.{ HttpServletResponse, HttpServletRequest, HttpServlet }
 import scala.util.Failure
 import skinny.SkinnyEnv
+import skinny.logging.Logging
 
 /**
  * Partially patched ScalatraBase for Skinny Framework.
@@ -17,7 +19,7 @@ import skinny.SkinnyEnv
  * So We've patched ScalatraBase to ignore "org.scalatra.ScalatraFilter.afterFilters.Run" only for Filters.
  * Hope Scalatra to support ignoring "org.scalatra.ScalatraFilter.afterFilters.Run" option to 3rd party.
  */
-trait SkinnyScalatraBase extends ScalatraBase {
+trait SkinnyScalatraBase extends ScalatraBase with Logging {
 
   override protected def executeRoutes() {
     var result: Any = null
@@ -120,4 +122,25 @@ trait SkinnyScalatraBase extends ScalatraBase {
     }
   }
 
+  /**
+   * Count execution of error filter registration.
+   */
+  private[this] lazy val errorMethodCallCountAtSkinnyScalatraBase: AtomicInteger = new AtomicInteger(0)
+
+  /**
+   * Detects error filter leak issue as an error.
+   */
+  protected def detectTooManyErrorFilterRegistrationnAsAnErrorAtSkinnyScalatraBase: Boolean = false
+
+  // https://github.com/scalatra/scalatra/blob/v2.3.1/core/src/main/scala/org/scalatra/ScalatraBase.scala#L333-L335
+  override def error(handler: ErrorHandler) {
+    val count = errorMethodCallCountAtSkinnyScalatraBase.incrementAndGet()
+    if (count > 500) {
+      val message = s"skinny's error filter registration for this controller has been evaluated $count times, this behavior will cause memory leak."
+      if (detectTooManyErrorFilterRegistrationnAsAnErrorAtSkinnyScalatraBase) throw new RuntimeException(message)
+      else logger.warn(message)
+    }
+    super.error(handler)
+  }
+
 }

framework/src/main/scala/skinny/filter/SkinnyFilterActivation.scala

@@ -6,7 +6,7 @@ import org.scalatra._
 /**
  * Activates skinny filters.
  */
-trait SkinnyFilterActivation { self: SkinnyControllerBase =>
+trait SkinnyFilterActivation extends SkinnyScalatraBase { self: SkinnyControllerBase =>
 
   sealed trait RenderingRequired
   case object WithRendering extends RenderingRequired
@@ -27,48 +27,55 @@ trait SkinnyFilterActivation { self: SkinnyControllerBase =>
     skinnyErrorFilters.update(WithoutRendering, mergedHandlers)
   }
 
-  /**
-   * Start applying all the filters
-   */
-  beforeAction() {
-    // combine error filters
-    val filtersTraverse: PartialFunction[Throwable, Any] = {
-      case (t: Throwable) =>
+  // mark as already initialized
+  private[this] var isFiltersTraverseAtSkinnyFilterActivationAlreadyLoaded: Boolean = false
 
-        skinnyErrorFilters.get(WithoutRendering).foreach { handlers =>
-          handlers.foreach { handler =>
-            // just apply this filter and return the existing body.
-            try handler.apply(t)
-            catch { case e: Exception => logger.error(s"Failed to apply SkinnyFilter (error: ${e.getMessage})", e) }
-          }
+  // combined error filters
+  private[this] lazy val filtersTraverseAtSkinnyFilterActivation: PartialFunction[Throwable, Any] = {
+    case (t: Throwable) =>
+
+      skinnyErrorFilters.get(WithoutRendering).foreach { handlers =>
+        handlers.foreach { handler =>
+          // just apply this filter and return the existing body.
+          try handler.apply(t)
+          catch { case e: Exception => logger.error(s"Failed to apply SkinnyFilter (error: ${e.getMessage})", e) }
         }
+      }
 
-        var rendered = false
+      var rendered = false
 
-        // rendering body
-        val body: Any = skinnyErrorFilters.get(WithRendering).map { handlers =>
-          handlers.foldLeft(null.asInstanceOf[Any]) {
-            case (body, handler) =>
-              // just apply this filter and return the existing body.
-              if (rendered) {
-                logger.error("This response is marked as rendered because other RenderingFilter already did the stuff.")
-                body
-              } else {
-                rendered = true
-                try handler.apply(t)
-                catch {
-                  case e: Exception =>
-                    logger.error(s"Failed to apply SkinnyFilter (error: ${e.getMessage})", e)
-                    body
-                }
+      // rendering body
+      val body: Any = skinnyErrorFilters.get(WithRendering).map { handlers =>
+        handlers.foldLeft(null.asInstanceOf[Any]) {
+          case (body, handler) =>
+            // just apply this filter and return the existing body.
+            if (rendered) {
+              logger.error("This response is marked as rendered because other RenderingFilter already did the stuff.")
+              body
+            } else {
+              rendered = true
+              try handler.apply(t)
+              catch {
+                case e: Exception =>
+                  logger.error(s"Failed to apply SkinnyFilter (error: ${e.getMessage})", e)
+                  body
               }
-          }
-        }.getOrElse(null)
+            }
+        }
+      }.getOrElse(null)
+
+      if (rendered) body else throw t
+  }
 
-        if (rendered) body else throw t
+  /**
+   * Start applying all the filters
+   */
+  beforeAction() {
+    if (!isFiltersTraverseAtSkinnyFilterActivationAlreadyLoaded) {
+      // register combined error filters
+      error(filtersTraverseAtSkinnyFilterActivation)
+      isFiltersTraverseAtSkinnyFilterActivationAlreadyLoaded = true
     }
-    // register combined error filters
-    error(filtersTraverse)
   }
 
 }

framework/src/test/scala/skinny/filter/SkinnyFilterActivationSpec.scala

@@ -0,0 +1,24 @@
+package skinny.filter
+
+import org.scalatra.test.scalatest.ScalatraFlatSpec
+import skinny.SkinnyApiController
+import skinny.routing.Routes
+
+class SkinnyFilterActivationSpec extends ScalatraFlatSpec with skinny.Logging {
+
+  object Controller extends SkinnyApiController with SkinnyFilterActivation with Routes {
+    override def detectTooManyErrorFilterRegistrationnAsAnErrorAtSkinnyScalatraBase = true
+    def index = "ok"
+    get("/")(index).as('index)
+  }
+  addFilter(Controller, "/*")
+
+  it should "not waste memory" in {
+    (1 to 500).foreach { _ =>
+      get("/") { status should equal(200) }
+    }
+    // if possible leak detected, this status will be 500
+    get("/") { status should equal(200) }
+  }
+
+}