Implementing v2/spans endpoint (#241)

embrace-io · May 9, 2024 · 646b536 · 646b536
1 parent 8c9a258
commit 646b536
Show file tree

Hide file tree

Showing 34 changed files with 510 additions and 605 deletions.
diff --git a/Sources/EmbraceCore/Internal/Embrace+Setup.swift b/Sources/EmbraceCore/Internal/Embrace+Setup.swift
@@ -30,14 +30,14 @@ extension Embrace {
         // endpoints
         let baseUrl = EMBDevice.isDebuggerAttached ?
             options.endpoints.developmentBaseURL : options.endpoints.baseURL
-        guard let sessionsURL = URL.sessionsEndpoint(basePath: baseUrl),
+        guard let spansURL = URL.spansEndpoint(basePath: baseUrl),
               let blobsURL = URL.blobsEndpoint(basePath: baseUrl),
               let logsURL = URL.logsEndpoint(basePath: baseUrl) else {
             ConsoleLog.error("Failed to initialize endpoints!")
             return nil
         }
 
-        let endpoints = EmbraceUpload.EndpointOptions(sessionsURL: sessionsURL, blobsURL: blobsURL, logsURL: logsURL)
+        let endpoints = EmbraceUpload.EndpointOptions(spansURL: spansURL, blobsURL: blobsURL, logsURL: logsURL)
 
         // cache
         guard let cacheUrl = EmbraceFileSystem.uploadsDirectoryPath(

diff --git a/Sources/EmbraceCore/Payload/Builders/SessionPayloadBuilder.swift b/Sources/EmbraceCore/Payload/Builders/SessionPayloadBuilder.swift
@@ -10,7 +10,7 @@ class SessionPayloadBuilder {
 
     static var resourceName = "emb.session.upload_index"
 
-    class func build(for sessionRecord: SessionRecord, storage: EmbraceStorage) -> SessionPayload {
+    class func build(for sessionRecord: SessionRecord, storage: EmbraceStorage) -> PayloadEnvelope<[SpanPayload]> {
         var resource: MetadataRecord?
 
         do {
@@ -42,15 +42,36 @@ class SessionPayloadBuilder {
         }
 
         // build spans
-        let (spans, spanSnapshots) = SpansPayloadBuilder.build(for: sessionRecord, storage: storage)
+        let (spans, spanSnapshots) = SpansPayloadBuilder.build(
+            for: sessionRecord,
+            storage: storage,
+            sessionNumber: counter
+        )
+
+        // build resources payload
+        var resources: [MetadataRecord] = []
+        do {
+            resources = try storage.fetchResourcesForSessionId(sessionRecord.id)
+        } catch {
+            ConsoleLog.error("Error fetching resources for session \(sessionRecord.id.toString)")
+        }
+        let resourcePayload =  ResourcePayload(from: resources)
+
+        // build metadata payload
+        var metadata: [MetadataRecord] = []
+        do {
+            metadata = try storage.fetchCustomPropertiesForSessionId(sessionRecord.id)
+        } catch {
+            ConsoleLog.error("Error fetching custom properties for session \(sessionRecord.id.toString)")
+        }
+        let metadataPayload =  MetadataPayload(from: metadata)
 
         // build payload
-        return SessionPayload(
-            from: sessionRecord,
-            resourceFetcher: storage,
+        return PayloadEnvelope(
             spans: spans,
             spanSnapshots: spanSnapshots,
-            counter: counter
+            resource: resourcePayload,
+            metadata: metadataPayload
         )
     }
 }
diff --git a/Sources/EmbraceCore/Payload/Builders/SpansPayloadBuilder.swift b/Sources/EmbraceCore/Payload/Builders/SpansPayloadBuilder.swift
@@ -13,7 +13,8 @@ class SpansPayloadBuilder {
 
     class func build(
         for sessionRecord: SessionRecord,
-        storage: EmbraceStorage
+        storage: EmbraceStorage,
+        sessionNumber: Int = -1
     ) -> (spans: [SpanPayload], spanSnapshots: [SpanPayload]) {
 
         let endTime = sessionRecord.endTime ?? sessionRecord.lastHeartbeatTime
@@ -39,7 +40,11 @@ class SpansPayloadBuilder {
         var spanSnapshots: [SpanPayload] = []
 
         // fetch and add session span first
-        if let sessionSpanPayload = buildSessionSpanPayload(for: sessionRecord, storage: storage) {
+        if let sessionSpanPayload = buildSessionSpanPayload(
+            for: sessionRecord,
+            storage: storage,
+            sessionNumber: sessionNumber
+        ) {
             spans.append(sessionSpanPayload)
         }
 
@@ -69,19 +74,21 @@ class SpansPayloadBuilder {
         return (spans, spanSnapshots)
     }
 
-    class func buildSessionSpanPayload(for sessionRecord: SessionRecord, storage: EmbraceStorage) -> SpanPayload? {
+    class func buildSessionSpanPayload(
+        for sessionRecord: SessionRecord,
+        storage: EmbraceStorage,
+        sessionNumber: Int
+    ) -> SpanPayload? {
         do {
+            var spanData: SpanData?
             let sessionSpan = try storage.fetchSpan(id: sessionRecord.spanId, traceId: sessionRecord.traceId)
+
             if let rawData = sessionSpan?.data {
-                let sessionSpanData = try JSONDecoder().decode(SpanData.self, from: rawData)
-                return SpanPayload(
-                    from: sessionSpanData,
-                    endTime: sessionRecord.endTime ?? sessionRecord.lastHeartbeatTime
-                )
-            } else {
-                return SessionSpanUtils.payload(from: sessionRecord)
+                spanData = try JSONDecoder().decode(SpanData.self, from: rawData)
             }
 
+            return SessionSpanUtils.payload(from: sessionRecord, spanData: spanData, sessionNumber: sessionNumber)
+
         } catch {
             ConsoleLog.warning("Error fetching span for session \(sessionRecord.id):\n\(error.localizedDescription)")
         }

diff --git a/Sources/EmbraceCore/Payload/PayloadEnvelope.swift b/Sources/EmbraceCore/Payload/PayloadEnvelope.swift
@@ -4,7 +4,7 @@
 
 import Foundation
 
-struct PayloadEnvelope<T: Codable>: Codable {
+struct PayloadEnvelope<T: Encodable>: Encodable {
     var resource: ResourcePayload
     var metadata: MetadataPayload
     var version: String = "1.0" // TODO: Make this the actual version
@@ -20,3 +20,13 @@ extension PayloadEnvelope<[LogPayload]> {
         self.metadata = metadata
     }
 }
+
+extension PayloadEnvelope<[SpanPayload]> {
+    init(spans: [SpanPayload], spanSnapshots: [SpanPayload], resource: ResourcePayload, metadata: MetadataPayload) {
+        type = "spans"
+        self.data["spans"] = spans
+        self.data["span_snapshots"] = spanSnapshots
+        self.resource = resource
+        self.metadata = metadata
+    }
+}
diff --git a/Sources/EmbraceCore/Payload/SessionInfoPayload.swift b/Sources/EmbraceCore/Payload/SessionInfoPayload.swift
diff --git a/Sources/EmbraceCore/Payload/SessionPayload.swift b/Sources/EmbraceCore/Payload/SessionPayload.swift
diff --git a/Sources/EmbraceCore/Payload/Spans/SpanEventPayload.swift b/Sources/EmbraceCore/Payload/Spans/SpanEventPayload.swift
@@ -8,7 +8,7 @@ import EmbraceOTel
 struct SpanEventPayload: Encodable {
     let name: String
     let timestamp: Int
-    let attributes: [String: Any]
+    let attributes: [Attribute]
 
     enum CodingKeys: String, CodingKey {
         case name

diff --git a/Sources/EmbraceCore/Payload/Spans/SpanLinkPayload.swift b/Sources/EmbraceCore/Payload/Spans/SpanLinkPayload.swift
@@ -7,7 +7,7 @@ import EmbraceOTel
 struct SpanLinkPayload: Encodable {
     let traceId: String
     let spanId: String
-    let attributes: [String: Any]
+    let attributes: [Attribute]
 
     enum CodingKeys: String, CodingKey {
         case traceId = "trace_id"

diff --git a/Sources/EmbraceCore/Payload/Spans/SpanPayload.swift b/Sources/EmbraceCore/Payload/Spans/SpanPayload.swift
@@ -13,7 +13,7 @@ struct SpanPayload: Encodable {
     let status: String
     let startTime: Int
     let endTime: Int?
-    let attributes: [String: Any]
+    let attributes: [Attribute]
     let events: [SpanEventPayload]
     let links: [SpanLinkPayload]
 
@@ -48,11 +48,11 @@ struct SpanPayload: Encodable {
             self.endTime = nil
         }
 
-        var dict = PayloadUtils.convertSpanAttributes(span.attributes)
+        var attributeArray = PayloadUtils.convertSpanAttributes(span.attributes)
         if failed {
-            dict["emb.error_code"] = "failure"
+            attributeArray.append(Attribute(key: "emb.error_code", value: "failure"))
         }
-        self.attributes = dict
+        self.attributes = attributeArray
     }
 
     func encode(to encoder: Encoder) throws {

diff --git a/Sources/EmbraceCore/Payload/Utils/PayloadUtils.swift b/Sources/EmbraceCore/Payload/Utils/PayloadUtils.swift
@@ -43,16 +43,15 @@ class PayloadUtils {
         return []
     }
 
-    static func convertSpanAttributes(_ attributes: [String: AttributeValue]) -> [String: Any] {
-        var result: [String: Any] = [:]
+    static func convertSpanAttributes(_ attributes: [String: AttributeValue]) -> [Attribute] {
+        var result: [Attribute] = []
 
         for (key, value) in attributes {
             switch value {
-            case .bool(let boolValue): result[key] = boolValue
-            case .double(let doubleValue): result[key] = doubleValue
-            case .int(let intValue): result[key] = intValue
-            case .string(let stringValue): result[key] = stringValue
-            default: continue
+            case .boolArray, .intArray, .doubleArray, .stringArray:
+                continue
+            default:
+                result.append(Attribute(key: key, value: value.description))
             }
         }
 

diff --git a/Sources/EmbraceCore/Session/DataRecovery/UnsentDataHandler.swift b/Sources/EmbraceCore/Session/DataRecovery/UnsentDataHandler.swift
@@ -196,8 +196,8 @@ class UnsentDataHandler {
             return
         }
 
-        // upload session
-        upload.uploadSession(id: session.id.toString, data: payloadData) { result in
+        // upload session spans
+        upload.uploadSpans(id: session.id.toString, data: payloadData) { result in
             switch result {
             case .success:
                 do {

diff --git a/Sources/EmbraceCore/Session/SessionController.swift b/Sources/EmbraceCore/Session/SessionController.swift
@@ -50,7 +50,9 @@ class SessionController: SessionControllable {
         self.heartbeat = SessionHeartbeat(queue: heartbeatQueue, interval: heartbeatInterval)
 
         self.heartbeat.callback = { [weak self] in
-            self?.currentSession?.lastHeartbeatTime = Date()
+            let heartbeat = Date()
+            self?.currentSession?.lastHeartbeatTime = heartbeat
+            SessionSpanUtils.setHeartbeat(span: self?.currentSessionSpan, heartbeat: heartbeat)
             self?.save()
         }
     }
@@ -75,9 +77,12 @@ class SessionController: SessionControllable {
 
         return lock.locked {
 
+            // detect cold start
+            let isColdStart = withinColdStartInterval(startTime: startTime)
+
             // create session span
             let newId = SessionIdentifier.random
-            let span = SessionSpanUtils.span(id: newId, startTime: startTime)
+            let span = SessionSpanUtils.span(id: newId, startTime: startTime, state: state, coldStart: isColdStart)
             currentSessionSpan = span
 
             // create session record
@@ -89,7 +94,7 @@ class SessionController: SessionControllable {
                 spanId: span.context.spanId.hexString,
                 startTime: startTime
             )
-            session.coldStart = withinColdStartInterval(startTime: startTime)
+            session.coldStart = isColdStart
             currentSession = session
 
             // save session record
@@ -119,6 +124,8 @@ class SessionController: SessionControllable {
 
             let now = Date()
             currentSessionSpan?.end(time: now)
+            SessionSpanUtils.setCleanExit(span: currentSessionSpan, cleanExit: true)
+
             currentSession?.endTime = now
             currentSession?.cleanExit = true
 
@@ -136,11 +143,13 @@ class SessionController: SessionControllable {
     }
 
     func update(state: SessionState) {
+        SessionSpanUtils.setState(span: currentSessionSpan, state: state)
         currentSession?.state = state.rawValue
         save()
     }
 
     func update(appTerminated: Bool) {
+        SessionSpanUtils.setTerminated(span: currentSessionSpan, terminated: appTerminated)
         currentSession?.appTerminated = appTerminated
         save()
     }