Make imports more time- and memory- efficient.

Author: baffles

codepulse/src/main/scala/com/secdec/codepulse/data/bytecode/CodeForestBuilder.scala

@@ -44,16 +44,16 @@ class CodeForestBuilder {
 	}
 
 	def result = {
-		val lb = List.newBuilder[TreeNodeData]
-		for (root <- roots) root.visitTree { node =>
-			val id = node.id
-			val name = node.name
-			val parentId = node.parentId
-			val kind = node.kind
-			val size = node.size
-			lb += TreeNodeData(id, parentId, name, kind, size)
+		roots.toStream flatMap { root =>
+			root.streamTree map { node =>
+				val id = node.id
+				val name = node.name
+				val parentId = node.parentId
+				val kind = node.kind
+				val size = node.size
+				root -> TreeNodeData(id, parentId, name, kind, size)
+			}
 		}
-		lb.result()
 	}
 
 	/** Applies the `condensePathNodes` transformation to each root node

codepulse/src/main/scala/com/secdec/codepulse/data/bytecode/CodeTree2.scala

@@ -54,6 +54,11 @@ trait CodeTreeNode {
 		for (child <- children) child.visitTree(callback)
 	}
 
+	/* Stream this node and all of its descendants */
+	def streamTree(): Stream[CodeTreeNode] = {
+		this +: (children.toStream.flatMap(_.streamTree))
+	}
+
 	def findChild(predicate: CodeTreeNode => Boolean): Option[CodeTreeNode] = {
 		children.find(predicate)
 	}

codepulse/src/main/scala/com/secdec/codepulse/data/trace/TreeNodeData.scala

@@ -60,6 +60,9 @@ trait TreeNodeDataAccess {
 	def mapJsp(jspClass: String, nodeId: Int): Unit
 	def mapJsps(jsps: Iterable[(String, Int)]): Unit = jsps foreach { case (jspPath, nodeId) => mapJsp(jspPath, nodeId) }
 
+	case class BulkImportElement(data: TreeNodeData, traced: Option[Boolean])
+	def bulkImport(data: Iterable[BulkImportElement]): Unit
+
 	def storeNode(node: TreeNodeData): Unit
 	def storeNodes(nodes: Iterable[TreeNodeData]) = nodes foreach storeNode
 

codepulse/src/main/scala/com/secdec/codepulse/data/trace/TreeNodeImporter.scala

@@ -0,0 +1,57 @@
+/*
+ * 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.data.trace
+
+/** Helper utility for mass-importing tree node data into the data model.
+  *
+  * @author robertf
+  */
+class TreeNodeImporter(treeNodeData: TreeNodeDataAccess, bufferSize: Int) {
+	import treeNodeData.{ BulkImportElement => BufferRecord }
+
+	private val buffer = collection.mutable.ListBuffer.empty[BufferRecord]
+
+	private def checkAndFlush() {
+		if (buffer.size >= bufferSize)
+			flush
+	}
+
+	def flush() {
+		treeNodeData.bulkImport(buffer)
+		buffer.clear
+	}
+
+	def +=(data: TreeNodeData) { +=(data, None) }
+	def +=(data: TreeNodeData, traced: Boolean) { +=(data, Some(traced)) }
+	def +=(data: TreeNodeData, traced: Option[Boolean]) {
+		buffer += BufferRecord(data, traced)
+		checkAndFlush
+	}
+
+	def ++=(data: Iterable[(TreeNodeData, Option[Boolean])]) {
+		for (chunk <- data.grouped(bufferSize)) {
+			treeNodeData.bulkImport(chunk map { case (data, traced) => BufferRecord(data, traced) })
+		}
+	}
+}
+
+object TreeNodeImporter {
+	def apply(treeNodeData: TreeNodeDataAccess, bufferSize: Int = 2500) = new TreeNodeImporter(treeNodeData, bufferSize)
+}

codepulse/src/main/scala/com/secdec/codepulse/data/trace/slick/SlickTreeNodeDataAccess.scala

@@ -93,6 +93,11 @@ private[slick] class SlickTreeNodeDataAccess(dao: TreeNodeDataDao, db: Database)
 		dao storeJsps jsps
 	}
 
+	def bulkImport(data: Iterable[BulkImportElement]) = db withTransaction { implicit transaction =>
+		dao storeNodes data.map(_.data)
+		dao storeTracedValues data.map { element => element.data.id -> element.traced }
+	}
+
 	def storeNode(node: TreeNodeData) = db withTransaction { implicit transaction =>
 		dao storeNode node
 	}

codepulse/src/main/scala/com/secdec/codepulse/data/trace/slick/TreeNodeDataDao.scala

@@ -144,6 +144,13 @@ private[slick] class TreeNodeDataDao(val driver: JdbcProfile) {
 
 	def getTraced()(implicit session: Session) = tracedNodes.list
 
+	def storeTracedValues(values: Iterable[(Int, Option[Boolean])])(implicit session: Session) {
+		tracedNodes ++= values flatMap {
+			case (id, Some(traced)) => Some(id -> traced)
+			case _ => None
+		}
+	}
+
 	def updateTraced(id: Int, traced: Option[Boolean])(implicit session: Session) {
 		traced match {
 			case Some(traced) =>

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

@@ -20,22 +20,26 @@
 package com.secdec.codepulse.tracer
 
 import java.io.File
+
+import scala.concurrent.ExecutionContext.Implicits.global
+import scala.concurrent.Future
+
+import org.apache.commons.io.FilenameUtils
+
 import com.secdec.codepulse.data.bytecode.AsmVisitors
 import com.secdec.codepulse.data.bytecode.CodeForestBuilder
+import com.secdec.codepulse.data.bytecode.CodeTreeNodeKind
+import com.secdec.codepulse.data.jsp.JasperJspAdapter
+import com.secdec.codepulse.data.trace.TraceData
 import com.secdec.codepulse.data.trace.TraceId
+import com.secdec.codepulse.data.trace.TreeNodeData
+import com.secdec.codepulse.data.trace.TreeNodeImporter
 import com.secdec.codepulse.tracer.export.TraceImporter
-import com.secdec.codepulse.util.RichFile.RichFile
 import com.secdec.codepulse.util.ZipEntryChecker
+
 import net.liftweb.common.Box
 import net.liftweb.common.Failure
-import org.apache.commons.io.FilenameUtils
-import com.secdec.codepulse.data.jsp.JasperJspAdapter
-import com.secdec.codepulse.data.trace.TraceData
-import scala.concurrent.Future
-import scala.concurrent.ExecutionContext.Implicits.global
 import net.liftweb.common.Full
-import com.secdec.codepulse.data.bytecode.CodeTreeNodeKind
-import com.secdec.codepulse.data.trace.TreeNodeData
 
 object TraceUploadData {
 
@@ -142,22 +146,20 @@ object TraceUploadData {
 		if (treemapNodes.isEmpty) {
 			throw new NoSuchElementException("No method data found in analyzed upload file.")
 		} else {
-			val treeNodeData = traceData.treeNodeData
-			import treeNodeData.ExtendedTreeNodeData
-
-			val nodeMap = treemapNodes.map(n => n.id -> n).toMap
-			def getRoot(n: TreeNodeData): TreeNodeData = n.parentId match {
-				case Some(parent) => getRoot(nodeMap(parent))
-				case None => n
+			val importer = TreeNodeImporter(traceData.treeNodeData)
+
+			importer ++= treemapNodes map {
+				case (root, node) =>
+					node -> (node.kind match {
+						case CodeTreeNodeKind.Grp | CodeTreeNodeKind.Pkg => Some(tracedGroups contains root.name)
+						case _ => None
+					})
 			}
 
-			treeNodeData.storeNodes(treemapNodes)
-			for (n <- treemapNodes) n.traced = n.kind match {
-				case CodeTreeNodeKind.Grp | CodeTreeNodeKind.Pkg => Some(tracedGroups contains getRoot(n).label)
-				case _ => None
-			}
-			treeNodeData.mapJsps(jspCorrelations)
-			treeNodeData.mapMethodSignatures(methodCorrelations)
+			importer.flush
+
+			traceData.treeNodeData.mapJsps(jspCorrelations)
+			traceData.treeNodeData.mapMethodSignatures(methodCorrelations)
 
 			// The `creationDate` for trace data detected in this manner
 			// should use its default value ('now'), as this is when the data

codepulse/src/main/scala/com/secdec/codepulse/tracer/export/TraceImportReaderV1.scala

@@ -71,14 +71,7 @@ object TraceImportReaderV1 extends TraceImportReader with TraceImportHelpers wit
 			if (jp.nextToken != START_ARRAY)
 				throw new TraceImportException(s"Unexpected token ${jp.getCurrentToken}; expected START_ARRAY.")
 
-			case class BufferRecord(data: TreeNodeData, traced: Option[Boolean])
-			val buffer = collection.mutable.ListBuffer.empty[BufferRecord]
-			def flushBuffer() {
-				treeNodeData.storeNodes(buffer.map(_.data))
-				for (BufferRecord(d, traced) <- buffer) d.traced = traced
-				buffer.clear
-			}
-			def checkAndFlush() { if (buffer.size >= 500) flushBuffer }
+			val importer = TreeNodeImporter(treeNodeData)
 
 			while (jp.nextValue != END_ARRAY) {
 				if (jp.getCurrentToken != START_OBJECT)
@@ -118,18 +111,16 @@ object TraceImportReaderV1 extends TraceImportReader with TraceImportHelpers wit
 					}
 				}
 
-				buffer += BufferRecord(TreeNodeData(
+				importer += (TreeNodeData(
 					id getOrElse { throw new TraceImportException("Missing ID for tree node.") },
 					parentId,
 					label getOrElse { throw new TraceImportException("Missing label for tree node.") },
 					kind.getOrElse { throw new TraceImportException("Missing or invalid kind for tree node.") },
 					size),
 					traced)
-
-				checkAndFlush
 			}
 
-			flushBuffer
+			importer.flush
 
 			if (jp.nextToken != null)
 				throw new TraceImportException(s"Unexpected token ${jp.getCurrentToken}; expected EOF.")