Introduce concurreny limiter abstraction, a trimmed down

OmSaran · OmSaran · commit 7ad076c27324 · 2022-08-09T08:42:08.000-07:00
version of token bucket.

Add tests
diff --git a/Package.resolved b/Package.resolved
@@ -15,13 +15,8 @@
         "repositoryURL": "https://github.com/apple/swift-nio.git",
         "state": {
           "branch": null,
-<<<<<<< HEAD
           "revision": "ece5057615d1bee848341eceafdf04ca54d60177",
           "version": "2.41.0"
-=======
-          "revision": "124119f0bb12384cef35aa041d7c3a686108722d",
-          "version": "2.40.0"
->>>>>>> 66928b7 (To support linux better, replace operation queue with dispatch queue in BatchingFutureOperationQueue)
         }
       },
       {
diff --git a/Package.swift b/Package.swift
@@ -12,11 +12,7 @@ let package = Package(
             targets: ["TSFCAS", "TSFCASFileTree", "TSFCASUtilities"]),
     ],
     dependencies: [
-<<<<<<< HEAD
         .package(url: "https://github.com/apple/swift-nio.git", from: "2.32.0"),
-=======
-        .package(url: "https://github.com/apple/swift-nio.git", from: "2.40.0"),
->>>>>>> d4b5e97 (Remove OperationQueue implementation altogether and use equivalent dispatchQoS)
         .package(url: "https://github.com/apple/swift-protobuf.git", from: "1.8.0"),
         .package(url: "https://github.com/apple/swift-tools-support-core.git", from: "0.2.3"),
     ],
diff --git a/Sources/TSFFutures/BatchingFutureOperationQueue.swift b/Sources/TSFFutures/BatchingFutureOperationQueue.swift
@@ -8,20 +8,18 @@
 
 import Dispatch
 import Foundation
-import NIOConcurrencyHelpers
 import NIO
-
+import NIOConcurrencyHelpers
 import TSCUtility
 
-
 /// Run the given computations on a given array in batches, exercising
 /// a specified amount of parallelism.
 ///
 /// - Discussion:
 ///     For some blocking operations (such as file system accesses) executing
 ///     them on the NIO loops is very expensive since it blocks the event
 ///     processing machinery. Here we use extra threads for such operations.
