Refactor how span status is peristed and sent to sdk span from wrappers

embrace-android-sdk/src/main/java/io/embrace/android/embracesdk/internal/payload/SpanMapper.kt

@@ -7,6 +7,7 @@ import io.embrace.android.embracesdk.internal.clock.millisToNanos
 import io.embrace.android.embracesdk.internal.spans.EmbraceSpanData
 import io.embrace.android.embracesdk.internal.spans.hasFixedAttribute
 import io.embrace.android.embracesdk.internal.spans.setFixedAttribute
+import io.embrace.android.embracesdk.internal.spans.toStatus
 import io.embrace.android.embracesdk.spans.EmbraceSpanEvent
 import io.opentelemetry.api.trace.SpanId
 import io.opentelemetry.api.trace.StatusCode
@@ -18,12 +19,7 @@ internal fun EmbraceSpanData.toNewPayload() = Span(
     name = name,
     startTimeUnixNano = startTimeNanos,
     endTimeUnixNano = endTimeNanos,
-    status = when (status) {
-        StatusCode.UNSET -> Span.Status.UNSET
-        StatusCode.OK -> Span.Status.OK
-        StatusCode.ERROR -> Span.Status.ERROR
-        else -> Span.Status.UNSET
-    },
+    status = status.toStatus(),
     events = events.map(EmbraceSpanEvent::toNewPayload),
     attributes = attributes.toNewPayload()
 )

embrace-android-sdk/src/main/java/io/embrace/android/embracesdk/internal/spans/EmbraceExtensions.kt

@@ -6,7 +6,6 @@ import io.embrace.android.embracesdk.arch.schema.TelemetryType
 import io.embrace.android.embracesdk.opentelemetry.exceptionStacktrace
 import io.embrace.android.embracesdk.spans.EmbraceSpan
 import io.embrace.android.embracesdk.spans.EmbraceSpanEvent
-import io.embrace.android.embracesdk.spans.ErrorCode
 import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.common.AttributesBuilder
@@ -98,25 +97,6 @@ internal fun LogRecordBuilder.setFixedAttribute(fixedAttribute: FixedAttribute):
     return this
 }
 
-/**
- * Ends the given [Span], and setting the correct properties per the optional [ErrorCode] passed in. If [errorCode]
- * is not specified, it means the [Span] completed successfully, and no [ErrorCode] will be set.
- */
-internal fun Span.endSpan(errorCode: ErrorCode? = null, endTimeMs: Long? = null): Span {
-    if (errorCode != null) {
-        setStatus(StatusCode.ERROR)
-        setFixedAttribute(errorCode.fromErrorCode())
-    }
-
-    if (endTimeMs != null) {
-        end(endTimeMs, TimeUnit.MILLISECONDS)
-    } else {
-        end()
-    }
-
-    return this
-}
-
 /**
  * Returns the attributes as a new Map<String, String>
  */
@@ -195,3 +175,11 @@ internal fun io.embrace.android.embracesdk.Severity.toOtelSeverity(): Severity =
 internal fun String.isValidLongValueAttribute() = longValueAttributes.contains(this)
 
 internal val longValueAttributes = setOf(exceptionStacktrace.key)
+
+internal fun StatusCode.toStatus(): io.embrace.android.embracesdk.internal.payload.Span.Status {
+    return when (this) {
+        StatusCode.UNSET -> io.embrace.android.embracesdk.internal.payload.Span.Status.UNSET
+        StatusCode.OK -> io.embrace.android.embracesdk.internal.payload.Span.Status.OK
+        StatusCode.ERROR -> io.embrace.android.embracesdk.internal.payload.Span.Status.ERROR
+    }
+}

embrace-android-sdk/src/main/java/io/embrace/android/embracesdk/internal/spans/EmbraceSpanImpl.kt

@@ -2,6 +2,7 @@ package io.embrace.android.embracesdk.internal.spans
 
 import io.embrace.android.embracesdk.arch.schema.EmbType
 import io.embrace.android.embracesdk.arch.schema.EmbraceAttributeKey
+import io.embrace.android.embracesdk.arch.schema.ErrorCodeAttribute
 import io.embrace.android.embracesdk.arch.schema.FixedAttribute
 import io.embrace.android.embracesdk.internal.clock.millisToNanos
 import io.embrace.android.embracesdk.internal.clock.nanosToMillis
@@ -20,6 +21,7 @@ import io.embrace.android.embracesdk.spans.ErrorCode
 import io.embrace.android.embracesdk.spans.PersistableEmbraceSpan
 import io.opentelemetry.api.common.Attributes
 import io.opentelemetry.api.trace.SpanContext
+import io.opentelemetry.api.trace.StatusCode
 import io.opentelemetry.context.Context
 import io.opentelemetry.sdk.common.Clock
 import java.util.concurrent.ConcurrentHashMap
@@ -126,13 +128,18 @@ internal class EmbraceSpanImpl(
                         TimeUnit.NANOSECONDS
                     )
                 }
