Fixing memory leak in SpanStore

michaelhyatt · May 8, 2022 · e7af3eb · e7af3eb
1 parent b1a88e5
commit e7af3eb
Show file tree

Hide file tree

Showing 12 changed files with 60 additions and 51 deletions.
diff --git a/apm-mule3-agent/pom.xml b/apm-mule3-agent/pom.xml
@@ -5,7 +5,7 @@
 	<groupId>co.elastic.apm</groupId>
 	<artifactId>apm-mule3-agent</artifactId>
 
-	<version>1.19.1</version>
+	<version>1.21.1</version>
 	<packaging>${packaging}</packaging>
 	<name>Mule apm-mule3-agent Application</name>
 
@@ -34,7 +34,7 @@
 		<mule.version>3.9.0</mule.version>
 		<mule.tools.version>1.2</mule.tools.version>
 
-		<elastic-apm.version>1.19.0</elastic-apm.version>
+		<elastic-apm.version>1.21.0</elastic-apm.version>
 		<version.byte-buddy>1.10.8</version.byte-buddy>
 
 		<maven.compiler.source>1.8</maven.compiler.source>

diff --git a/apm-mule3-agent/src/main/java/co/elastic/apm/mule/utils/SpanStore.java b/apm-mule3-agent/src/main/java/co/elastic/apm/mule/utils/SpanStore.java
@@ -9,6 +9,7 @@
 import org.mule.api.context.notification.ServerNotification;
 import org.mule.api.processor.MessageProcessor;
 import org.mule.context.notification.MessageProcessorNotification;
+import org.python.modules.synchronize;
 
 import co.elastic.apm.api.Span;
 
