code refactor and javadoc update

opensearch-project · Jul 27, 2023 · 27d5511 · 27d5511
1 parent aaf19a0
commit 27d5511
Show file tree

Hide file tree

Showing 3 changed files with 39 additions and 31 deletions.
diff --git a/...elemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/TraceEventListener.java b/...elemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/TraceEventListener.java
@@ -11,4 +11,4 @@
 /**
  * Combines both RunnableEventListener and SpanEventListener. Usually both are used in conjunction.
  */
-public interface /**/ TraceEventListener extends RunnableEventListener, SpanEventListener {}
+public interface TraceEventListener extends RunnableEventListener, SpanEventListener {}
diff --git a/...lemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/TraceEventsRunnable.java b/...lemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/TraceEventsRunnable.java
@@ -40,39 +40,14 @@ public class TraceEventsRunnable implements Runnable {
     @Override
     public void run() {
         try {
-            if (traceEventsService.isTracingEnabled()) {
-                Span span = traceEventsService.getTracer().getCurrentSpan();
-                // repeat it for all the spans in the hierarchy
-                while (span != null) {
-                    if (!span.hasEnded()) {
-                        Span finalSpan = span;
-                        traceEventsService.executeListeners(
-                            span,
-                            traceEventListener -> traceEventListener.onRunnableStart(finalSpan, Thread.currentThread())
-                        );
-                    }
-                    span = span.getParentSpan();
-                }
-            }
+           invokeOnRunnableStart(traceEventsService);
         } catch (Exception e) {
             logger.debug("Error in onRunnableStart", e);
         } finally {
             delegate.run();
         }
         try {
-            if (traceEventsService.isTracingEnabled()) {
-                Span span = traceEventsService.getTracer().getCurrentSpan();
-                while (span != null) {
-                    if (!span.hasEnded()) {
-                        Span finalSpan = span;
-                        traceEventsService.executeListeners(
-                            span,
-                            traceEventListener -> traceEventListener.onRunnableComplete(finalSpan, Thread.currentThread())
-                        );
-                    }
-                    span = span.getParentSpan();
-                }
-            }
+            invokeOnRunnableComplete(traceEventsService);
         } catch (Exception e) {
             logger.debug("Error in onRunnableEnd", e);
         }
@@ -86,4 +61,37 @@ public void run() {
     public Runnable unwrap() {
         return delegate;
     }
+
+    public static void invokeOnRunnableStart(TraceEventsService traceEventsService) {
+        if (traceEventsService.isTracingEnabled()) {
+            Span span = traceEventsService.getTracer().getCurrentSpan();
+            // repeat it for all the spans in the hierarchy
+            while (span != null) {
+                if (!span.hasEnded()) {
+                    Span finalSpan = span;
+                    traceEventsService.executeListeners(
+                        span,
+                        traceEventListener -> traceEventListener.onRunnableStart(finalSpan, Thread.currentThread())
+                    );
+                }
+                span = span.getParentSpan();
+            }
+        }
+    }
+
+    public static void invokeOnRunnableComplete(TraceEventsService traceEventsService) {
+        if (traceEventsService.isTracingEnabled()) {
+            Span span = traceEventsService.getTracer().getCurrentSpan();
+            while (span != null) {
+                if (!span.hasEnded()) {
+                    Span finalSpan = span;
+                    traceEventsService.executeListeners(
+                        span,
+                        traceEventListener -> traceEventListener.onRunnableComplete(finalSpan, Thread.currentThread())
+                    );
+                }
+                span = span.getParentSpan();
+            }
+        }
+    }
 }
diff --git a/...elemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/TraceEventsService.java b/...elemetry/src/main/java/org/opensearch/telemetry/tracing/listeners/TraceEventsService.java
@@ -130,8 +130,8 @@ public boolean isTracingEnabled() {
 
     /**
      * Wraps the given Runnable with trace event listeners registered with {@link TraceEventsService}
-     * Note: Runnable should be wrapped using this method only after context has been restored so that when
-     * {@link TraceEventsRunnable#run()} is called, it has the right context with current Span information.
+     * Note: Runnable should be wrapped using this method only after thread context has been restored so that when
+     * {@link TraceEventsRunnable#run()} is called, it has the right thread context with current Span information.
      * @param runnable the Runnable to wrap
      * @return the wrapped TraceEventsRunnable
      */
@@ -199,7 +199,7 @@ public static DiagnosticSpan wrapWithDiagnosticSpan(Span span) {
      * @param span associated span
      * @param listenerMethod the listener method to be invoked
      */
-    void executeListeners(Span span, Consumer<TraceEventListener> listenerMethod) {
+    public void executeListeners(Span span, Consumer<TraceEventListener> listenerMethod) {
         if (span == null || traceEventListeners == null) {
             return;
         }