Merge pull request #140 from eyalfa/fix_wrapped_callable_tasks_hierarchy

Fix wrapped callable tasks hierarchy

Author: edengal

timbermill-java/timbermill-api/pom.xml

@@ -16,7 +16,7 @@
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <java.source-target.version>1.8</java.source-target.version>
-        <timbermill.version>2.4.3</timbermill.version>
+        <timbermill.version>2.4.4</timbermill.version>
         <slf4jVersion>1.7.28</slf4jVersion>
     </properties>
 

timbermill-java/timbermill-api/src/main/java/com/datorama/oss/timbermill/EventLogger.java

@@ -180,23 +180,43 @@ String getCurrentTaskId() {
 		return taskIdStack.empty() ? null : taskIdStack.peek();
 	}
 
+	private class Scope implements AutoCloseable {
+		final Stack<String> origStack = taskIdStack;
+
+		Scope(String taskId) {
+			taskIdStack = new Stack<String>();
+			if(null != taskId)
+				taskIdStack.push(taskId);
+		}
+
+		<T> T apply(Callable<T> c) throws Exception {
+			return c.call();
+		}
+		<I, O> O apply(Function<I, O> f, I in) {
+			return f.apply(in);
+		}
+
+		@Override
+		public void close() {
+			taskIdStack = origStack;
+		}
+	}
+
 	<T> Callable<T> wrapCallable(Callable<T> callable) {
-		final Stack<String> origTaskIdStack = taskIdStack;
+		final String currentTaskId = getCurrentTaskId();
 		return () -> {
-			get().taskIdStack = (Stack<String>) origTaskIdStack.clone();
-			T call = callable.call();
-			get().taskIdStack.clear();
-			return call;
+			try(Scope scope = new Scope(currentTaskId)) {
+				return scope.apply(callable);
+			}
 		};
 	}
 
 	<T, R> Function<T, R> wrapFunction(Function<T, R> function) {
-		final Stack<String> origTaskIdStack = taskIdStack;
-		return t -> {
-			get().taskIdStack = (Stack<String>) origTaskIdStack.clone();
-			R call = function.apply(t);
-			get().taskIdStack.clear();
-			return call;
+		final String currentTaskId = getCurrentTaskId();
+		return (in) -> {
+			try(Scope scope = new Scope(currentTaskId)) {
+				return scope.apply(function, in);
+			}
 		};
 	}
 

timbermill-java/timbermill-api/src/test/java/com/datorama/oss/timbermill/EventLoggerTest.java

@@ -1,10 +1,13 @@
 package com.datorama.oss.timbermill;
 
 import java.util.List;
+import java.util.concurrent.Callable;
 import java.util.concurrent.TimeUnit;
+import java.util.function.Function;
 
 import org.awaitility.Awaitility;
 import org.junit.After;
+import org.junit.Before;
 import org.junit.BeforeClass;
 import org.junit.Test;
 
@@ -39,6 +42,7 @@ public static void init(){
 
 	@After
 	public void tearDown() {
+		el.clearStack();
 		EventLogger.exit();
 		mockPipe.close();
 	}
@@ -60,6 +64,48 @@ public void testSimpleEventLogger() {
 		assertEquals(TEST, startEvent.getStrings().get(BOOTSTRAP));
 	}
 
+	@Test
+	public void testWrapCallable() throws Exception {
+		String parentTaskId = el.startEvent(QUERY, EMPTY_LOG_PARAMS);
+		Callable<String> wrappedCallable = el.wrapCallable(() -> {
+			return el.getCurrentTaskId();
+		});
+
+		String childTaskId = el.startEvent(QUERY, EMPTY_LOG_PARAMS);
+
+		String callableVisibleTaskId = wrappedCallable.call();
+
+		assertEquals(callableVisibleTaskId, parentTaskId);
+	}
+
+	@Test
+	public void testWrapCallableOnEmptyStack() throws Exception {
+		Callable<String> wrappedCallable = el.wrapCallable(() -> {
+			String parentTaskId = TimberLogger.getCurrentTaskId();
+			String selfTaskId = TimberLogger.start("some_nifty_task");
+			TimberLogger.success();
+			return parentTaskId;
+		});
+
+		String callableVisibleTaskId = wrappedCallable.call();
+
+		assertNull(callableVisibleTaskId);
+	}
+
+	@Test
+	public void testWrapFunction() throws Exception {
+		String parentTaskId = el.startEvent(QUERY, EMPTY_LOG_PARAMS);
+		Function<Integer, String> wrappedFunction = el.<Integer, String>wrapFunction(ignored -> {
+			return el.getCurrentTaskId();
+		});
+
+		String childTaskId = el.startEvent(QUERY, EMPTY_LOG_PARAMS);
+
+		String callableVisibleTaskId = wrappedFunction.apply(42);
+
+		assertEquals(callableVisibleTaskId, parentTaskId);
+	}
+
 	@Test
 	public void testFailedEventLogger() {
 		el.startEvent(QUERY, EMPTY_LOG_PARAMS);

timbermill-java/timbermill-client/pom.xml

@@ -15,7 +15,7 @@
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <java.source-target.version>1.8</java.source-target.version>
-        <timbermill.version>2.4.3</timbermill.version>
+        <timbermill.version>2.4.4</timbermill.version>
         <aspectj.version>1.9.6</aspectj.version>
         <org.springframework.version>5.3.4</org.springframework.version>
     </properties>
@@ -124,7 +124,7 @@
                 <plugin>
                     <groupId>org.codehaus.mojo</groupId>
                     <artifactId>aspectj-maven-plugin</artifactId>
-                    <version>1.7</version>
+                    <version>1.14.0</version>
                     <configuration>
                         <showWeaveInfo>true</showWeaveInfo>
                         <source>${java.source-target.version}</source>

timbermill-java/timbermill-local/pom.xml

@@ -14,7 +14,7 @@
 
     <properties>
         <java.source-target.version>1.8</java.source-target.version>
-        <timbermill.version>2.4.3</timbermill.version>
+        <timbermill.version>2.4.4</timbermill.version>
         <kamon.bundle.version>2.0.4</kamon.bundle.version>
     </properties>
 

timbermill-java/timbermill-server/pom.xml

@@ -6,6 +6,7 @@
         <groupId>org.springframework.boot</groupId>
         <artifactId>spring-boot-starter-parent</artifactId>
         <version>2.3.3.RELEASE</version>
+        <relativePath/>
     </parent>
     <modelVersion>4.0.0</modelVersion>
 
@@ -16,7 +17,7 @@
     <properties>
         <project.build.sourceEncoding>UTF-8</project.build.sourceEncoding>
         <java.source-target.version>1.8</java.source-target.version>
-        <timbermill.version>2.4.3</timbermill.version>
+        <timbermill.version>2.4.4</timbermill.version>
         <slf4jVersion>1.7.28</slf4jVersion>
     </properties>