[EMBR-5543] Span auto termination (#132)

Sources/EmbraceCaptureService/CaptureService.swift

@@ -96,7 +96,7 @@ extension CaptureService {
             return nil
         }
 
-        return otel?.buildSpan(name: name, type: type, attributes: attributes)
+        return otel?.buildSpan(name: name, type: type, attributes: attributes, autoTerminationCode: nil)
     }
 
     /// Adds the given event to the session.

Sources/EmbraceCore/Capture/Network/URLSessionTaskHandler.swift

@@ -103,7 +103,8 @@ final class DefaultURLSessionTaskHandler: URLSessionTaskHandler {
             let networkSpan = otel.buildSpan(
                 name: name,
                 type: .networkRequest,
-                attributes: attributes
+                attributes: attributes,
+                autoTerminationCode: nil
             )
 
             // This should be modified if we start doing this for streaming tasks.

Sources/EmbraceCore/Capture/UX/View/Protocols/CaptureServices+UIViewController.swift

@@ -63,7 +63,12 @@ extension CaptureServices {
             throw parentSpanNotFoundError
         }
 
-        guard let builder = viewCaptureService.otel?.buildSpan(name: name, type: type, attributes: attributes) else {
+        guard let builder = viewCaptureService.otel?.buildSpan(
+            name: name, 
+            type: type,
+            attributes: attributes,
+            autoTerminationCode: nil
+        ) else {
             return nil
         }
 
@@ -88,7 +93,12 @@ extension CaptureServices {
             throw parentSpanNotFoundError
         }
 
-        guard let builder = viewCaptureService.otel?.buildSpan(name: name, type: type, attributes: attributes) else {
+        guard let builder = viewCaptureService.otel?.buildSpan(
+            name: name,
+            type: type,
+            attributes: attributes,
+            autoTerminationCode: nil
+        ) else {
             return
         }
 

Sources/EmbraceCore/Capture/UX/View/UIViewControllerHandler.swift

@@ -322,7 +322,9 @@ class UIViewControllerHandler {
             attributes: [
                 SpanSemantics.View.keyViewTitle: vc.emb_viewName,
                 SpanSemantics.View.keyViewName: vc.className
-            ])
+            ],
+            autoTerminationCode: nil
+        )
 
         if let parent = parent {
             builder.setParent(parent)

Sources/EmbraceCore/Public/Embrace+OTel.swift

@@ -30,14 +30,22 @@ extension Embrace: EmbraceOpenTelemetry {
     /// - Parameters:
     ///    - name: The name of the span.
     ///    - type: The type of the span. Will be set as the `emb.type` attribute.
+    ///    - autoTerminationCode: `SpanErrorCode` to be used to automatically close this span if the current session ends while the span is open.
     ///    - attributes: A dictionary of attributes to set on the span.
     /// - Returns: An OpenTelemetry `SpanBuilder`.
     public func buildSpan(
         name: String,
         type: SpanType = .performance,
-        attributes: [String: String] = [:]
+        attributes: [String: String] = [:],
+        autoTerminationCode: SpanErrorCode? = nil
     ) -> SpanBuilder {
-        otel.buildSpan(name: name, type: type, attributes: attributes)
+
+        if let autoTerminationCode = autoTerminationCode {
+            var attributes = attributes
+            attributes[SpanSemantics.keyAutoTerminationCode] = autoTerminationCode.rawValue
+        }
+
+        return otel.buildSpan(name: name, type: type, attributes: attributes)
     }
 
     /// Record a span after the fact
@@ -58,7 +66,7 @@ extension Embrace: EmbraceOpenTelemetry {
         endTime: Date,
         attributes: [String: String],
         events: [RecordingSpanEvent],
-        errorCode: ErrorCode?
+        errorCode: SpanErrorCode?
     ) {
         let builder = otel
             .buildSpan(name: name, type: type, attributes: attributes)
@@ -100,7 +108,7 @@ extension Embrace: EmbraceOpenTelemetry {
     /// - Parameter span: A `Span` object that implements `ReadableSpan`.
     public func flush(_ span: Span) {
         if let span = span as? ReadableSpan {
-            _ = exporter.export(spans: [span.toSpanData()])
+            EmbraceOTel.processor?.flush(span: span)
         } else {
             Embrace.logger.debug("Tried to flush a non-ReadableSpan object")
         }

Sources/EmbraceCore/Session/SessionController.swift

@@ -162,6 +162,9 @@ class SessionController: SessionControllable {
             }
             // -
 
+            // auto terminate spans
+            EmbraceOTel.processor?.autoTerminateSpans()
+
             // post notification
             notificationCenter.post(name: .embraceSessionWillEnd, object: currentSession)
 

Sources/EmbraceOTelInternal/EmbraceOTel.swift

@@ -14,6 +14,8 @@ public final class EmbraceOTel: NSObject {
     let instrumentationName = "EmbraceOpenTelemetry"
     let instrumentationVersion = "semver:\(EmbraceMeta.sdkVersion)"
 
+    public private(set) static var processor: SingleSpanProcessor?
+
     /// Setup the OpenTelemetryApi
     /// - Parameter: spanProcessor The processor in which to run during the lifetime of each Span
     public static func setup(spanProcessors: [SpanProcessor]) {
@@ -24,6 +26,8 @@ public final class EmbraceOTel: NSObject {
                 spanProcessors: spanProcessors
             )
         )
+
+        processor = spanProcessors.first(where: { $0 is SingleSpanProcessor }) as? SingleSpanProcessor
     }
 
     public static func setup(logSharedState: EmbraceLogSharedState) {

Sources/EmbraceOTelInternal/EmbraceOpenTelemetry.swift

@@ -5,11 +5,14 @@
 import Foundation
 import OpenTelemetryApi
 import EmbraceCommonInternal
+import EmbraceSemantics
 
 public protocol EmbraceOpenTelemetry: AnyObject {
     func buildSpan(name: String,
                    type: SpanType,
-                   attributes: [String: String]) -> SpanBuilder
+                   attributes: [String: String],
+                   autoTerminationCode: SpanErrorCode?
+    ) -> SpanBuilder
 
     func recordCompletedSpan(
         name: String,
@@ -19,7 +22,7 @@ public protocol EmbraceOpenTelemetry: AnyObject {
         endTime: Date,
         attributes: [String: String],
         events: [RecordingSpanEvent],
-        errorCode: ErrorCode?
+        errorCode: SpanErrorCode?
     )
 
     func add(events: [SpanEvent])

Sources/EmbraceOTelInternal/Trace/EmbraceSemantics/Span/Span+Embrace.swift

@@ -21,11 +21,11 @@ extension Span {
         setAttribute(key: SpanSemantics.keyIsPrivateSpan, value: "true")
     }
 
-    public func end(errorCode: ErrorCode? = nil, time: Date = Date()) {
+    public func end(errorCode: SpanErrorCode? = nil, time: Date = Date()) {
         end(error: nil, errorCode: errorCode, time: time)
     }
 
-    public func end(error: Error?, errorCode: ErrorCode? = nil, time: Date = Date()) {
+    public func end(error: Error?, errorCode: SpanErrorCode? = nil, time: Date = Date()) {
         var errorCode = errorCode
 
         // get attributes from error

Sources/EmbraceOTelInternal/Trace/EmbraceSemantics/Span/SpanData+Embrace.swift

@@ -20,11 +20,11 @@ extension SpanData {
         return .performance
     }
 
-    var errorCode: ErrorCode? {
+    var errorCode: SpanErrorCode? {
         guard let value = attributes[SpanSemantics.keyErrorCode] else {
             return nil
         }
-        return ErrorCode(rawValue: value.description)
+        return SpanErrorCode(rawValue: value.description)
     }
 
     public func toJSON() throws -> Data {

Sources/EmbraceOTelInternal/Trace/EmbraceSemantics/SpanBuilder/SpanBuilder+Embrace.swift

@@ -22,7 +22,7 @@ extension SpanBuilder {
         return self
     }
 
-    @discardableResult public func error(errorCode: ErrorCode) -> Self {
+    @discardableResult public func error(errorCode: SpanErrorCode) -> Self {
         setAttribute(key: SpanSemantics.keyErrorCode, value: errorCode.rawValue)
     }
 

Sources/EmbraceOTelInternal/Trace/Tracer/Span/Processor/SingleSpanProcessor.swift

@@ -5,21 +5,37 @@
 import Foundation
 import OpenTelemetryApi
 import OpenTelemetrySdk
+import EmbraceSemantics
+import EmbraceCommonInternal
 
 /// A really simple implementation of the SpanProcessor that converts the ExportableSpan to SpanData
 /// and passes it to the configured exporter in both `onStart` and `onEnd`
-public struct SingleSpanProcessor: SpanProcessor {
+public class SingleSpanProcessor: SpanProcessor {
 
     let spanExporter: SpanExporter
     private let processorQueue = DispatchQueue(label: "io.embrace.spanprocessor", qos: .utility)
 
+    @ThreadSafe var autoTerminationSpans: [SpanId: SpanAutoTerminationData] = [:]
+
     /// Returns a new SingleSpanProcessor that converts spans to SpanData and forwards them to
     /// the given spanExporter.
     /// - Parameter spanExporter: the SpanExporter to where the Spans are pushed.
     public init(spanExporter: SpanExporter) {
         self.spanExporter = spanExporter
     }
 
+    public func autoTerminateSpans() {
+        let now = Date()
+
+        for data in autoTerminationSpans.values {
+            data.span.setAttribute(key: SpanSemantics.keyErrorCode, value: data.code)
+            data.span.status = .error(description: data.code)
+            data.span.end(time: now)
+        }
+
+        autoTerminationSpans.removeAll()
+    }
+
     public let isStartRequired: Bool = true
 
     public let isEndRequired: Bool = true
@@ -29,14 +45,22 @@ public struct SingleSpanProcessor: SpanProcessor {
 
         let data = span.toSpanData()
 
+        // cache if flagged for auto termination
+        if let code = autoTerminationCode(for: data, parentId: data.parentSpanId) {
+            autoTerminationSpans[data.spanId] = SpanAutoTerminationData(
+                span: span,
+                spanData: data,
+                code: code, 
+                parentId: data.parentSpanId
+            )
+        }
+
         processorQueue.async {
             _ = exporter.export(spans: [data])
         }
     }
 
     public func onEnd(span: OpenTelemetrySdk.ReadableSpan) {
-        let exporter = self.spanExporter
-
         var data = span.toSpanData()
         if data.hasEnded && data.status == .unset {
             if let errorCode = data.errorCode {
@@ -47,7 +71,25 @@ public struct SingleSpanProcessor: SpanProcessor {
         }
 
         processorQueue.async {
-            _ = exporter.export(spans: [data])
+            _ = self.spanExporter.export(spans: [data])
+        }
+    }
+
+    public func flush(span: OpenTelemetrySdk.ReadableSpan) {
+        let data = span.toSpanData()
+
+        // update cache if needed
+        if let code = autoTerminationCode(for: data, parentId: data.parentSpanId) {
+            autoTerminationSpans[data.spanId] = SpanAutoTerminationData(
+                span: span,
+                spanData: data,
+                code: code,
+                parentId: data.parentSpanId
+            )
+        }
+
+        processorQueue.sync {
+            _ = self.spanExporter.export(spans: [data])
         }
     }
 
@@ -60,4 +102,33 @@ public struct SingleSpanProcessor: SpanProcessor {
             spanExporter.shutdown()
         }
     }
+
+    // finds the auto termination code from the span's attributes
+    // also tries to find it from it's parent spans
+    private func autoTerminationCode(for data: SpanData, parentId: SpanId? = nil) -> String? {
+        if let code = data.attributes[SpanSemantics.keyAutoTerminationCode] {
+            return code.description
+        }
+
+        if let parentId = parentId,
+           let parentData = autoTerminationSpans[parentId] {
+            return autoTerminationCode(for: parentData.spanData, parentId: parentData.parentId)
+        }
+
+        return nil
+    }
+}
+
+struct SpanAutoTerminationData {
+    let span: ReadableSpan
+    let spanData: SpanData
+    let code: String
+    let parentId: SpanId?
+
+    init(span: ReadableSpan, spanData: SpanData, code: String, parentId: SpanId? = nil) {
+        self.span = span
+        self.spanData = spanData
+        self.code = code
+        self.parentId = parentId
+    }
 }

Sources/EmbraceOTelInternal/Trace/EmbraceSemantics/Span/ErrorCode.swift → Sources/EmbraceSemantics/Span/SpanErrorCode.swift

@@ -3,7 +3,7 @@
 //
 
 /// Embrace specific error status for spans
-public enum ErrorCode: String {
+public enum SpanErrorCode: String {
     /// Span ended in an expected, but less than optimal state
     case failure
 

Sources/EmbraceSemantics/Span/SpanSemantics.swift

@@ -10,4 +10,6 @@ public struct SpanSemantics {
 
     public static let keyNSErrorMessage = "error.message"
     public static let keyNSErrorCode = "error.code"
+
+    public static let keyAutoTerminationCode = "emb.auto_termination.code"
 }