Fix zio (#9)

* fix print error

* fix zio operation name

* ignore main fiber

README.md

@@ -24,16 +24,17 @@ SkyWalking Extension Plugins for Scala 3
 > Other small versions of the library supported by this plugin may also work, but they have not been tested.
 
 ## Available Configurations
-| key                                             | description                                                                           |
-|-------------------------------------------------|---------------------------------------------------------------------------------------|
-| `plugin.ziov2.trace_fiber_fork`                 | Create spans for `ZIO.fork`, default is `false`.                                      |
-| `plugin.calibanv2.url_prefix`                   | Add a custom prefix to the graphql operation, default is `Caliban/GraphQL/`.          |
-| `plugin.calibanv2.ignore_url_prefixes`          | Ignore operation name that start with this prefix, i.e. no span will be created.      |
-| `plugin.calibanv2.collect_variables`            | Collect request variables.                                                            |
-| `plugin.calibanv2.variables_length_threshold`   | How many characters to keep and send to the OAP backend.                              |
-| `plugin.ziohttpv2.ignore_url_prefixes`          | Ignore request path that should start with this prefix, i.e. no span will be created. |
-| `plugin.ziohttpv2.collect_http_params`          | Collect http query params.                                                            |
-| `plugin.ziohttpv2.http_params_length_threshold` | How many characters to keep and send to the OAP backend.                              |
+| key                                             | description                                                                                                                                                             |
+|-------------------------------------------------|-------------------------------------------------------------------------------------------------------------------------------------------------------------------------|
+| `plugin.ziov2.trace_fiber_fork`                 | Create spans for `ZIO.fork`, default is `false`.                                                                                                                        |
+| `plugin.ziov2.ignore_fiber_regexes`             | Ignore ZIO Fibers that match this regex, i.e. no span will be created, default is `.*Application\.run.*,.*ZHttpServer\.start.*`. **Allow commas to separate multiple**. |
+| `plugin.calibanv2.url_prefix`                   | Add a custom prefix to the graphql operation, default is `Caliban/GraphQL/`.                                                                                            |
+| `plugin.calibanv2.ignore_url_prefixes`          | Ignore operation names starting with this prefix, i.e. no span will be created. **Allow commas to separate multiple**.                                                  |
+| `plugin.calibanv2.collect_variables`            | Collect request variables.                                                                                                                                              |
+| `plugin.calibanv2.variables_length_threshold`   | How many characters to keep and send to the OAP backend.                                                                                                                |
+| `plugin.ziohttpv2.ignore_url_prefixes`          | Ignore request paths starting with this prefix, i.e. no span will be created. **Allow commas to separate multiple**.                                                    |
+| `plugin.ziohttpv2.collect_http_params`          | Collect http query params.                                                                                                                                              |
+| `plugin.ziohttpv2.http_params_length_threshold` | How many characters to keep and send to the OAP backend.                                                                                                                |
 
 The prefix should be added when passing command line parameters, such as: `-Dskywalking.plugin.calibanv2.url_prefix=GQL/`
 

plugins/caliban-v2x-plugin/src/main/scala/org/bitlap/skywalking/apm/plugin/caliban/v2/TracingCaliban.scala

@@ -45,14 +45,7 @@ object TracingCaliban:
   ): ZIO[R, Nothing, GraphQLResponse[CalibanError]] =
     ret.onExit {
       case Exit.Success(value) =>
-        ZIO.attempt {
-          span.foreach(a => AgentUtils.stopAsync(a))
-          if value.errors.nonEmpty then {
-            val ex: Option[CalibanError] = value.errors.headOption
-            span.foreach(_.log(ex.getOrElse(CalibanError.ExecutionError("Effect failure"))))
-          }
-          ContextManager.stopSpan()
-        }.ignore
+        ZIO.attempt(tagCalibanError(span, value)).ignore
 
       case Exit.Failure(cause) =>
         ZIO.attempt {
@@ -82,16 +75,20 @@ object TracingCaliban:
     span: Option[AbstractSpan],
     result: GraphQLResponse[CalibanError]
   ): ZIO[Any, Nothing, Unit] =
-    ZIO.attempt {
-      span.foreach(a => AgentUtils.stopAsync(a))
+    ZIO.attempt(tagCalibanError(span, result)).ignore
 
-      if result.errors.nonEmpty then {
-        val ex: Option[CalibanError] = result.errors.headOption
-        span.foreach(_.log(ex.getOrElse(CalibanError.ExecutionError("Effect failure"))))
-      }
+  private def tagCalibanError(span: Option[AbstractSpan], result: GraphQLResponse[CalibanError]): Unit = {
+    span.foreach(a => AgentUtils.stopAsync(a))
+
+    if result.errors.nonEmpty then {
+      val ex: Option[CalibanError] = result.errors.headOption
+      span.foreach(_.log(ex.getOrElse(CalibanError.ExecutionError("Effect failure"))))
+    } else if result.data == null then {
+      span.foreach(_.log(CalibanError.ExecutionError("Effect failure")))
+    }
 
-      ContextManager.stopSpan()
-    }.ignore
+    ContextManager.stopSpan()
+  }
 
   def beforeOverallGraphQLRequest(graphQLRequest: GraphQLRequest): Option[AbstractSpan] =
     checkRequest(graphQLRequest) {

plugins/zio-v2x-plugin/src/main/scala/org/bitlap/skywalking/apm/plugin/zio/v2x/ZioPluginConfig.java

@@ -2,15 +2,16 @@
 
 import org.apache.skywalking.apm.agent.core.boot.PluginConfig;
 
-/** @author
- *    梦境迷离
- *  @version 1.0,2023/5/23
+/**
+ * @author 梦境迷离
+ * @version 1.0, 2023/5/23
  */
 public class ZioPluginConfig {
     public static class Plugin {
         @PluginConfig(root = ZioPluginConfig.class)
         public static class ZioV2 {
             public static boolean TRACE_FIBER_FORK = false;
+            public static String IGNORE_FIBER_REGEXES = ".*Application\\.run.*,.*ZHttpServer\\.start.*";
 
         }
     }

plugins/zio-v2x-plugin/src/main/scala/org/bitlap/skywalking/apm/plugin/zio/v2x/define/ZioFiberRuntimeInstrumentation.scala

@@ -58,8 +58,11 @@ object ZioFiberRuntimeInstrumentation:
 
   final val RUN_METHOD_INTERCEPTOR: String = classOf[ZioFiberRuntimeInterceptor].getTypeName
 
+  final val EXECUTOR_INTERCEPTOR: String = classOf[SetContextOnNewFiber].getTypeName
+
   final val methodInterceptors: Map[String, ElementMatcher[MethodDescription]] =
     Map(
+      EXECUTOR_INTERCEPTOR          -> named("drainQueueLaterOnExecutor").and(takesArguments(1)),
       RUN_METHOD_INTERCEPTOR + "_0" -> named("run").and(takesArguments(0)),
       RUN_METHOD_INTERCEPTOR + "_1" -> named("run").and(takesArguments(1))
     )

plugins/zio-v2x-plugin/src/main/scala/org/bitlap/skywalking/apm/plugin/zio/v2x/interceptor/ZioFiberRuntimeInterceptor.scala

@@ -14,25 +14,38 @@ import org.apache.skywalking.apm.agent.core.plugin.interceptor.enhance.*
 import org.apache.skywalking.apm.network.trace.component.ComponentsDefine
 import org.bitlap.skywalking.apm.plugin.common.*
 import org.bitlap.skywalking.apm.plugin.zcommon.*
+import org.bitlap.skywalking.apm.plugin.zio.v2x.ZioPluginConfig
 
 /** @author
  *    梦境迷离
  *  @version 1.0,2023/5/16
  */
 final class ZioFiberRuntimeInterceptor extends InstanceMethodsAroundInterceptor:
 
+  private val SpanSwitch = "spanSwitch"
+
   override def beforeMethod(
     objInst: EnhancedInstance,
     method: Method,
     allArguments: Array[Object],
     argumentsTypes: Array[Class[?]],
     result: MethodInterceptResult
   ): Unit =
-    val fiberRuntime = objInst.asInstanceOf[FiberRuntime[?, ?]]
-    val currentSpan  = ContextManager.createLocalSpan(AgentUtils.generateFiberOperationName("ZIO"))
-    currentSpan.setComponent(ComponentsDefine.JDK_THREADING)
-    TagUtils.setZioTags(currentSpan, fiberRuntime.id, objInst)
-    AgentUtils.continuedSnapshot(objInst)
+    val fiberRuntime      = objInst.asInstanceOf[FiberRuntime[?, ?]]
+    val location          = fiberRuntime.location.toString
+    val mainMethodRegexes = ZioPluginConfig.Plugin.ZioV2.IGNORE_FIBER_REGEXES.split(",").toList.filter(_.nonEmpty)
+    val matchRegex        = mainMethodRegexes.map(_.r).exists(_.matches(location))
+
+    ContextManager.getRuntimeContext.put(SpanSwitch, !matchRegex)
+
+    if location != null && location != "" && !matchRegex then {
+      val currentSpan = ContextManager.createLocalSpan(
+        AgentUtils.generateFiberOperationName("ZIO", Some(fiberRuntime.location.toString))
+      )
+      currentSpan.setComponent(ComponentsDefine.JDK_THREADING)
+      TagUtils.setZioTags(currentSpan, fiberRuntime.id, objInst)
+      AgentUtils.continuedSnapshot(objInst)
+    }
 
   override def afterMethod(
     objInst: EnhancedInstance,
@@ -41,7 +54,11 @@ final class ZioFiberRuntimeInterceptor extends InstanceMethodsAroundInterceptor:
     argumentsTypes: Array[Class[?]],
     ret: Object
   ): Object =
-    AgentUtils.stopIfActive()
+    val switch = ContextManager.getRuntimeContext.get(SpanSwitch, classOf[Boolean])
+    if switch then {
+      AgentUtils.stopIfActive()
+    }
+
     ret
 
   override def handleMethodException(

plugins/zio-v2x-plugin/src/main/scala/org/bitlap/skywalking/apm/plugin/zio/v2x/interceptor/ZioUnsafeForkInterceptor.scala

@@ -35,7 +35,9 @@ final class ZioUnsafeForkInterceptor extends InstanceMethodsAroundInterceptor:
     ContextManager.getRuntimeContext.put(ForkSwitch, switch)
     if switch then {
       val fiberRuntime = allArguments(2).asInstanceOf[FiberRuntime[?, ?]]
-      val currentSpan  = ContextManager.createLocalSpan(AgentUtils.generateFiberForkOperationName("ZIO"))
+      val currentSpan = ContextManager.createLocalSpan(
+        AgentUtils.generateFiberForkOperationName("ZIO", Some(fiberRuntime.id.location.toString))
+      )
       currentSpan.setComponent(ComponentsDefine.JDK_THREADING)
       TagUtils.setZioTags(currentSpan, fiberRuntime.id, objInst)
       AgentUtils.continuedSnapshot(fiberRuntime)

shared/plugin-common/src/main/scala/org/bitlap/skywalking/apm/plugin/common/AgentUtils.scala

@@ -81,6 +81,12 @@ object AgentUtils:
     LOGGER.error(s"Span Operation Error!", e)
     if ContextManager.isActive then ContextManager.activeSpan.log(e)
 
-  def generateFiberOperationName(tpe: String) = s"$tpe/IOFiber/${Thread.currentThread.getName}"
+  def generateFiberOperationName(tpe: String, id: Option[String] = None) =
+    id.fold(
+      s"$tpe/Fiber/${Thread.currentThread().getName}"
+    )(i => s"$tpe/Fiber/$i")
 
-  def generateFiberForkOperationName(tpe: String) = s"$tpe/ForkFiber/${Thread.currentThread.getName}"
+  def generateFiberForkOperationName(tpe: String, id: Option[String] = None) =
+    id.fold(
+      s"$tpe/Fiber/${Thread.currentThread().getName}/fork"
+    )(i => s"$tpe/Fiber/$i/fork")

shared/plugin-common/src/main/scala/org/bitlap/skywalking/apm/plugin/common/CustomTag.scala

@@ -12,3 +12,4 @@ enum CustomTag[T](val tag: StringTag):
   case FiberLocation    extends CustomTag(new StringTag(102, "zio.fiber.location"))
   case CalibanVariables extends CustomTag(new StringTag(103, "caliban.variables"))
   case FiberClassName   extends CustomTag(new StringTag(104, "zio.fiber.className"))
+  case CurrentThread    extends CustomTag(new StringTag(105, "zio.fiber.currentThread"))

shared/zio-plugin-common/src/main/scala/org/bitlap/skywalking/apm/plugin/zcommon/TagUtils.scala

@@ -31,6 +31,7 @@ object TagUtils:
       CustomTag.FiberStartTime.tag.set(span, formatDate(fiberId.startTimeMillis))
       CustomTag.FiberLocation.tag.set(span, fiberId.location.toString)
       CustomTag.FiberClassName.tag.set(span, objInst.getClass.getName)
+      CustomTag.CurrentThread.tag.set(span, Thread.currentThread().getName)
     } match
       case Failure(ex) => span.errorOccurred.log(ex)
       case Success(_)  =>

version.sbt

@@ -1 +1 @@
-ThisBuild / version := "0.1.9-SNAPSHOT"
+ThisBuild / version := "0.1.8"