Add TaggedScheduler, couple lifetime of tasks and orchestrator (#2398)

src/mbgl/util/thread_pool.cpp

@@ -11,11 +11,8 @@ namespace mbgl {
 ThreadedSchedulerBase::~ThreadedSchedulerBase() = default;
 
 void ThreadedSchedulerBase::terminate() {
-    // Run any leftover render jobs
-    runRenderJobs();
-
     {
-        std::lock_guard<std::mutex> lock(mutex);
+        std::lock_guard<std::mutex> lock(workerMutex);
         terminated = true;
     }
 
@@ -37,48 +34,64 @@ std::thread ThreadedSchedulerBase::makeSchedulerThread(size_t index) {
         owningThreadPool.set(this);
 
         while (true) {
-            std::unique_lock<std::mutex> lock(mutex);
-            if (queue.empty() && !pendingItems) {
-                cvEmpty.notify_all();
+            std::unique_lock<std::mutex> conditionLock(workerMutex);
+            if (!terminated && taskCount == 0) {
+                cvAvailable.wait(conditionLock);
             }
 
-            cvAvailable.wait(lock, [this] { return !queue.empty() || terminated; });
-
             if (terminated) {
                 platform::detachThread();
-                return;
+                break;
             }
 
-            auto function = std::move(queue.front());
-            queue.pop();
+            // Let other threads run
+            conditionLock.unlock();
 
-            if (function) {
-                pendingItems++;
+            std::vector<std::shared_ptr<Queue>> pending;
+            {
+                // 1. Gather buckets for us to visit this iteration
+                std::lock_guard<std::mutex> lock(taggedQueueLock);
+                for (const auto& [tag, queue] : taggedQueue) {
+                    pending.push_back(queue);
+                }
             }
 
-            lock.unlock();
+            // 2. Visit a task from each
+            for (auto& q : pending) {
+                std::function<void()> tasklet;
+                {
+                    std::lock_guard<std::mutex> lock(q->lock);
+                    if (q->queue.size()) {
+                        q->runningCount++;
+                        tasklet = std::move(q->queue.front());
+                        q->queue.pop();
+                    }
+                    if (!tasklet) continue;
+                }
+
+                assert(taskCount > 0);
+                taskCount--;
 
-            if (function) {
                 try {
-                    function();
+                    tasklet();
+                    tasklet = {}; // destroy the function and release its captures before unblocking `waitForEmpty`
 
-                    // destroy the function and release its captures before unblocking `waitForEmpty`
-                    function = {};
-                    if (!--pendingItems) {
-                        std::unique_lock<std::mutex> inner_lock(mutex);
-                        if (queue.empty()) {
-                            cvEmpty.notify_all();
+                    if (!--q->runningCount) {
+                        std::lock_guard<std::mutex> lock(q->lock);
+                        if (q->queue.empty()) {
+                            q->cv.notify_all();
                         }
                     }
                 } catch (...) {
-                    std::unique_lock<std::mutex> inner_lock(mutex);
+                    std::lock_guard<std::mutex> lock(q->lock);
                     if (handler) {
                         handler(std::current_exception());
                     }
 
-                    function = {};
-                    if (!--pendingItems && queue.empty()) {
-                        cvEmpty.notify_all();
+                    tasklet = {};
+
+                    if (!--q->runningCount && q->queue.empty()) {
+                        q->cv.notify_all();
                     }
 
                     if (handler) {
@@ -92,47 +105,61 @@ std::thread ThreadedSchedulerBase::makeSchedulerThread(size_t index) {
 }
 
 void ThreadedSchedulerBase::schedule(std::function<void()>&& fn) {
-    assert(fn);
-    if (fn) {
-        {
-            // We need to block if adding adding a new task from a thread not controlled by this
-            // pool.  Tasks are added by other tasks, so we must not block a thread we do control
-            // or `waitForEmpty` will deadlock.
-            std::unique_lock<std::mutex> addLock(addMutex, std::defer_lock);
-            if (!thisThreadIsOwned()) {
-                addLock.lock();
-            }
-            std::lock_guard<std::mutex> lock(mutex);
-            queue.push(std::move(fn));
-        }
-        cvAvailable.notify_one();
-    }
+    schedule(uniqueID, std::move(fn));
 }
 
-std::size_t ThreadedSchedulerBase::waitForEmpty(Milliseconds timeout) {
+void ThreadedSchedulerBase::schedule(const util::SimpleIdentity tag, std::function<void()>&& fn) {
+    assert(fn);
+    if (!fn) return;
+
+    std::shared_ptr<Queue> q;
+    {
+        std::lock_guard<std::mutex> lock(taggedQueueLock);
+        auto it = taggedQueue.find(tag);
+        if (it == taggedQueue.end()) {
+            q = std::make_shared<Queue>();
+            taggedQueue.insert({tag, q});
+        } else {
+            q = it->second;
+        }
+    }
+
+    {
+        std::lock_guard<std::mutex> lock(q->lock);
+        q->queue.push(std::move(fn));
+        taskCount++;
+    }
+
+    cvAvailable.notify_one();
+}
+
+void ThreadedSchedulerBase::waitForEmpty(const util::SimpleIdentity tag) {
     // Must not be called from a thread in our pool, or we would deadlock
     assert(!thisThreadIsOwned());
     if (!thisThreadIsOwned()) {
-        const auto startTime = util::MonotonicTimer::now();
-        const auto isDone = [&] {
-            return queue.empty() && pendingItems == 0;
-        };
-        // Block any other threads from adding new items
-        std::scoped_lock<std::mutex> addLock(addMutex);
-        std::unique_lock<std::mutex> lock(mutex);
-        while (!isDone()) {
-            if (timeout > Milliseconds::zero()) {
-                const auto elapsed = util::MonotonicTimer::now() - startTime;
-                if (timeout <= elapsed || !cvEmpty.wait_for(lock, timeout - elapsed, isDone)) {
-                    break;
-                }
-            } else {
-                cvEmpty.wait(lock, isDone);
+        const auto tagToFind = tag.isEmpty() ? uniqueID : tag;
+
+        std::shared_ptr<Queue> q;
+        {
+            std::lock_guard<std::mutex> lock(taggedQueueLock);
+            auto it = taggedQueue.find(tagToFind);
+            if (it == taggedQueue.end()) {
+                return;
             }
+            q = it->second;
+        }
+
+        std::unique_lock<std::mutex> queueLock(q->lock);
+        while (q->queue.size() + q->runningCount) {
+            q->cv.wait(queueLock);
+        }
+
+        // After waiting for the queue to empty, go ahead and erase it from the map.
+        {
+            std::lock_guard<std::mutex> lock(taggedQueueLock);
+            taggedQueue.erase(tagToFind);
         }
-        return queue.size() + pendingItems;
     }
-    return 0;
 }
 
 } // namespace mbgl