Merge 338c92f into 2a6422c

src/main/java/io/opentracing/contrib/hazelcast/TracingCallable.java

@@ -17,12 +17,21 @@
 import static io.opentracing.contrib.hazelcast.TracingHelper.extract;
 import static io.opentracing.contrib.hazelcast.TracingHelper.nullableClass;
 
+import com.hazelcast.core.HazelcastInstance;
+import com.hazelcast.core.HazelcastInstanceAware;
+import com.hazelcast.instance.Node;
+import com.hazelcast.spi.NodeAware;
+import com.hazelcast.spi.serialization.SerializationService;
+import com.hazelcast.spi.serialization.SerializationServiceAware;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
+
+import java.io.Serializable;
 import java.util.Map;
 import java.util.concurrent.Callable;
 
-public class TracingCallable<V> implements Callable<V> {
+public class TracingCallable<V> implements Callable<V>, Serializable, HazelcastInstanceAware,
+        NodeAware, SerializationServiceAware {
 
   private final Callable<V> callable;
   private final boolean traceWithActiveSpanOnly;
@@ -42,4 +51,24 @@ public V call() throws Exception {
     span.setTag("callable", nullableClass(callable));
     return decorateExceptionally2(callable::call, span);
   }
+
+  @Override
+  public void setHazelcastInstance(HazelcastInstance hazelcastInstance) {
+    if (callable instanceof HazelcastInstanceAware) {
+      ((HazelcastInstanceAware) callable).setHazelcastInstance(hazelcastInstance);
+    }
+  }
+
+  @Override public void setNode(Node node) {
+    if (callable instanceof NodeAware) {
+      ((NodeAware) callable).setNode(node);
+    }
+  }
+
+  @Override public void setSerializationService(
+          SerializationService serializationService) {
+    if (callable instanceof SerializationServiceAware) {
+      ((SerializationServiceAware) callable).setSerializationService(serializationService);
+    }
+  }
 }

src/main/java/io/opentracing/contrib/hazelcast/TracingRunnable.java

@@ -17,12 +17,19 @@
 import static io.opentracing.contrib.hazelcast.TracingHelper.extract;
 import static io.opentracing.contrib.hazelcast.TracingHelper.nullableClass;
 
+import com.hazelcast.core.HazelcastInstance;
+import com.hazelcast.core.HazelcastInstanceAware;
+import com.hazelcast.instance.Node;
+import com.hazelcast.spi.NodeAware;
+import com.hazelcast.spi.serialization.SerializationService;
+import com.hazelcast.spi.serialization.SerializationServiceAware;
 import io.opentracing.Span;
 import io.opentracing.SpanContext;
 import java.io.Serializable;
 import java.util.Map;
 
-public class TracingRunnable implements Runnable, Serializable {
+public class TracingRunnable implements Runnable, Serializable,
+        HazelcastInstanceAware, NodeAware, SerializationServiceAware {
 
   private final Runnable runnable;
   private final boolean traceWithActiveSpanOnly;
@@ -44,5 +51,23 @@ public void run() {
     decorateAction(runnable::run, span);
   }
 
+  @Override
+  public void setHazelcastInstance(HazelcastInstance hazelcastInstance) {
+    if (runnable instanceof HazelcastInstanceAware) {
+      ((HazelcastInstanceAware) runnable).setHazelcastInstance(hazelcastInstance);
+    }
+  }
 
+  @Override public void setNode(Node node) {
+    if (runnable instanceof NodeAware) {
+      ((NodeAware) runnable).setNode(node);
+    }
+  }
+
+  @Override public void setSerializationService(
+          SerializationService serializationService) {
+    if (runnable instanceof SerializationServiceAware) {
+      ((SerializationServiceAware) runnable).setSerializationService(serializationService);
+    }
+  }
 }

src/test/java/io/opentracing/contrib/hazelcast/AwareCallable.java

@@ -0,0 +1,56 @@
+/*
+ * Copyright 2018-2019 The OpenTracing Authors
+ *
+ * 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 io.opentracing.contrib.hazelcast;
+
+import com.hazelcast.core.HazelcastInstance;
+import com.hazelcast.core.HazelcastInstanceAware;
+import com.hazelcast.instance.Node;
+import com.hazelcast.spi.NodeAware;
+import com.hazelcast.spi.serialization.SerializationService;
+import com.hazelcast.spi.serialization.SerializationServiceAware;
+
+import java.io.Serializable;
+import java.util.Objects;
+import java.util.concurrent.Callable;
+
+public class AwareCallable<V> implements Callable<V>, Serializable,
+        HazelcastInstanceAware, NodeAware, SerializationServiceAware {
+
+    private HazelcastInstance hazelcastInstance;
+    private Node node;
+    private SerializationService serializationService;
+
+    @Override
+    public void setHazelcastInstance(HazelcastInstance hazelcastInstance) {
+        this.hazelcastInstance = hazelcastInstance;
+    }
+
+    @Override public void setNode(Node node) {
+
+        this.node = node;
+    }
+
+    @Override public void setSerializationService(
+            SerializationService serializationService) {
+
+        this.serializationService = serializationService;
+    }
+
+    @Override public V call() {
+        Objects.requireNonNull(hazelcastInstance);
+        Objects.requireNonNull(node);
+        Objects.requireNonNull(serializationService);
+        return null;
+    }
+}

src/test/java/io/opentracing/contrib/hazelcast/AwareRunnable.java

@@ -0,0 +1,54 @@
+/*
+ * Copyright 2018-2019 The OpenTracing Authors
+ *
+ * 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 io.opentracing.contrib.hazelcast;
+
+import com.hazelcast.core.HazelcastInstance;
+import com.hazelcast.core.HazelcastInstanceAware;
+import com.hazelcast.instance.Node;
+import com.hazelcast.spi.NodeAware;
+import com.hazelcast.spi.serialization.SerializationService;
+import com.hazelcast.spi.serialization.SerializationServiceAware;
+
+import java.io.Serializable;
+import java.util.Objects;
+
+public class AwareRunnable implements Runnable, Serializable,
+        HazelcastInstanceAware, NodeAware, SerializationServiceAware {
+
+    private HazelcastInstance hazelcastInstance;
+    private Node node;
+    private SerializationService serializationService;
+
+    @Override
+    public void setHazelcastInstance(HazelcastInstance hazelcastInstance) {
+        this.hazelcastInstance = hazelcastInstance;
+    }
+
+    @Override public void setNode(Node node) {
+
+        this.node = node;
+    }
+
+    @Override public void setSerializationService(
+            SerializationService serializationService) {
+
+        this.serializationService = serializationService;
+    }
+
+    @Override public void run() {
+        Objects.requireNonNull(hazelcastInstance);
+        Objects.requireNonNull(node);
+        Objects.requireNonNull(serializationService);
+    }
+}

src/test/java/io/opentracing/contrib/hazelcast/TracingTest.java

@@ -296,6 +296,32 @@ public void testRingbuffer() throws Exception {
     assertNull(tracer.activeSpan());
   }
 
+  @Test
+  public void testAwareRunnable() {
+    IExecutorService ex = hazelcast.getExecutorService("executor");
+    ex.submit(new AwareRunnable());
+
+    await().atMost(15, TimeUnit.SECONDS).until(reportedSpansSize(), equalTo(2));
+
+    List<MockSpan> spans = tracer.finishedSpans();
+    assertEquals(2, spans.size());
+    checkSpans(spans);
+    assertNull(tracer.activeSpan());
+  }
+
+  @Test
+  public void testAwareCallable() {
+    IExecutorService ex = hazelcast.getExecutorService("executor");
+    ex.submit(new AwareCallable<>());
+
+    await().atMost(15, TimeUnit.SECONDS).until(reportedSpansSize(), equalTo(2));
+
+    List<MockSpan> spans = tracer.finishedSpans();
+    assertEquals(2, spans.size());
+    checkSpans(spans);
+    assertNull(tracer.activeSpan());
+  }
+
   private Callable<Integer> reportedSpansSize() {
     return () -> tracer.finishedSpans().size();
   }