-                spanToStop.endSpan(errorCode, attemptedEndTimeMs)
+
+                if (errorCode != null) {
+                    setStatus(StatusCode.ERROR)
+                    spanToStop.setFixedAttribute(errorCode.fromErrorCode())
+                } else if (status == Span.Status.ERROR) {
+                    spanToStop.setFixedAttribute(ErrorCodeAttribute.Failure)
+                }
+
+                spanToStop.end(attemptedEndTimeMs, TimeUnit.MILLISECONDS)
                 successful = !isRecording
                 if (successful) {
                     spanId?.let { spanRepository.trackedSpanStopped(it) }
-                    if (errorCode != null) {
-                        status = Span.Status.ERROR
-                    }
                     spanEndTimeMs = attemptedEndTimeMs
                 }
             }
@@ -187,6 +194,15 @@ internal class EmbraceSpanImpl(
         return false
     }
 
+    override fun setStatus(statusCode: StatusCode, description: String) {
+        startedSpan.get()?.let { sdkSpan ->
+            synchronized(startedSpan) {
+                status = statusCode.toStatus()
+                sdkSpan.setStatus(statusCode, description)
+            }
+        }
+    }
+
     override fun addAttribute(key: String, value: String): Boolean {
         if (customAttributes.size < MAX_ATTRIBUTE_COUNT && attributeValid(key, value)) {
             synchronized(customAttributes) {

embrace-android-sdk/src/main/java/io/embrace/android/embracesdk/opentelemetry/EmbSpan.kt

@@ -1,7 +1,6 @@
 package io.embrace.android.embracesdk.opentelemetry
 
 import io.embrace.android.embracesdk.internal.spans.toStringMap
-import io.embrace.android.embracesdk.spans.ErrorCode
 import io.embrace.android.embracesdk.spans.PersistableEmbraceSpan
 import io.opentelemetry.api.common.AttributeKey
 import io.opentelemetry.api.common.Attributes
@@ -12,16 +11,12 @@ import io.opentelemetry.context.Context
 import io.opentelemetry.context.Scope
 import io.opentelemetry.sdk.common.Clock
 import java.util.concurrent.TimeUnit
-import java.util.concurrent.atomic.AtomicReference
 
 internal class EmbSpan(
     private val embraceSpan: PersistableEmbraceSpan,
     private val clock: Clock
 ) : Span {
 
-    private val pendingStatus: AtomicReference<StatusCode> = AtomicReference(StatusCode.UNSET)
-    private var pendingStatusDescription: String? = null
-
     override fun <T : Any> setAttribute(key: AttributeKey<T>, value: T): Span {
         embraceSpan.addAttribute(key = key.key, value = value.toString())
         return this
@@ -45,12 +40,7 @@ internal class EmbSpan(
 
     override fun setStatus(statusCode: StatusCode, description: String): Span {
         if (isRecording) {
-            synchronized(pendingStatus) {
-                if (isRecording) {
-                    pendingStatus.set(statusCode)
-                    pendingStatusDescription = description
-                }
-            }
+            embraceSpan.setStatus(statusCode, description)
         }
         return this
     }
@@ -69,19 +59,7 @@ internal class EmbSpan(
 
     override fun end(timestamp: Long, unit: TimeUnit) {
         if (isRecording) {
-            val endTimeMs = unit.toMillis(timestamp)
-            synchronized(pendingStatus) {
-                val finalStatus = pendingStatus.get()
-                setStatus(finalStatus)
-                when (finalStatus) {
-                    StatusCode.ERROR -> {
-                        embraceSpan.stop(errorCode = ErrorCode.FAILURE, endTimeMs = endTimeMs)
-                    }
-                    else -> {
-                        embraceSpan.stop(endTimeMs = endTimeMs)
-                    }
-                }
-            }
+            embraceSpan.stop(endTimeMs = unit.toMillis(timestamp))
         }
     }
 

embrace-android-sdk/src/main/java/io/embrace/android/embracesdk/spans/PersistableEmbraceSpan.kt

@@ -4,6 +4,7 @@ import io.embrace.android.embracesdk.arch.schema.EmbType
 import io.embrace.android.embracesdk.arch.schema.EmbraceAttributeKey
 import io.embrace.android.embracesdk.arch.schema.FixedAttribute
 import io.embrace.android.embracesdk.internal.payload.Span
+import io.opentelemetry.api.trace.StatusCode
 import io.opentelemetry.context.Context
 import io.opentelemetry.context.ContextKey
 import io.opentelemetry.context.ImplicitContextKeyed
@@ -49,6 +50,11 @@ internal interface PersistableEmbraceSpan : EmbraceSpan, ImplicitContextKeyed {
      */
     fun removeEvents(type: EmbType): Boolean
 
+    /**
+     * Set the [StatusCode] and status description of the wrapped Span
+     */
+    fun setStatus(statusCode: StatusCode, description: String = "")
+
     override fun storeInContext(context: Context): Context = context.with(embraceSpanContextKey, this)
 }
 

embrace-android-sdk/src/test/java/io/embrace/android/embracesdk/fakes/FakePersistableEmbraceSpan.kt

@@ -2,21 +2,23 @@ package io.embrace.android.embracesdk.fakes
 
 import io.embrace.android.embracesdk.arch.schema.EmbType
 import io.embrace.android.embracesdk.arch.schema.EmbraceAttributeKey
+import io.embrace.android.embracesdk.arch.schema.ErrorCodeAttribute
 import io.embrace.android.embracesdk.arch.schema.FixedAttribute
 import io.embrace.android.embracesdk.arch.schema.TelemetryType
 import io.embrace.android.embracesdk.internal.clock.millisToNanos
 import io.embrace.android.embracesdk.internal.clock.normalizeTimestampAsMillis
 import io.embrace.android.embracesdk.internal.payload.Span
 import io.embrace.android.embracesdk.internal.payload.toNewPayload
 import io.embrace.android.embracesdk.internal.spans.EmbraceSpanImpl.Companion.EXCEPTION_EVENT_NAME
-import io.embrace.android.embracesdk.internal.spans.endSpan
 import io.embrace.android.embracesdk.internal.spans.hasFixedAttribute
+import io.embrace.android.embracesdk.internal.spans.toStatus
 import io.embrace.android.embracesdk.spans.EmbraceSpan
 import io.embrace.android.embracesdk.spans.EmbraceSpanEvent
 import io.embrace.android.embracesdk.spans.ErrorCode
 import io.embrace.android.embracesdk.spans.PersistableEmbraceSpan
 import io.embrace.android.embracesdk.spans.getEmbraceSpan
 import io.opentelemetry.api.trace.SpanContext
+import io.opentelemetry.api.trace.StatusCode
 import io.opentelemetry.context.Context
 import java.util.concurrent.ConcurrentLinkedQueue
 import java.util.concurrent.TimeUnit
@@ -34,6 +36,7 @@ internal class FakePersistableEmbraceSpan(
     var spanStartTimeMs: Long? = null
     var spanEndTimeMs: Long? = null
     var status = Span.Status.UNSET
+    var statusDescription = ""
     var errorCode: ErrorCode? = null
     val attributes = mutableMapOf(type.toEmbraceKeyValuePair())
     val events = ConcurrentLinkedQueue<EmbraceSpanEvent>()
@@ -67,16 +70,21 @@ internal class FakePersistableEmbraceSpan(
         return true
     }
 
-    override fun stop(code: ErrorCode?, endTimeMs: Long?): Boolean {
+    override fun stop(errorCode: ErrorCode?, endTimeMs: Long?): Boolean {
         if (isRecording) {
-            checkNotNull(sdkSpan).endSpan(errorCode, endTimeMs)
-            errorCode = code
-            if (code != null) {
-                val error = code.fromErrorCode()
+            this.errorCode = errorCode
+            if (errorCode != null) {
+                setStatus(StatusCode.ERROR)
+            }
+
+            if (status == Span.Status.ERROR) {
+                val error = errorCode?.fromErrorCode() ?: ErrorCodeAttribute.Failure
                 setSystemAttribute(error.key, error.value)
-                status = Span.Status.ERROR
             }
-            spanEndTimeMs = endTimeMs ?: fakeClock.now()
+
+            val timestamp = endTimeMs ?: fakeClock.now()
+            checkNotNull(sdkSpan).end(timestamp, TimeUnit.MILLISECONDS)
+            spanEndTimeMs = timestamp
         }
         return true
     }
@@ -102,6 +110,11 @@ internal class FakePersistableEmbraceSpan(
         return true
     }
 
+    override fun setStatus(statusCode: StatusCode, description: String) {
+        status = statusCode.toStatus()
+        statusDescription = description
+    }
+
     override fun addAttribute(key: String, value: String): Boolean {
         attributes[key] = value
         return true

embrace-android-sdk/src/test/java/io/embrace/android/embracesdk/fakes/FakeSpan.kt

@@ -28,7 +28,8 @@ internal class FakeSpan(
     private var statusDescription: String = ""
 
     override fun <T : Any> setAttribute(key: AttributeKey<T>, value: T): Span {
-        TODO("Not yet implemented")
+        fakeSpanBuilder.setAttribute(key, value)
+        return this
     }
 
     override fun addEvent(name: String, attributes: Attributes): Span {

embrace-android-sdk/src/test/java/io/embrace/android/embracesdk/fakes/FakeSpanBuilder.kt

@@ -17,6 +17,7 @@ internal class FakeSpanBuilder(
     var spanKind: SpanKind? = null
     var parentContext: Context = Context.root()
     var startTimestampMs: Long? = null
+    var attributes = mutableMapOf<Any, Any>()
 
     override fun setParent(context: Context): SpanBuilder {
         parentContext = context
@@ -65,6 +66,7 @@ internal class FakeSpanBuilder(
     override fun startSpan(): FakeSpan = FakeSpan(this)
 
     override fun <T : Any> setAttribute(key: AttributeKey<T>, value: T): SpanBuilder {
-        TODO("Not yet implemented")
+        attributes[key] = value
+        return this
     }
 }