@@ -22,7 +23,7 @@
  */
 public class SpanStore {
 
-	private Map<String, Map<Optional<MessageProcessor>, Span>> map = Collections.synchronizedMap(new WeakHashMap<>());
+	private static Map<String, Map<Optional<MessageProcessor>, Span>> map = Collections.synchronizedMap(new WeakHashMap<String, Map<Optional<MessageProcessor>, Span>>());
 
 	/**
 	 * Store a {@link co.elastic.apm.api.Span} or a
@@ -38,13 +39,15 @@ public void storeTransactionOrSpan(String key, ServerNotification notification,
 
 		Optional<MessageProcessor> key2 = getKey2(notification);
 
-		Map<Optional<MessageProcessor>, Span> innerMap = map.get(key);
+		synchronized(map) {
+			Map<Optional<MessageProcessor>, Span> innerMap = map.get(key);
 
-		if (innerMap == null)
-			innerMap = new HashMap<>();
+			if (innerMap == null)
+				innerMap = new HashMap<>();
 
-		innerMap.put(key2, span);
-		map.put(key, innerMap);
+			innerMap.put(key2, span);
+			map.put(key, innerMap);
+		}
 	}
 
 	/**
@@ -59,11 +62,17 @@ public Span getTransactionOrSpan(String key, ServerNotification notification) {
 
 		Optional<MessageProcessor> key2 = getKey2(notification);
 
-		Map<Optional<MessageProcessor>, Span> stack = map.get(key);
-		Span span = stack.remove(key2);
+		Span span = null;
+
+		synchronized(map) {
+			Map<Optional<MessageProcessor>, Span> stack = map.get(key);
+			span = stack.remove(key2);
 
-		if (stack.size() == 0)
-			map.remove(key);
+			if (stack.size() == 0)
+				map.remove(key);
+			else
+				map.put(key, stack);
+		}
 
 		return span;
 	}

diff --git a/apm-mule3.8-agent/pom.xml b/apm-mule3.8-agent/pom.xml
@@ -5,7 +5,7 @@
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>co.elastic.apm</groupId>
 	<artifactId>apm-mule3.8-agent</artifactId>
-	<version>1.19.1</version>
+	<version>1.21.1</version>
 	<packaging>${packaging}</packaging>
 	<name>Mule apm-mule3-agent Application built for Mule 3.8</name>
 
@@ -34,7 +34,7 @@
 		<mule.version>3.8.1</mule.version>
 		<mule.tools.version>1.2</mule.tools.version>
 
-		<elastic-apm.version>1.19.0</elastic-apm.version>
+		<elastic-apm.version>1.21.0</elastic-apm.version>
 		<version.byte-buddy>1.10.8</version.byte-buddy>
 
 		<maven.compiler.source>1.8</maven.compiler.source>

diff --git a/apm-mule3.8-agent/src/main/java/co/elastic/apm/mule/utils/SpanStore.java b/apm-mule3.8-agent/src/main/java/co/elastic/apm/mule/utils/SpanStore.java
@@ -9,6 +9,7 @@
 import org.mule.api.context.notification.ServerNotification;
 import org.mule.api.processor.MessageProcessor;
 import org.mule.context.notification.MessageProcessorNotification;
+import org.python.modules.synchronize;
 
 import co.elastic.apm.api.Span;
 
@@ -22,29 +23,31 @@
  */
 public class SpanStore {
 
-	private Map<String, Map<Optional<MessageProcessor>, Span>> map = Collections.synchronizedMap(new WeakHashMap<>());
+	private static Map<String, Map<Optional<MessageProcessor>, Span>> map = Collections.synchronizedMap(new WeakHashMap<String, Map<Optional<MessageProcessor>, Span>>());
 
 	/**
 	 * Store a {@link co.elastic.apm.api.Span} or a
 	 * {@link co.elastic.apm.api.Transaction} in a {@link java.util.Map} using a
 	 * {@link String} key
 	 * 
 	 * @param key
-	 *            rootMessageId from MuleMessage
+	 *             rootMessageId from MuleMessage
 	 * @param span
-	 *            Span or Transaction object to store
+	 *             Span or Transaction object to store
 	 */
 	public void storeTransactionOrSpan(String key, ServerNotification notification, Span span) {
 
 		Optional<MessageProcessor> key2 = getKey2(notification);
 
-		Map<Optional<MessageProcessor>, Span> innerMap = map.get(key);
+		synchronized(map) {
+			Map<Optional<MessageProcessor>, Span> innerMap = map.get(key);
 
-		if (innerMap == null)
-			innerMap = new HashMap<>();
+			if (innerMap == null)
+				innerMap = new HashMap<>();
 
-		innerMap.put(key2, span);
-		map.put(key, innerMap);
+			innerMap.put(key2, span);
+			map.put(key, innerMap);
+		}
 	}
 
 	/**
@@ -59,12 +62,18 @@ public Span getTransactionOrSpan(String key, ServerNotification notification) {
 
 		Optional<MessageProcessor> key2 = getKey2(notification);
 
-		Map<Optional<MessageProcessor>, Span> stack = map.get(key);
-		Span span = stack.remove(key2);
-
-		if (stack.size() == 0)
-			map.remove(key);
-
+		Span span = null;
+
+		synchronized(map) {
+			Map<Optional<MessageProcessor>, Span> stack = map.get(key);
+			span = stack.remove(key2);
+
+			if (stack.size() == 0)
+				map.remove(key);
+			else
+				map.put(key, stack);
+		}
+
 		return span;
 	}
 
@@ -74,7 +83,7 @@ public Span getTransactionOrSpan(String key, ServerNotification notification) {
 	 * can throw NullPointerException, if called on an empty map.
 	 * 
 	 * @param notification
-	 *            typically, rootMessageId from MuleMessage
+	 *                     typically, rootMessageId from MuleMessage
 	 * @return
 	 */
 	public int depth(String key) {
@@ -99,7 +108,7 @@ public Span getTopLevelTransaction(String key) {
 	private Optional<MessageProcessor> getKey2(ServerNotification notification) {
 
 		Optional<MessageProcessor> key2 = Optional.empty();
-		
+
 		if (notification instanceof MessageProcessorNotification) {
 			MessageProcessor processor = ((MessageProcessorNotification) notification).getProcessor();
 			key2 = Optional.of(processor);

diff --git a/example-app/component1/.classpath b/example-app/component1/.classpath
@@ -6,12 +6,7 @@
   <classpathentry kind="src" path="src/test/java" output="target/test-classes"/>
   <classpathentry kind="src" path="src/test/resources" output="target/test-classes" excluding="**/*.java"/>
   <classpathentry kind="output" path="target/classes"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-agent-api/1.19.0/apm-agent-api-1.19.0.jar"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-agent-attach/1.19.0/apm-agent-attach-1.19.0.jar"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-mule3-agent/1.19.1/apm-mule3-agent-1.19.1.jar"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/elastic-apm-agent/1.19.0/elastic-apm-agent-1.19.0.jar"/>
-  <classpathentry kind="var" path="M2_REPO/net/java/dev/jna/jna/5.3.1/jna-5.3.1.jar"/>
-  <classpathentry kind="var" path="M2_REPO/net/java/dev/jna/jna-platform/5.3.1/jna-platform-5.3.1.jar"/>
+  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-mule3-agent/1.21.1/apm-mule3-agent-1.21.1.jar"/>
   <classpathentry kind="con" path="MULE_RUNTIME"/>
   <classpathentry kind="con" path="org.eclipse.jdt.launching.JRE_CONTAINER"/>
 </classpath>
diff --git a/example-app/component1/mule-project.xml b/example-app/component1/mule-project.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
-<mule-project xmlns="http://www.mulesoft.com/tooling/project" runtimeId="org.mule.tooling.server.3.9.0.ee" schemaVersion="5.3.0.0">
+<mule-project xmlns="http://www.mulesoft.com/tooling/project" runtimeId="org.mule.tooling.server.3.9.0" schemaVersion="5.3.0.0">
     <name>component1</name>
     <description></description>
 </mule-project>
diff --git a/example-app/component1/pom.xml b/example-app/component1/pom.xml
@@ -1,6 +1,5 @@
 <?xml version="1.0" encoding="UTF-8" standalone="no"?>
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-	xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
 
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>com.mycompany</groupId>
@@ -16,7 +15,7 @@
 		<mule.version>3.9.0</mule.version>
 		<mule.tools.version>1.7</mule.tools.version>
 
-		<apm-mule3-agent.version>1.19.1</apm-mule3-agent.version>
+		<apm-mule3-agent.version>1.21.1</apm-mule3-agent.version>
 
 		<maven.test.skip>true</maven.test.skip>
 	</properties>

diff --git a/example-app/component1/src/main/app/mule-app.properties b/example-app/component1/src/main/app/mule-app.properties
@@ -19,6 +19,7 @@ elastic.apm.service_name=component1
 elastic.apm.service_version=v1.0.0
 elastic.apm.stack_trace_limit=3
 elastic.apm.span_frames_min_duration=0ms
+elastic.apm.cloud_provider=NONE
 
 # Mule apm specific properties
 elastic.apm.mule.capture_input_properties=true

diff --git a/example-app/component2/.classpath b/example-app/component2/.classpath
@@ -6,12 +6,7 @@
   <classpathentry kind="src" path="src/test/java" output="target/test-classes"/>
   <classpathentry kind="src" path="src/test/resources" output="target/test-classes" excluding="**/*.java"/>
   <classpathentry kind="output" path="target/classes"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-agent-api/1.19.0/apm-agent-api-1.19.0.jar"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-agent-attach/1.19.0/apm-agent-attach-1.19.0.jar"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-mule3-agent/1.19.1/apm-mule3-agent-1.19.1.jar"/>
-  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/elastic-apm-agent/1.19.0/elastic-apm-agent-1.19.0.jar"/>
-  <classpathentry kind="var" path="M2_REPO/net/java/dev/jna/jna/5.3.1/jna-5.3.1.jar"/>
-  <classpathentry kind="var" path="M2_REPO/net/java/dev/jna/jna-platform/5.3.1/jna-platform-5.3.1.jar"/>
+  <classpathentry kind="var" path="M2_REPO/co/elastic/apm/apm-mule3-agent/1.21.1/apm-mule3-agent-1.21.1.jar"/>
   <classpathentry kind="con" path="MULE_RUNTIME"/>
   <classpathentry kind="con" path="org.eclipse.jdt.launching.JRE_CONTAINER"/>
 </classpath>
diff --git a/example-app/component2/mule-project.xml b/example-app/component2/mule-project.xml
@@ -1,5 +1,5 @@
 <?xml version="1.0" encoding="UTF-8" standalone="yes"?>
-<mule-project xmlns="http://www.mulesoft.com/tooling/project" runtimeId="org.mule.tooling.server.3.9.0.ee" schemaVersion="5.3.0.0">
+<mule-project xmlns="http://www.mulesoft.com/tooling/project" runtimeId="org.mule.tooling.server.3.9.0" schemaVersion="5.3.0.0">
     <name>component2</name>
     <description></description>
 </mule-project>
diff --git a/example-app/component2/pom.xml b/example-app/component2/pom.xml
@@ -1,6 +1,5 @@
-<?xml version="1.0" encoding="UTF-8"?>
-<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
-	xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
+<?xml version="1.0" encoding="UTF-8" standalone="no"?>
+<project xmlns="http://maven.apache.org/POM/4.0.0" xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance" xsi:schemaLocation="http://maven.apache.org/POM/4.0.0 http://maven.apache.org/maven-v4_0_0.xsd">
 
 	<modelVersion>4.0.0</modelVersion>
 	<groupId>com.mycompany</groupId>
@@ -16,7 +15,7 @@
 		<mule.version>3.9.0</mule.version>
 		<mule.tools.version>1.7</mule.tools.version>
 
-		<apm-mule3-agent.version>1.19.1</apm-mule3-agent.version>
+		<apm-mule3-agent.version>1.21.1</apm-mule3-agent.version>
 
 		<maven.test.skip>true</maven.test.skip>
 	</properties>

diff --git a/example-app/component2/src/main/app/mule-app.properties b/example-app/component2/src/main/app/mule-app.properties
@@ -10,6 +10,8 @@ elastic.apm.service_name=component2
 elastic.apm.service_version=v1.0.0
 elastic.apm.stack_trace_limit=3
 elastic.apm.span_frames_min_duration=0ms
+elastic.apm.cloud_provider=NONE
+
 
 # Mule apm specific properties
 elastic.apm.mule.capture_input_properties=true