-public class LLBBatchingFutureOperationQueue {
+public struct LLBBatchingFutureOperationQueue {
 
     /// Threads capable of running futures.
     public let group: LLBFuturesDispatchGroup
@@ -35,28 +33,33 @@ public class LLBBatchingFutureOperationQueue {
 
     /// Maximum number of operations executed concurrently.
     public var maxOpCount: Int {
-        get { lock.withLock { maxOpCount_ } }
-        set { scheduleMoreTasks { maxOpCount_ = newValue } }
+        get { concurrencyLimiter.maximumConcurrency }
+        set { concurrencyLimiter.maximumConcurrency = Self.bridged(maxOpCount: newValue) }
     }
 
     /// Return the number of operations currently queued.
-    public var opCount: Int { lock.withLock { opCount_ } }
+    public var opCount: Int { concurrencyLimiter.sharesInUse }
+
+    /// Name to be used for dispatch queue
+    private let name: String
 
-    /// Queue of outstanding operations
-    private let dispatchQueue: DispatchQueue
+    /// QoS passed to DispatchQueue
+    private let qos: DispatchQoS
 
     /// Lock protecting state.
     private let lock = NIOConcurrencyHelpers.Lock()
 
-    private var maxOpCount_: Int
-
-    private var opCount_: Int
+    /// Limits number of concurrent operations being executed
+    private let concurrencyLimiter: ConcurrencyLimiter
 
     /// The queue of operations to run.
     private var workQueue = NIO.CircularBuffer<DispatchWorkItem>()
 
     @available(*, deprecated, message: "'qualityOfService' is deprecated: Use 'dispatchQoS'")
-    public convenience init(name: String, group: LLBFuturesDispatchGroup, maxConcurrentOperationCount maxOpCount: Int, qualityOfService: QualityOfService) {
+    public init(
+        name: String, group: LLBFuturesDispatchGroup, maxConcurrentOperationCount maxOpCount: Int,
+        qualityOfService: QualityOfService
+    ) {
         let dispatchQoS: DispatchQoS
 
         switch qualityOfService {
@@ -81,94 +84,72 @@ public class LLBBatchingFutureOperationQueue {
     ///    - group:     Threads capable of running futures.
     ///    - maxConcurrentOperationCount:
     ///                 Operations to execute in parallel.
-    public convenience init(name: String, group: LLBFuturesDispatchGroup, maxConcurrentOperationCount maxOpCount: Int) {
+    public init(name: String, group: LLBFuturesDispatchGroup, maxConcurrentOperationCount maxOpCount: Int) {
         self.init(name: name, group: group, maxConcurrentOperationCount: maxOpCount, dispatchQoS: .default)
     }
 
-    public init(name: String, group: LLBFuturesDispatchGroup, maxConcurrentOperationCount maxOpCnt: Int, dispatchQoS: DispatchQoS) {
+    public init(
+        name: String, group: LLBFuturesDispatchGroup, maxConcurrentOperationCount maxOpCnt: Int,
+        dispatchQoS: DispatchQoS
+    ) {
         self.group = group
-        self.dispatchQueue = DispatchQueue(label: name, qos: dispatchQoS, attributes: .concurrent)
-        self.opCount_ = 0
-        self.maxOpCount_ = maxOpCnt
+        self.name = name
+        self.qos = dispatchQoS
+
+        self.concurrencyLimiter = ConcurrencyLimiter(maximumConcurrency: Self.bridged(maxOpCount: maxOpCnt))
     }
 
     public func execute<T>(_ body: @escaping () throws -> T) -> LLBFuture<T> {
-        let promise = group.any().makePromise(of: T.self)
+        return self.concurrencyLimiter.withReplenishableLimit(eventLoop: group.any()) { eventLoop in
+            let promise = eventLoop.makePromise(of: T.self)
 
-        let workItem = DispatchWorkItem {
-            promise.fulfill(body)
-            self.scheduleMoreTasks {
-                self.opCount_ -= 1
+            DispatchQueue(label: self.name, qos: self.qos).async {
+                promise.fulfill(body)
             }
-        }
 
-        self.scheduleMoreTasks {
-            workQueue.append(workItem)
+            return promise.futureResult
         }
-
-        return promise.futureResult
     }
 
     public func execute<T>(_ body: @escaping () -> LLBFuture<T>) -> LLBFuture<T> {
-        let promise = group.any().makePromise(of: T.self)
+        return self.concurrencyLimiter.withReplenishableLimit(eventLoop: group.any()) { eventLoop in
+            let promise = eventLoop.makePromise(of: T.self)
 
-        let workItem = DispatchWorkItem {
-            let f = body()
-            f.cascade(to: promise)
-
-            _ = try? f.wait()
-
-            self.scheduleMoreTasks {
-                self.opCount_ -= 1
+            DispatchQueue(label: self.name, qos: self.qos).async {
+                body().cascade(to: promise)
             }
-        }
 
-        self.scheduleMoreTasks {
-            workQueue.append(workItem)
+            return promise.futureResult
         }
-
-        return promise.futureResult
     }
 
     /// Order-preserving parallel execution. Wait for everything to complete.
     @inlinable
-    public func execute<A,T>(_ args: [A], minStride: Int = 1, _ body: @escaping (ArraySlice<A>) throws -> [T]) -> LLBFuture<[T]> {
+    public func execute<A, T>(_ args: [A], minStride: Int = 1, _ body: @escaping (ArraySlice<A>) throws -> [T])
+        -> LLBFuture<[T]>
+    {
         let futures: [LLBFuture<[T]>] = executeNoWait(args, minStride: minStride, body)
         let loop = futures.first?.eventLoop ?? group.next()
-        return LLBFuture<[T]>.whenAllSucceed(futures, on: loop).map{$0.flatMap{$0}}
+        return LLBFuture<[T]>.whenAllSucceed(futures, on: loop).map { $0.flatMap { $0 } }
     }
 
     /// Order-preserving parallel execution.
     /// Do not wait for all executions to complete, returning individual futures.
     @inlinable
-    public func executeNoWait<A,T>(_ args: [A], minStride: Int = 1, maxStride: Int = Int.max, _ body: @escaping (ArraySlice<A>) throws -> [T]) -> [LLBFuture<[T]>] {
-        let batches: [ArraySlice<A>] = args.tsc_sliceBy(maxStride: max(minStride, min(maxStride, args.count / maxOpCount)))
-        return batches.map{arg in execute{try body(arg)}}
+    public func executeNoWait<A, T>(
+        _ args: [A], minStride: Int = 1, maxStride: Int = Int.max, _ body: @escaping (ArraySlice<A>) throws -> [T]
+    ) -> [LLBFuture<[T]>] {
+        let batches: [ArraySlice<A>] = args.tsc_sliceBy(
+            maxStride: max(minStride, min(maxStride, args.count / maxOpCount)))
+        return batches.map { arg in execute { try body(arg) } }
     }
 
-    private func scheduleMoreTasks(performUnderLock: () -> Void) {
-        let toExecute: [DispatchWorkItem] = lock.withLock {
-            performUnderLock()
-
-            var scheduleItems: [DispatchWorkItem] = []
-
-            while opCount_ < maxOpCount_ {
-
-                // Schedule a new operation, if available.
-                guard let op = workQueue.popFirst() else {
-                    break
-                }
-
-                self.opCount_ += 1
-                scheduleItems.append(op)
-            }
-
-            return scheduleItems
-        }
-
-        for workItem in toExecute {
-            dispatchQueue.async(execute: workItem)
+    private static func bridged(maxOpCount: Int) -> Int {
+        switch maxOpCount {
+        case OperationQueue.defaultMaxConcurrentOperationCount:
+            return System.coreCount
+        default:
+            return maxOpCount
         }
     }
-
 }
diff --git a/Sources/TSFFutures/ConcurrencyLimiter.swift b/Sources/TSFFutures/ConcurrencyLimiter.swift
@@ -0,0 +1,111 @@
+//===----------------------------------------------------------------------===//
+//
+// This source file is part of the Swift.org open source project
+//
+// Copyright (c) 2022 Apple Inc. and the Swift project authors
+// Licensed under Apache License v2.0 with Runtime Library Exception
+//
+// See https://swift.org/LICENSE.txt for license information
+// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+//
+//===----------------------------------------------------------------------===//
+
+import NIO
+import NIOConcurrencyHelpers
+
+internal final class ConcurrencyLimiter {
+    public var maximumConcurrency: Int {
+        get { lock.withLock { maximumConcurrency_ } }
+        set {
+            var waiters: [Waiter] = []
+            lock.withLockVoid {
+                let differenceInCapacity = (newValue - maximumConcurrency_)
+                unitsLeft += differenceInCapacity
+                maximumConcurrency_ = newValue
+                waiters = tryFulfillSomeLocked()
+            }
+            waiters.fulfillWaiters()
+        }
+    }
+    public var sharesInUse: Int { lock.withLock { maximumConcurrency_ - unitsLeft } }
+
+    private var unitsLeft: Int  // protected by `self.lock`
+    private var waiters: CircularBuffer<Waiter> = []  // protected by `self.lock`
+    private let lock = Lock()
+    private var maximumConcurrency_: Int  // protected by `self.lock`
+
+    public init(maximumConcurrency: Int) {
+        precondition(maximumConcurrency >= 0)
+
+        self.maximumConcurrency_ = maximumConcurrency
+        self.unitsLeft = maximumConcurrency
+    }
+
+    /// Reserves 1 unit of concurrency, executes body after which it restores the 1 unit.
+    public func withReplenishableLimit<T>(
+        eventLoop: EventLoop,
+        _ body: @escaping (EventLoop) -> EventLoopFuture<T>
+    ) -> EventLoopFuture<T> {
+        return self.withdraw(eventLoop: eventLoop).flatMap { lease in
+            body(eventLoop).always { _ in
+                self.replenish(lease)
+            }
+        }
+    }
+
+    private func tryFulfillSomeLocked() -> [Waiter] {
+        var toSucceed: [Waiter] = []
+        let unitsLeftAtStart = self.unitsLeft
+
+        while !self.waiters.isEmpty, self.unitsLeft >= 1 {
+            let waiter = self.waiters.removeFirst()
+
+            self.unitsLeft -= 1
+            assert(self.unitsLeft >= 0)
+            toSucceed.append(waiter)
+        }
+
+        assert(unitsLeftAtStart - toSucceed.count == self.unitsLeft)
+        return toSucceed
+    }
+
+    private func replenish(_ lease: Lease) {
+        self.lock.withLock { () -> [Waiter] in
+            self.unitsLeft += 1
+            assert(self.unitsLeft <= self.maximumConcurrency_)
+            return self.tryFulfillSomeLocked()
+        }.fulfillWaiters()
+    }
+
+    /// Reserve 1 unit of the limit if available
+    private func withdraw(eventLoop: EventLoop) -> EventLoopFuture<Lease> {
+        let future = self.lock.withLock { () -> EventLoopFuture<Lease> in
+            if self.waiters.isEmpty && self.unitsLeft >= 1 {
+                self.unitsLeft -= 1
+
+                return eventLoop.makeSucceededFuture(Lease())
+            }
+
+            let promise = eventLoop.makePromise(of: Lease.self)
+            self.waiters.append(Waiter(promise: promise))
+
+            return promise.futureResult
+        }
+
+        return future
+    }
+
+    fileprivate struct Waiter {
+        var promise: EventLoopPromise<Lease>
+    }
+
+    fileprivate struct Lease {}
+}
+
+extension Array where Element == ConcurrencyLimiter.Waiter {
+    fileprivate func fulfillWaiters() {
+        self.forEach { waiter in
+            return waiter.promise.succeed(.init())
+        }
+    }
+}
diff --git a/Tests/TSFFuturesTests/BatchingFutureOperationQueue.swift b/Tests/TSFFuturesTests/BatchingFutureOperationQueue.swift
diff --git a/Tests/TSFFuturesTests/ConcurrencyLimiterTests.swift b/Tests/TSFFuturesTests/ConcurrencyLimiterTests.swift

Original file line number	Diff line number	Diff line change
`@@ -15,13 +15,8 @@`
`15`	`15`	`"repositoryURL": "https://github.com/apple/swift-nio.git",`
`16`	`16`	`"state": {`
`17`	`17`	`"branch": null,`
`18`		`-<<<<<<< HEAD`
`19`	`18`	`"revision": "ece5057615d1bee848341eceafdf04ca54d60177",`
`20`	`19`	`"version": "2.41.0"`
`21`		`-=======`
`22`		`- "revision": "124119f0bb12384cef35aa041d7c3a686108722d",`
`23`		`- "version": "2.40.0"`
`24`		`->>>>>>> 66928b7 (To support linux better, replace operation queue with dispatch queue in BatchingFutureOperationQueue)`
`25`	`20`	`}`
`26`	`21`	`},`
`27`	`22`	`{`