Merge pull request #3 from stgatilov/master

General performance changes & bugfix

Heapy/Heapy.cpp

@@ -98,7 +98,7 @@ extern "C" int main(int argc, char* argv[]){
 	// Start our new process with a suspended main thread.
 	std::cout << "Starting process with heap profiling enabled..." << std::endl;
 	std::cout << "Target exe path: " << injectionTarget << std::endl;
-	std::cout << "Target execommand line: " << commandLine;
+	std::cout << "Target execommand line: " << commandLine << std::endl;
 	std::cout << "Dll to inject: " << dllPath << std::endl;
 
 
@@ -109,6 +109,7 @@ extern "C" int main(int argc, char* argv[]){
 
 	// CreatePRocessA can modify input arg so do this to be safe.
 	std::vector<char> commandLineMutable(commandLine.begin(), commandLine.end()); 
+	commandLineMutable.push_back(0);
 
 	if(CreateProcessA(NULL, commandLineMutable.data(), NULL, NULL, 0, flags, NULL, 
 		             (LPSTR)".", &si, &pi) == 0){

HeapyInject/HeapProfiler.cpp

@@ -12,7 +12,13 @@ StackTrace::StackTrace() : hash(0){
 }
 
 void StackTrace::trace(){
-	CaptureStackBackTrace(0, backtraceSize, backtrace, &hash);
+	int framesCnt = CaptureStackBackTrace(0, backtraceSize, backtrace, 0);
+	// Compute simple polynomial hash of the stack trace.
+	// Note: CaptureStackBackTrace returns plain sum of all pointers as BackTraceHash.
+	const size_t BASE = sizeof(size_t) > 4 ? 11400714819323198485ULL : 2654435769U;
+	hash = 0;
+	for (int i = 0; i < framesCnt; i++)
+		hash = hash * BASE + (size_t)backtrace[i];
 }
 
 void StackTrace::print(std::ostream &stream) const {
@@ -54,16 +60,23 @@ void StackTrace::print(std::ostream &stream) const {
 void HeapProfiler::malloc(void *ptr, size_t size, const StackTrace &trace){
 	std::lock_guard<std::mutex> lk(mutex);
 
+	if (ptrs.find(ptr) != ptrs.end())
+		return;   //two buffers at same address!
+
 	// Locate or create this stacktrace in the allocations map.
 	if(stackTraces.find(trace.hash) == stackTraces.end()){
-		stackTraces[trace.hash].trace = trace;
+		auto &stack = stackTraces[trace.hash];
+		stack.trace = trace;
+		stack.totalSize = 0;
 	}
 
 	// Store the size for this allocation this stacktraces allocation map.
-	stackTraces[trace.hash].allocations[ptr] = size;
+	stackTraces[trace.hash].totalSize += size;
 
 	// Store the stracktrace hash of this allocation in the pointers map.
-	ptrs[ptr] = trace.hash;
+	auto &ptrInfo = ptrs[ptr];
+	ptrInfo.size = size;
+	ptrInfo.stack = trace.hash;
 }
 
 void HeapProfiler::free(void *ptr, const StackTrace &trace){
@@ -73,8 +86,8 @@ void HeapProfiler::free(void *ptr, const StackTrace &trace){
 	// allocating stack traces map.
 	auto it = ptrs.find(ptr);
 	if(it != ptrs.end()){
-		StackHash stackHash = it->second;
-		stackTraces[stackHash].allocations.erase(ptr); 
+		const PointerInfo &info = it->second;
+		stackTraces[info.stack].totalSize -= info.size;
 		ptrs.erase(it);
 	}else{
 		// Do anything with wild pointer frees?
@@ -85,14 +98,8 @@ void HeapProfiler::getAllocationSiteReport(std::vector<std::pair<StackTrace, siz
 	std::lock_guard<std::mutex> lk(mutex);
 	allocs.clear();
 
-	// For each allocation point.
-	for(auto &traceInfo : stackTraces){
-		// Sum up the size of all the allocations made.
-		size_t sumOfAlloced = 0;
-		for(auto &alloc : traceInfo.second.allocations)
-			sumOfAlloced += alloc.second;
-
-		// Add to alloation site report.
-		allocs.push_back(std::make_pair(traceInfo.second.trace, sumOfAlloced));
+	for(auto it = stackTraces.begin(); it != stackTraces.end(); it++){
+		const auto &info = it->second;
+		allocs.push_back(std::make_pair(info.trace, info.totalSize));
 	}
 }

HeapyInject/HeapProfiler.h

@@ -6,7 +6,7 @@
 #include <mutex>
 
 const int backtraceSize = 64;
-typedef unsigned long StackHash;
+typedef size_t StackHash;
 
 struct StackTrace{
 	void *backtrace[backtraceSize];
@@ -27,11 +27,16 @@ class HeapProfiler{
 	void getAllocationSiteReport(std::vector<std::pair<StackTrace, size_t>> &allocs);
 private:
 	std::mutex mutex;
-	struct TraceInfo{
+	struct CallStackInfo {
 		StackTrace trace;
-		std::unordered_map<void *, size_t> allocations;
+		size_t totalSize;
 	};
-	std::unordered_map<StackHash, TraceInfo> stackTraces;
-	std::unordered_map<void*, StackHash> ptrs;
+	struct PointerInfo {
+		StackHash stack;
+		size_t size;
+	};
+
+	std::unordered_map<StackHash, CallStackInfo> stackTraces;
+	std::unordered_map<void*, PointerInfo> ptrs;
 
 };

HeapyInject/HeapyInject.cpp

@@ -77,7 +77,7 @@ void  __cdecl freeHook(void * p){
 	originalFrees[N](p);
 	if(preventSelfProfile.shouldProfile()){
 		StackTrace trace;
-		trace.trace();
+		//trace.trace();
 		heapProfiler->free(p, trace);
 	}
 }