Better management lifetime of treebuilder's cached data.

Author: baffles

codepulse/src/main/scala/com/secdec/codepulse/tracer/AkkaTracingTarget.scala

@@ -33,7 +33,6 @@ import reactive.EventStream
 import com.secdec.codepulse.data.jsp.JspMapper
 import com.secdec.codepulse.data.trace.TraceId
 import com.secdec.codepulse.data.trace.TraceData
-import com.secdec.codepulse.data.trace.TreeBuilder
 import akka.actor.PoisonPill
 
 //sealed trait TracingTargetEvent
@@ -77,7 +76,6 @@ trait TracingTarget {
 
 	def traceData: TraceData
 	def transientData: TransientTraceData
-	def treeBuilder: TreeBuilder
 
 	def requestDeletion(): Unit
 	def notifyFinishedLoading(): Unit
@@ -102,7 +100,7 @@ object AkkaTracingTarget {
 	// implicit Timeout used for the Akka ask (?) method in getState
 	private implicit val timeout = new Timeout(5000)
 
-	private class TracingTargetImpl(val id: TraceId, actor: ActorRef, val traceData: TraceData, val transientData: TransientTraceData, val treeBuilder: TreeBuilder) extends TracingTarget with AskSupport {
+	private class TracingTargetImpl(val id: TraceId, actor: ActorRef, val traceData: TraceData, val transientData: TransientTraceData) extends TracingTarget with AskSupport {
 		def subscribeToStateChanges(sub: EventStream[TracingTargetState] => Unit) = actor ! Subscribe(sub)
 		def requestNewTraceConnection() = actor ! RequestTraceConnect
 		def requestTraceEnd() = actor ! RequestTraceEnd
@@ -114,10 +112,10 @@ object AkkaTracingTarget {
 		def notifyFinishedLoading() = actor ! FinishedLoading
 	}
 
-	def apply(actorSystem: ActorSystem, traceId: TraceId, traceData: TraceData, transientTraceData: TransientTraceData, treeBuilder: TreeBuilder, jspMapper: Option[JspMapper]): TracingTarget = {
+	def apply(actorSystem: ActorSystem, traceId: TraceId, traceData: TraceData, transientTraceData: TransientTraceData, jspMapper: Option[JspMapper]): TracingTarget = {
 		val props = Props { new AkkaTracingTarget(traceId, traceData, transientTraceData, jspMapper) }
 		val actorRef = actorSystem.actorOf(props)
-		new TracingTargetImpl(traceId, actorRef, traceData, transientTraceData, treeBuilder)
+		new TracingTargetImpl(traceId, actorRef, traceData, transientTraceData)
 	}
 
 }

codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceAPIServer.scala

@@ -39,7 +39,7 @@ import net.liftweb.json.JsonDSL._
 import net.liftweb.util.Helpers.AsBoolean
 import net.liftweb.util.Helpers.AsInt
 
-class TraceAPIServer(manager: TraceManager) extends RestHelper with Loggable {
+class TraceAPIServer(manager: TraceManager, treeBuilderManager: TreeBuilderManager) extends RestHelper with Loggable {
 
 	import ActivityRequest._
 
@@ -289,14 +289,14 @@ class TraceAPIServer(manager: TraceManager) extends RestHelper with Loggable {
 
 		// GET the trace's package tree as json
 		case Paths.PackageTree(target) Get req =>
-			PackageTreeStreamer.streamPackageTree(target.treeBuilder.packageTree)
+			PackageTreeStreamer.streamPackageTree(treeBuilderManager.get(target.id).packageTree)
 
 		// stream a projected treemap based on the POSTed selected packages
 		case Paths.Treemap(target) Post req =>
 			req.param("nodes") match {
 				case Full(packages) =>
 					val ids = packages.split(',').flatMap(AsInt.unapply).toSet
-					val tree = target.treeBuilder.projectTree(ids)
+					val tree = treeBuilderManager.get(target.id).projectTree(ids)
 					TreemapDataStreamer.streamTreemapData(target.traceData.treeNodeData, tree)
 
 				case _ => BadResponse()

codepulse/src/main/scala/com/secdec/codepulse/tracer/TraceManager.scala

@@ -35,7 +35,6 @@ import com.secdec.codepulse.data.jsp.JasperJspMapper
 import com.secdec.codepulse.data.trace.TraceId
 import com.secdec.codepulse.data.trace.TraceDataProvider
 import com.secdec.codepulse.data.trace.TraceData
-import com.secdec.codepulse.data.trace.TreeBuilder
 
 object TraceManager {
 	lazy val defaultActorSystem = {
@@ -88,9 +87,7 @@ class TraceManager(val actorSystem: ActorSystem) extends Observing {
 	private def registerTrace(traceId: TraceId, traceData: TraceData, jspMapper: Option[JspMapper]) = {
 		registerTraceId(traceId)
 
-		val treeBuilder = new TreeBuilder(traceData.treeNodeData)
-
-		val target = AkkaTracingTarget(actorSystem, traceId, traceData, transientDataProvider get traceId, treeBuilder, jspMapper)
+		val target = AkkaTracingTarget(actorSystem, traceId, traceData, transientDataProvider get traceId, jspMapper)
 		traces.put(traceId, target)
 
 		// cause a traceListUpdate when this trace's name changes

codepulse/src/main/scala/com/secdec/codepulse/tracer/TreeBuilderManager.scala

@@ -0,0 +1,60 @@
+/*
+ * Code Pulse: A real-time code coverage testing tool. For more information
+ * see http://code-pulse.com
+ *
+ * Copyright (C) 2014 Applied Visions - http://securedecisions.avi.com
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package com.secdec.codepulse.tracer
+
+import com.secdec.codepulse.data.trace.TraceDataProvider
+import com.secdec.codepulse.data.trace.TraceId
+import com.secdec.codepulse.data.trace.TreeBuilder
+
+/** Manages a single TreeBuilder instance, such that we're only holding an
+  * entire trace tree for a single trace at a time.
+  *
+  * @author robertf
+  */
+class TreeBuilderManager(dataProvider: TraceDataProvider) {
+	private val lock = new Object
+	private var cachedId = None: Option[TraceId]
+	private var current = None: Option[TreeBuilder]
+
+	private def changeTrace(newTrace: TraceId) = {
+		cachedId = None
+		val data = dataProvider.getTrace(newTrace)
+		val builder = new TreeBuilder(data.treeNodeData)
+		current = Some(builder)
+		cachedId = Some(newTrace)
+		builder
+	}
+
+	def get(trace: TraceId): TreeBuilder = {
+		lock.synchronized {
+			if (cachedId.exists(_ == trace))
+				current.getOrElse(changeTrace(trace))
+			else
+				changeTrace(trace)
+		}
+	}
+
+	def clear() {
+		lock.synchronized {
+			cachedId = None
+			current = None
+		}
+	}
+}

codepulse/src/main/scala/com/secdec/codepulse/tracer/package.scala

@@ -43,20 +43,25 @@ package object tracer {
 	val traceManager = new BootVar[TraceManager]
 	val traceDataProvider = new BootVar[TraceDataProvider]
 	val transientTraceDataProvider = new BootVar[TransientTraceDataProvider]
+	val treeBuilderManager = new BootVar[TreeBuilderManager]
 	val traceFileUploadServer = new BootVar[TraceFileUploadHandler]
 	val traceAPIServer = new BootVar[TraceAPIServer]
 
 	def boot() {
 		val as = TraceManager.defaultActorSystem
 
-		traceDataProvider set new SlickH2TraceDataProvider(TraceDataProvider.DefaultStorageDir, as)
+		val dataProvider = new SlickH2TraceDataProvider(TraceDataProvider.DefaultStorageDir, as)
+		traceDataProvider set dataProvider
 		transientTraceDataProvider set new TransientTraceDataProvider
 
+		val tbm = new TreeBuilderManager(dataProvider)
+		treeBuilderManager set tbm
+
 		val tm = new TraceManager(as)
 
 		traceActorSystem set as
 		traceManager set tm
 		traceFileUploadServer set new TraceFileUploadHandler(tm).initializeServer
-		traceAPIServer set new TraceAPIServer(tm).initializeServer
+		traceAPIServer set new TraceAPIServer(tm, tbm).initializeServer
 	}
 }