Adds TracedScheduledExecutorService

src/main/java/io/opentracing/contrib/concurrent/TracedScheduledExecutorService.java

@@ -0,0 +1,45 @@
+package io.opentracing.contrib.concurrent;
+
+import io.opentracing.Tracer;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.ScheduledExecutorService;
+import java.util.concurrent.ScheduledFuture;
+import java.util.concurrent.TimeUnit;
+
+/**
+ * @author Jose Montoya
+ */
+public class TracedScheduledExecutorService extends TracedExecutorService implements ScheduledExecutorService {
+
+	private final ScheduledExecutorService delegate;
+
+	public TracedScheduledExecutorService(ScheduledExecutorService delegate, Tracer tracer) {
+		super(delegate, tracer);
+		this.delegate = delegate;
+	}
+
+	@Override
+	public ScheduledFuture<?> schedule(Runnable runnable, long delay, TimeUnit timeUnit) {
+		return delegate.schedule(tracer.activeSpan() == null ? runnable :
+				new TracedRunnable(runnable, tracer), delay, timeUnit);
+	}
+
+	@Override
+	public <T> ScheduledFuture<T> schedule(Callable<T> callable, long delay, TimeUnit timeUnit) {
+		return delegate.schedule(tracer.activeSpan() == null ? callable :
+				new TracedCallable<T>(callable, tracer), delay, timeUnit);
+	}
+
+	@Override
+	public ScheduledFuture<?> scheduleAtFixedRate(Runnable runnable, long initialDelay, long period, TimeUnit timeUnit) {
+		return delegate.scheduleAtFixedRate(tracer.activeSpan() == null ? runnable :
+				new TracedRunnable(runnable, tracer), initialDelay, period, timeUnit);
+	}
+
+	@Override
+	public ScheduledFuture<?> scheduleWithFixedDelay(Runnable runnable, long initialDelay, long delay, TimeUnit timeUnit) {
+		return delegate.scheduleWithFixedDelay(tracer.activeSpan() == null ? runnable :
+				new TracedRunnable(runnable, tracer), initialDelay, delay, timeUnit);
+	}
+}

src/test/java/io/opentracing/contrib/concurrent/TracedScheduledExecutorServiceTest.java

@@ -0,0 +1,75 @@
+package io.opentracing.contrib.concurrent;
+
+import io.opentracing.mock.MockSpan;
+import org.junit.Test;
+
+import java.util.concurrent.*;
+
+import static org.junit.Assert.assertEquals;
+
+/**
+ * @author Jose Montoya
+ */
+public class TracedScheduledExecutorServiceTest extends AbstractConcurrentTest {
+	private static final int NUMBER_OF_THREADS = 4;
+
+	protected ScheduledExecutorService toTraced(ScheduledExecutorService scheduledExecutorService) {
+		return new TracedScheduledExecutorService(scheduledExecutorService, mockTracer);
+	}
+
+	@Test
+	public void scheduleRunnableTest() throws InterruptedException {
+		ScheduledExecutorService executorService = toTraced(Executors.newScheduledThreadPool(NUMBER_OF_THREADS));
+
+		MockSpan parentSpan = mockTracer.buildSpan("foo").startManual();
+		mockTracer.scopeManager().activate(parentSpan, true);
+		executorService.schedule(new TestRunnable(), 300, TimeUnit.MILLISECONDS);
+
+		countDownLatch.await();
+		assertParentSpan(parentSpan);
+		assertEquals(1, mockTracer.finishedSpans().size());
+	}
+
+	@Test
+	public void scheduleCallableTest() throws InterruptedException {
+		ScheduledExecutorService executorService = toTraced(Executors.newScheduledThreadPool(NUMBER_OF_THREADS));
+
+		MockSpan parentSpan = mockTracer.buildSpan("foo").startManual();
+		mockTracer.scopeManager().activate(parentSpan, true);
+		executorService.schedule(new TestCallable(), 300, TimeUnit.MILLISECONDS);
+
+		countDownLatch.await();
+		assertParentSpan(parentSpan);
+		assertEquals(1, mockTracer.finishedSpans().size());
+	}
+
+	@Test
+	public void scheduleAtFixedRateTest() throws InterruptedException {
+		countDownLatch = new CountDownLatch(2);
+		ScheduledExecutorService executorService = toTraced(Executors.newScheduledThreadPool(NUMBER_OF_THREADS));
+
+		MockSpan parentSpan = mockTracer.buildSpan("foo").startManual();
+		mockTracer.scopeManager().activate(parentSpan, true);
+		executorService.scheduleAtFixedRate(new TestRunnable(), 0, 300, TimeUnit.MILLISECONDS);
+
+		countDownLatch.await();
+		executorService.shutdown();
+		assertParentSpan(parentSpan);
+		assertEquals(2, mockTracer.finishedSpans().size());
+	}
+
+	@Test
+	public void scheduleWithFixedDelayTest() throws InterruptedException {
+		countDownLatch = new CountDownLatch(2);
+		ScheduledExecutorService executorService = toTraced(Executors.newScheduledThreadPool(NUMBER_OF_THREADS));
+
+		MockSpan parentSpan = mockTracer.buildSpan("foo").startManual();
+		mockTracer.scopeManager().activate(parentSpan, true);
+		executorService.scheduleWithFixedDelay(new TestRunnable(), 0, 300, TimeUnit.MILLISECONDS);
+
+		countDownLatch.await();
+		executorService.shutdown();
+		assertParentSpan(parentSpan);
+		assertEquals(2, mockTracer.finishedSpans().size());
+	}
+}