Fixes logic in EmbraceStorage+Options such that limit is applied to S…

…panType (#232)

Was originally SpanType per Trace, which will realistically never hit

Also introduces a default max limit if none is configured in the EmbraceStorage.Options
Adds tests to cover limit pruning logic

Sources/EmbraceCommon/SpanType/SpanType.swift

@@ -6,7 +6,7 @@
 /// - This struct will be serialized into the Span's `emb.type` attribute.
 /// - This struct is encoded as a String with the format `<primary>.<secondary>`.
 /// - The primary category is required, but the secondary category is optional.
-public struct SpanType: Equatable {
+public struct SpanType: Equatable, Hashable {
     let primary: Primary
 
     let secondary: String?

Sources/EmbraceStorage/EmbraceStorage+Options.swift

@@ -3,6 +3,7 @@
 //
 
 import Foundation
+import EmbraceCommon
 
 public extension EmbraceStorage {
     enum StorageMechanism {
@@ -15,7 +16,7 @@ public extension EmbraceStorage {
         let storageMechanism: StorageMechanism
 
         /// Dictionary containing the storage limits per span type
-        public var spanLimits: [String: Int] = [:]
+        public var spanLimits: [SpanType: Int] = [:]
 
         /// Determines how many `MetadataRecords` of the `.resource` type can be present at any given time.
         public var resourcesLimit: Int = 100

Sources/EmbraceStorage/Records/EmbraceStorage+Span.swift

@@ -8,6 +8,8 @@ import GRDB
 
 extension EmbraceStorage {
 
+    static let defaultSpanLimitByType = 1500
+
     /// Adds a span to the storage synchronously.
     /// - Parameters:
     ///   - id: Identifier of the span
@@ -142,43 +144,36 @@ fileprivate extension EmbraceStorage {
         }
 
         // check limit and delete if necessary
-        if let limit = options.spanLimits[span.type.rawValue] {
-
-            let count = try spanCount(db: db, traceId: span.traceId, type: span.type)
-            if count >= limit {
-                let spansToDelete = try fetchSpans(
-                    db: db, traceId:
-                        span.traceId,
-                    type: span.type,
-                    limit: count - limit + 1
-                )
-
-                for spanToDelete in spansToDelete {
-                    try spanToDelete.delete(db)
-                }
+        // default to 1500 if limit is not set
+        let limit = options.spanLimits[span.type, default: Self.defaultSpanLimitByType]
+
+        let count = try spanCount(db: db, type: span.type)
+        if count >= limit {
+            let spansToDelete = try fetchSpans(
+                db: db,
+                type: span.type,
+                limit: count - limit + 1
+            )
+
+            for spanToDelete in spansToDelete {
+                try spanToDelete.delete(db)
             }
         }
 
         try span.insert(db)
     }
 
-    func spanInTraceByTypeRequest(traceId: String, type: SpanType?) -> QueryInterfaceRequest<SpanRecord> {
-        var filter = SpanRecord.filter(SpanRecord.Schema.traceId == traceId)
-
-        if let type = type {
-            filter = filter.filter(SpanRecord.Schema.type == type.rawValue)
-        }
-
-        return filter
+    func requestSpans(of type: SpanType) -> QueryInterfaceRequest<SpanRecord> {
+        return SpanRecord.filter(SpanRecord.Schema.type == type.rawValue)
     }
 
-    func spanCount(db: Database, traceId: String, type: SpanType?) throws -> Int {
-        return try spanInTraceByTypeRequest(traceId: traceId, type: type)
+    func spanCount(db: Database, type: SpanType) throws -> Int {
+        return try requestSpans(of: type)
             .fetchCount(db)
     }
 
-    func fetchSpans(db: Database, traceId: String, type: SpanType?, limit: Int?) throws -> [SpanRecord] {
-        var request = spanInTraceByTypeRequest(traceId: traceId, type: type)
+    func fetchSpans(db: Database, type: SpanType, limit: Int?) throws -> [SpanRecord] {
+        var request = requestSpans(of: type)
             .order(SpanRecord.Schema.startTime)
 
         if let limit = limit {

Tests/EmbraceStorageTests/Records/SpanRecord/EmbraceStorage+SpanTests.swift

@@ -0,0 +1,136 @@
+//
+//  Copyright © 2023 Embrace Mobile, Inc. All rights reserved.
+//
+
+import XCTest
+@testable import EmbraceStorage
+import OpenTelemetryApi
+
+final class EmbraceStorage_SpanTests: XCTestCase {
+
+    var storage: EmbraceStorage!
+
+    override func setUpWithError() throws {
+        storage = try EmbraceStorage.createInMemoryDb()
+    }
+
+    override func tearDownWithError() throws {
+        try storage.teardown()
+        storage = nil
+    }
+
+    func test_upsertSpan_appliesConfiguredLimitForType() throws {
+        storage.options.spanLimits[.performance] = 3
+
+        for i in 0..<3 {
+            // given inserted record
+            let span = SpanRecord(
+                id: SpanId.random().hexString,
+                name: "example \(i)",
+                traceId: TraceId.random().hexString,
+                type: .performance,
+                data: Data(),
+                startTime: Date()
+            )
+
+            try storage.upsertSpan(span)
+        }
+
+        try storage.upsertSpan(
+            SpanRecord(
+                id: SpanId.random().hexString,
+                name: "newest",
+                traceId: TraceId.random().hexString,
+                type: .performance,
+                data: Data(),
+                startTime: Date()
+            )
+        )
+
+        let allRecords = try storage.dbQueue.read { db in
+            try SpanRecord.order(SpanRecord.Schema.startTime).fetchAll(db)
+        }
+        XCTAssertEqual(allRecords.count, 3)
+        XCTAssertEqual(allRecords.map(\.name), ["example 1", "example 2", "newest"])
+    }
+
+    func test_upsertSpan_limitIsUniqueToSpecificType() throws {
+        storage.options.spanLimits[.performance] = 3
+        storage.options.spanLimits[.networkHTTP] = 1
+
+        // insert 3 .performance spans
+        for i in 0..<3 {
+            let span = SpanRecord(
+                id: SpanId.random().hexString,
+                name: "performance \(i)",
+                traceId: TraceId.random().hexString,
+                type: .performance,
+                data: Data(),
+                startTime: Date()
+            )
+
+            try storage.upsertSpan(span)
+        }
+
+        // insert 3 .networkHTTP spans
+        for i in 0..<3 {
+            let span = SpanRecord(
+                id: SpanId.random().hexString,
+                name: "network \(i)",
+                traceId: TraceId.random().hexString,
+                type: .networkHTTP,
+                data: Data(),
+                startTime: Date()
+            )
+
+            try storage.upsertSpan(span)
+        }
+
+        let allRecords = try storage.dbQueue.read { db in
+            try SpanRecord.order(SpanRecord.Schema.startTime).fetchAll(db)
+        }
+        XCTAssertEqual(allRecords.count, 4)
+        XCTAssertEqual(
+            allRecords.map(\.name),
+            [
+                "performance 0",
+                "performance 1",
+                "performance 2",
+                "network 2"
+            ]
+        )
+    }
+
+    func test_upsertSpan_appliesDefaultLimit() throws {
+        for i in 0..<(EmbraceStorage.defaultSpanLimitByType + 10) {
+            // given inserted record
+            let span = SpanRecord(
+                id: SpanId.random().hexString,
+                name: "example \(i)",
+                traceId: TraceId.random().hexString,
+                type: .performance,
+                data: Data(),
+                startTime: Date()
+            )
+
+            try storage.upsertSpan(span)
+        }
+
+        try storage.upsertSpan(
+            SpanRecord(
+                id: SpanId.random().hexString,
+                name: "newest",
+                traceId: TraceId.random().hexString,
+                type: .performance,
+                data: Data(),
+                startTime: Date()
+            )
+        )
+
+        let allRecords = try storage.dbQueue.read { db in
+            try SpanRecord.order(SpanRecord.Schema.startTime).fetchAll(db)
+        }
+        XCTAssertEqual(allRecords.count, EmbraceStorage.defaultSpanLimitByType) // 1500 is default limit
+    }
+
+}