Return span handler (#261)

src/main/kotlin/io/emeraldpay/dshackle/config/spans/ProviderSpanHandler.kt

@@ -3,48 +3,82 @@ package io.emeraldpay.dshackle.config.spans
 import brave.handler.MutableSpan
 import brave.handler.SpanHandler
 import brave.propagation.TraceContext
-import org.springframework.beans.factory.annotation.Value
+import com.github.benmanes.caffeine.cache.Caffeine
+import org.slf4j.LoggerFactory
 import org.springframework.boot.autoconfigure.condition.ConditionalOnProperty
 import org.springframework.stereotype.Component
-import java.util.concurrent.TimeUnit
+import java.time.Duration
 
 @Component
 @ConditionalOnProperty(value = ["spring.zipkin.enabled"], havingValue = "true")
 class ProviderSpanHandler(
-    @Value("\${spans.collect.collect-only-errors:false}")
+    private val spanExportableList: List<SpanExportable>,
-    private val onlyErrors: Boolean? = false,
+    private val zipkinSpanHandler: SpanHandler,
-    @Value("\${spans.collect.long-span-threshold}")
-    private val longSpanThreshold: Long? = null
 ) : SpanHandler() {
+    private val log = LoggerFactory.getLogger(ProviderSpanHandler::class.java)
 
-    override fun end(context: TraceContext?, span: MutableSpan, cause: Cause?): Boolean {
+    private val spans = Caffeine
-        val duration = TimeUnit.MILLISECONDS.convert(span.finishTimestamp() - span.startTimestamp(), TimeUnit.MICROSECONDS)
+        .newBuilder()
+        .expireAfterWrite(Duration.ofMinutes(2))
+        .build<String, MutableList<MutableSpan>>()
 
-        val isError = span.tags().containsKey("error")
+    override fun end(context: TraceContext, span: MutableSpan, cause: Cause): Boolean {
-
+        if (span.traceId().length > 20) {
-        // If onlyErrors is true and there is an error, return true.
+            val key = span.parentId() ?: span.id()
-        if (onlyErrors == true && isError) {
+            val spanList = spans.asMap().computeIfAbsent(key) { mutableListOf() }
-            return true
+            spanList.add(span)
         }
-
-        // If onlyErrors is true, there is no error, and the span is long, return true.
-        if (onlyErrors == true && !isError && longSpanThreshold != null && duration >= longSpanThreshold) {
-            return true
-        }
-
-        // If onlyErrors is false, check for the time threshold condition.
-        if (onlyErrors == false) {
-            // If longSpanThreshold is null, return true.
-            if (longSpanThreshold == null) {
-                return true
-            }
-            // If longSpanThreshold is set, only return true if the duration is >= time threshold.
-            else if (duration >= longSpanThreshold) {
-                return true
-            }
-        }
-
-        // Return false if none of the above conditions are met.
         return false
     }
+
+    fun sendSpans(traceContext: TraceContext) {
+        try {
+            sendSpansInternal(traceContext)
+        } catch (e: Exception) {
+            log.warn("Error while handling and sending spans - ${e.message}")
+        }
+    }
+
+    private fun sendSpansInternal(traceContext: TraceContext) {
+        val spansInfo = SpansInfo()
+
+        processSpans(traceContext.spanIdString(), spansInfo)
+
+        spansInfo.spans
+            .map { it.parentId() }
+            .forEach {
+                if (it != null) {
+                    spans.invalidate(it)
+                }
+            }
+
+        if (spansInfo.exportable) {
+            spansInfo.spans.forEach {
+                zipkinSpanHandler.end(traceContext, it, Cause.FINISHED)
+            }
+        }
+    }
+
+    private fun processSpans(spanId: String, spansInfo: SpansInfo) {
+        val currentSpans: List<MutableSpan>? = spans.getIfPresent(spanId)
+
+        currentSpans?.forEach {
+            processSpanInfo(it, spansInfo)
+            if (spanId != it.id()) {
+                processSpans(it.id(), spansInfo)
+            }
+        }
+    }
+
+    private fun processSpanInfo(span: MutableSpan, spansInfo: SpansInfo) {
+        spansInfo.spans.add(span)
+        if (spanExportableList.any { it.isExportable(span) }) {
+            spansInfo.exportable = true
+        }
+    }
+
+    private data class SpansInfo(
+        var exportable: Boolean = false,
+        val spans: MutableList<MutableSpan> = mutableListOf()
+    )
 }

src/main/kotlin/io/emeraldpay/dshackle/config/spans/export.kt

@@ -0,0 +1,36 @@
+package io.emeraldpay.dshackle.config.spans
+
+import brave.handler.MutableSpan
+import io.emeraldpay.dshackle.commons.SPAN_ERROR
+import io.emeraldpay.dshackle.commons.SPAN_NO_RESPONSE_MESSAGE
+import org.springframework.beans.factory.annotation.Value
+import org.springframework.stereotype.Component
+import java.util.concurrent.TimeUnit.MICROSECONDS
+import java.util.concurrent.TimeUnit.MILLISECONDS
+
+interface SpanExportable {
+    fun isExportable(span: MutableSpan): Boolean
+}
+
+@Component
+class ErrorSpanExportable : SpanExportable {
+    override fun isExportable(span: MutableSpan): Boolean = span.tags().containsKey(SPAN_ERROR)
+}
+
+@Component
+class NoResponseSpanExportable : SpanExportable {
+    override fun isExportable(span: MutableSpan): Boolean = span.tags().containsKey(SPAN_NO_RESPONSE_MESSAGE)
+}
+
+@Component
+class LongResponseSpanExportable(
+    @Value("\${spans.collect.long-span-threshold}")
+    private val longSpanThreshold: Long? = null
+) : SpanExportable {
+
+    override fun isExportable(span: MutableSpan): Boolean {
+        return MILLISECONDS.convert(
+            span.finishTimestamp() - span.startTimestamp(), MICROSECONDS
+        ) >= longSpanThreshold!!
+    }
+}

src/main/kotlin/io/emeraldpay/dshackle/rpc/BlockchainRpc.kt

@@ -16,17 +16,20 @@
  */
 package io.emeraldpay.dshackle.rpc
 
+import brave.Tracer
 import io.emeraldpay.api.proto.BlockchainOuterClass
 import io.emeraldpay.api.proto.Common
 import io.emeraldpay.api.proto.ReactorBlockchainGrpc
 import io.emeraldpay.dshackle.Chain
 import io.emeraldpay.dshackle.ChainValue
 import io.emeraldpay.dshackle.SilentException
+import io.emeraldpay.dshackle.config.spans.ProviderSpanHandler
 import io.micrometer.core.instrument.Counter
 import io.micrometer.core.instrument.Metrics
 import io.micrometer.core.instrument.Timer
 import org.apache.commons.lang3.RandomStringUtils
 import org.slf4j.LoggerFactory
+import org.springframework.beans.factory.annotation.Autowired
 import org.springframework.beans.factory.annotation.Qualifier
 import org.springframework.context.annotation.DependsOn
 import org.springframework.stereotype.Service
@@ -50,7 +53,10 @@ class BlockchainRpc(
     private val estimateFee: EstimateFee,
     private val subscribeNodeStatus: SubscribeNodeStatus,
     @Qualifier("rpcScheduler")
-    private val scheduler: Scheduler
+    private val scheduler: Scheduler,
+    @Autowired(required = false)
+    private val providerSpanHandler: ProviderSpanHandler?,
+    private val tracer: Tracer,
 ) : ReactorBlockchainGrpc.BlockchainImplBase() {
 
     private val log = LoggerFactory.getLogger(BlockchainRpc::class.java)
@@ -96,7 +102,11 @@ class BlockchainRpc(
                     itemMetrics.nativeItemResponseErr.increment()
                 }
             }
-        }.doOnError { failMetric.increment() }
+        }.doOnError {
+            failMetric.increment()
+        }.doFinally {
+            providerSpanHandler?.sendSpans(tracer.currentSpan().context())
+        }
     }
 
     override fun nativeSubscribe(request: Mono<BlockchainOuterClass.NativeSubscribeRequest>): Flux<BlockchainOuterClass.NativeSubscribeReplyItem> {

src/main/resources/application.yml

@@ -20,5 +20,4 @@ spring:
 
 spans:
   collect:
-    long-span-threshold: ${LONG_SPAN_THRESHOLD:1000}
+    long-span-threshold: ${LONG_SPAN_THRESHOLD:1000}
-    collect-only-errors: ${ONLY_ERROR_SPANS:true}

src/test/kotlin/io/emeraldpay/dshackle/config/spans/ProviderSpanHandlerTest.kt

@@ -1,85 +0,0 @@
-package io.emeraldpay.dshackle.config.spans
-
-import brave.handler.MutableSpan
-import org.junit.jupiter.api.Assertions.assertEquals
-import org.junit.jupiter.params.ParameterizedTest
-import org.junit.jupiter.params.provider.Arguments
-import org.junit.jupiter.params.provider.MethodSource
-import java.util.stream.Stream
-
-class ProviderSpanHandlerTest {
-
-    companion object {
-        @JvmStatic
-        fun data(): Stream<Arguments> {
-            return Stream.of(
-                Arguments.of(
-                    false, null,
-                    MutableSpan().apply {
-                        startTimestamp(0)
-                        finishTimestamp(1000000)
-                    },
-                    true
-                ),
-                Arguments.of(
-                    false, 1000L,
-                    MutableSpan().apply {
-                        startTimestamp(0)
-                        finishTimestamp(999999)
-                    },
-                    false
-                ),
-                Arguments.of(
-                    false, 1000L,
-                    MutableSpan().apply {
-                        startTimestamp(0)
-                        finishTimestamp(1000000)
-                    },
-                    true
-                ),
-                Arguments.of(
-                    true, null,
-                    MutableSpan().apply {
-                        startTimestamp(0)
-                        finishTimestamp(1000000)
-                    },
-                    false
-                ),
-                Arguments.of(
-                    true, null,
-                    MutableSpan().apply {
-                        startTimestamp(0)
-                        finishTimestamp(1000000)
-                        tag("error", "true")
-                    },
-                    true
-                ),
-                Arguments.of(
-                    false, null,
-                    MutableSpan().apply {
-                        startTimestamp(0)
-                        finishTimestamp(1000000)
-                        tag("error", "true")
-                    },
-                    true
-                ),
-                Arguments.of(
-                    true, 1000L,
-                    MutableSpan().apply {
-                        startTimestamp(0)
-                        finishTimestamp(1000000)
-                    },
-                    true
-                ),
-            )
-        }
-    }
-
-    @ParameterizedTest
-    @MethodSource("data")
-    fun end(onlyErrors: Boolean?, timeThreshold: Long?, span: MutableSpan, expected: Boolean) {
-        val handler = ProviderSpanHandler(onlyErrors, timeThreshold)
-        val res = handler.end(null, span, null)
-        assertEquals(expected, res)
-    }
-}