Merge pull request #50 from rryam/fix/throttled-concurrent-file-operations

rudrankriyam · web-flow · commit 3d8ccfa92076 · 2025-12-08T14:09:08.000+09:00
Add throttled concurrency for file operations
diff --git a/Sources/VecturaKit/Core/VecturaKit.swift b/Sources/VecturaKit/Core/VecturaKit.swift
@@ -141,17 +141,8 @@ public actor VecturaKit {
       documentIds.append(docId)
     }
 
-    // Save documents to storage (storage handles caching internally)
-    let storage = self.storageProvider
-    try await withThrowingTaskGroup(of: Void.self) { group in
-      for doc in documentsToSave {
-        group.addTask {
-          try await storage.saveDocument(doc)
-        }
-      }
-
-      try await group.waitForAll()
-    }
+    // Save documents to storage (storage provider handles batch concurrency)
+    try await storageProvider.saveDocuments(documentsToSave)
 
     // Notify search engine to index documents
     for doc in documentsToSave {
diff --git a/Sources/VecturaKit/Storage/FileStorageProvider.swift b/Sources/VecturaKit/Storage/FileStorageProvider.swift
@@ -99,12 +99,63 @@ extension FileStorageProvider: VecturaStorage {
       cache[document.id] = document
     }
   }
+
+  /// Saves multiple documents with throttled concurrent file writes.
+  ///
+  /// Uses nonisolated file I/O to achieve true parallelism, then updates
+  /// the cache after all writes complete.
+  public func saveDocuments(_ documents: [VecturaDocument]) async throws {
+    let directory = storageDirectory
+
+    // Perform concurrent file writes outside actor isolation
+    try await documents.concurrentForEach(maxConcurrency: Self.maxConcurrentFileOperations) { document in
+      try Self.writeDocumentToFile(document, in: directory)
+    }
+
+    // Update cache after all writes succeed
+    if cacheEnabled {
+      for document in documents {
+        cache[document.id] = document
+      }
+    }
+  }
+
+  /// Writes a document to disk without actor isolation.
+  ///
+  /// This allows concurrent file I/O when called from multiple tasks.
+  nonisolated private static func writeDocumentToFile(
+    _ document: VecturaDocument,
+    in directory: URL
+  ) throws {
+    let encoder = JSONEncoder()
+    encoder.outputFormatting = .prettyPrinted
+    let data = try encoder.encode(document)
+    let documentURL = directory.appendingPathComponent("\(document.id).json")
+
+    #if os(iOS) || os(tvOS) || os(watchOS) || os(visionOS)
+    try data.write(to: documentURL, options: [.atomic, .completeFileProtection])
+    #else
+    try data.write(to: documentURL, options: .atomic)
+    #endif
+
+    try FileManager.default.setAttributes(
+      [.posixPermissions: 0o600],
+      ofItemAtPath: documentURL.path(percentEncoded: false)
+    )
+  }
 }
 
 // MARK: - CachableVecturaStorage
 
 extension FileStorageProvider: CachableVecturaStorage {
+
+  /// Maximum number of concurrent file operations to prevent resource exhaustion.
+  private static let maxConcurrentFileOperations = 50
+
   /// Loads all documents from disk by reading JSON files (bypasses cache).
+  ///
+  /// Uses throttled concurrency to prevent file descriptor exhaustion
+  /// when loading large numbers of documents.
   public func loadDocumentsFromStorage() async throws -> [VecturaDocument] {
     let fileURLs = try FileManager.default.contentsOfDirectory(
       at: storageDirectory,
@@ -113,31 +164,17 @@ extension FileStorageProvider: CachableVecturaStorage {
 
     let jsonFileURLs = fileURLs.filter { $0.pathExtension.lowercased() == "json" }
 
-    return await withTaskGroup(of: VecturaDocument?.self, returning: [VecturaDocument].self) { group in
-      for fileURL in jsonFileURLs {
-        group.addTask {
-          do {
-            let data = try Data(contentsOf: fileURL)
-            let decoder = JSONDecoder()
-            return try decoder.decode(VecturaDocument.self, from: data)
-          } catch {
-            let path = fileURL.path(percentEncoded: false)
-            Self.logger.warning(
-              "Failed to load document at \(path): \(error.localizedDescription)"
-            )
-            return nil
-          }
-        }
-      }
-
-      var documents: [VecturaDocument] = []
-      documents.reserveCapacity(jsonFileURLs.count)
-      for await document in group {
-        if let document {
-          documents.append(document)
-        }
+    return await jsonFileURLs.concurrentMap(maxConcurrency: Self.maxConcurrentFileOperations) { fileURL in
+      do {
+        let data = try Data(contentsOf: fileURL)
+        return try JSONDecoder().decode(VecturaDocument.self, from: data)
+      } catch {
+        let path = fileURL.path(percentEncoded: false)
+        Self.logger.warning(
+          "Failed to load document at \(path): \(error.localizedDescription)"
+        )
+        return nil
       }
-      return documents
     }
   }
 
diff --git a/Sources/VecturaKit/Storage/VecturaStorage.swift b/Sources/VecturaKit/Storage/VecturaStorage.swift
@@ -33,6 +33,14 @@ public protocol VecturaStorage: Sendable {
   ///
   /// - Returns: The total document count.
   func getTotalDocumentCount() async throws -> Int
+
+  /// Saves multiple documents in batch.
+  ///
+  /// Storage providers can override this for optimized batch operations.
+  /// The default implementation calls saveDocument sequentially.
+  ///
+  /// - Parameter documents: The documents to save.
+  func saveDocuments(_ documents: [VecturaDocument]) async throws
 }
 
 // MARK: - Default Implementation
@@ -45,4 +53,13 @@ extension VecturaStorage {
   public func getTotalDocumentCount() async throws -> Int {
     return try await loadDocuments().count
   }
+
+  /// Default implementation that saves documents sequentially.
+  ///
+  /// Storage implementations can override this for concurrent I/O.
+  public func saveDocuments(_ documents: [VecturaDocument]) async throws {
+    for document in documents {
+      try await saveDocument(document)
+    }
+  }
 }
diff --git a/Sources/VecturaKit/Utilities/ConcurrencyExtensions.swift b/Sources/VecturaKit/Utilities/ConcurrencyExtensions.swift
@@ -0,0 +1,207 @@
+import Foundation
+
+// MARK: - Concurrent Collection Processing
+
+extension Sequence where Element: Sendable {
+
+  /// Concurrently maps elements with controlled parallelism using a sliding window pattern.
+  ///
+  /// This method processes elements concurrently while limiting the number of simultaneous
+  /// operations to prevent resource exhaustion (file handles, memory, etc.).
+  ///
+  /// - Parameters:
+  ///   - maxConcurrency: Maximum number of concurrent operations (default: 50)
+  ///   - transform: Async throwing closure that transforms each element. Return `nil` to skip.
+  /// - Returns: Array of non-nil transformed results
+  /// - Throws: Rethrows any error from the transform closure
+  ///
+  /// ## Example
+  ///
+  /// ```swift
+  /// let urls: [URL] = getFileURLs()
+  /// let documents = try await urls.concurrentMap(maxConcurrency: 50) { url in
+  ///   try? JSONDecoder().decode(Document.self, from: Data(contentsOf: url))
+  /// }
+  /// ```
+  ///
+  /// ## Performance Characteristics
+  ///
+  /// - Maintains exactly `maxConcurrency` tasks running at any time
+  /// - Uses structured concurrency with no semaphores or locks
+  /// - Memory efficient: processes results as they complete
+  @inlinable
+  public func concurrentMap<T: Sendable>(
+    maxConcurrency: Int = 50,
+    _ transform: @Sendable @escaping (Element) async throws -> T?
+  ) async rethrows -> [T] {
+    try await withThrowingTaskGroup(of: T?.self) { group in
+      var results: [T] = []
+      var iterator = makeIterator()
+
+      // Seed initial batch of tasks up to maxConcurrency
+      var activeCount = 0
+      while activeCount < maxConcurrency, let element = iterator.next() {
+        group.addTask { try await transform(element) }
+        activeCount += 1
+      }
+
+      // As each task completes, add the next one (sliding window)
+      while let result = try await group.next() {
+        if let value = result {
+          results.append(value)
+        }
+
+        // Add next task if elements remain
+        if let element = iterator.next() {
+          group.addTask { try await transform(element) }
+        }
+      }
+
+      return results
+    }
+  }
+
+  /// Concurrently maps elements with controlled parallelism (non-throwing version).
+  ///
+  /// - Parameters:
+  ///   - maxConcurrency: Maximum number of concurrent operations (default: 50)
+  ///   - transform: Async closure that transforms each element. Return `nil` to skip.
+  /// - Returns: Array of non-nil transformed results
+  @inlinable
+  public func concurrentMap<T: Sendable>(
+    maxConcurrency: Int = 50,
+    _ transform: @Sendable @escaping (Element) async -> T?
+  ) async -> [T] {
+    await withTaskGroup(of: T?.self) { group in
+      var results: [T] = []
+      var iterator = makeIterator()
+
+      // Seed initial batch
+      var activeCount = 0
+      while activeCount < maxConcurrency, let element = iterator.next() {
+        group.addTask { await transform(element) }
+        activeCount += 1
+      }
+
+      // Sliding window: add new task as each completes
+      for await result in group {
+        if let value = result {
+          results.append(value)
+        }
+
+        if let element = iterator.next() {
+          group.addTask { await transform(element) }
+        }
+      }
+
+      return results
+    }
+  }
+
+  /// Concurrently executes a side-effect closure on each element with controlled parallelism.
+  ///
+  /// Use this when you need to perform async operations for their side effects
+  /// (e.g., saving files, network requests) without collecting results.
+  ///
+  /// - Parameters:
+  ///   - maxConcurrency: Maximum number of concurrent operations (default: 50)
+  ///   - body: Async throwing closure to execute for each element
+  /// - Throws: Rethrows the first error encountered
+  ///
+  /// ## Example
+  ///
+  /// ```swift
+  /// try await documents.concurrentForEach(maxConcurrency: 20) { doc in
+  ///   try await storage.saveDocument(doc)
+  /// }
+  /// ```
+  @inlinable
+  public func concurrentForEach(
+    maxConcurrency: Int = 50,
+    _ body: @Sendable @escaping (Element) async throws -> Void
+  ) async throws {
+    try await withThrowingTaskGroup(of: Void.self) { group in
+      var iterator = makeIterator()
+
+      // Seed initial batch
+      var activeCount = 0
+      while activeCount < maxConcurrency, let element = iterator.next() {
+        group.addTask { try await body(element) }
+        activeCount += 1
+      }
+
+      // Sliding window
+      while try await group.next() != nil {
+        if let element = iterator.next() {
+          group.addTask { try await body(element) }
+        }
+      }
+    }
+  }
+
+  /// Concurrently maps elements while preserving the original order.
+  ///
+  /// Unlike `concurrentMap`, this method guarantees that output order matches input order.
+  /// Slightly higher memory overhead due to tracking indices.
+  ///
+  /// - Parameters:
+  ///   - maxConcurrency: Maximum number of concurrent operations (default: 50)
+  ///   - transform: Async throwing closure that transforms each element
+  /// - Returns: Array of transformed results in the same order as input
+  /// - Throws: Rethrows any error from the transform closure
+  ///
+  /// ## Example
+  ///
+  /// ```swift
+  /// let urls = [url1, url2, url3]
+  /// let data = try await urls.orderedConcurrentMap(maxConcurrency: 10) { url in
+  ///   try await fetchData(from: url)
+  /// }
+  /// // data[0] corresponds to url1, data[1] to url2, etc.
+  /// ```
+  @inlinable
+  public func orderedConcurrentMap<T: Sendable>(
+    maxConcurrency: Int = 50,
+    _ transform: @Sendable @escaping (Element) async throws -> T
+  ) async rethrows -> [T] {
+    let indexed = Array(self.enumerated())
+
+    let results = try await indexed.concurrentMap(maxConcurrency: maxConcurrency) { item -> (Int, T)? in
+      let result = try await transform(item.element)
+      return (item.offset, result)
+    }
+
+    // Sort by original index and extract values
+    return results
+      .sorted { $0.0 < $1.0 }
+      .map(\.1)
+  }
+}
+
+// MARK: - Collection Chunking
+
+extension Collection {
+
+  /// Splits the collection into chunks of the specified size.
+  ///
+  /// - Parameter size: Maximum size of each chunk
+  /// - Returns: Array of array slices, each containing up to `size` elements
+  ///
+  /// ## Example
+  ///
+  /// ```swift
+  /// let items = [1, 2, 3, 4, 5, 6, 7]
+  /// let chunks = items.chunked(into: 3)
+  /// // [[1, 2, 3], [4, 5, 6], [7]]
+  /// ```
+  @inlinable
+  public func chunked(into size: Int) -> [[Element]] {
+    guard size > 0 else { return [] }
+
+    return stride(from: 0, to: count, by: size).map { startOffset in
+      let start = index(startIndex, offsetBy: startOffset)
+      let end = index(start, offsetBy: size, limitedBy: endIndex) ?? endIndex
+      return Array(self[start..<end])
+    }
+  }
+}
