Reference safety for Worker after callbacks

Joining a WorkRequest now ensures that either the after callback has
already been executed, or will never be executed.

Fixes #1438

src/mbgl/util/work_request.cpp

@@ -4,22 +4,25 @@ namespace mbgl {
 
 WorkRequest::WorkRequest() = default;
 
-WorkRequest::WorkRequest(Future&& future)
-    : complete(std::move(future)) {
+WorkRequest::WorkRequest(Future&& future, JoinedFlag flag)
+    : complete(std::move(future)),
+      joinedFlag(flag) {
 }
 
 WorkRequest::WorkRequest(WorkRequest&& o)
-    : complete(std::move(o.complete)) {
+    : complete(std::move(o.complete)),
+      joinedFlag(std::move(o.joinedFlag)) {
 }
 
 WorkRequest::~WorkRequest() {
     if (complete.valid()) {
-        complete.get();
+        join();
     }
 }
 
 WorkRequest& WorkRequest::operator=(WorkRequest&& o) {
     complete = std::move(o.complete);
+    joinedFlag = std::move(o.joinedFlag);
     return *this;
 }
 
@@ -28,6 +31,7 @@ WorkRequest::operator bool() const {
 }
 
 void WorkRequest::join() {
+    *joinedFlag = true;
     complete.get();
 }
 

src/mbgl/util/work_request.hpp

@@ -10,9 +10,10 @@ namespace mbgl {
 class WorkRequest : public util::noncopyable {
 public:
     using Future = std::future<void>;
+    using JoinedFlag = std::shared_ptr<std::atomic<bool>>;
 
     WorkRequest();
-    WorkRequest(Future&&);
+    WorkRequest(Future&&, JoinedFlag);
     WorkRequest(WorkRequest&&);
     ~WorkRequest();
 
@@ -24,6 +25,7 @@ public:
 
 private:
     Future complete;
+    JoinedFlag joinedFlag;
 };
 
 }

src/mbgl/util/worker.cpp

@@ -27,10 +27,17 @@ WorkRequest Worker::send(Fn work, Fn after) {
     std::packaged_task<void ()> task(work);
     std::future<void> future = task.get_future();
 
-    threads[current]->invokeWithResult(&Worker::Impl::doWork, std::move(after), task);
+    std::shared_ptr<std::atomic<bool>> joined = std::make_shared<std::atomic<bool>>();
+    *joined = false;
+
+    threads[current]->invokeWithResult(&Worker::Impl::doWork, [joined, after] {
+        if (!*joined) {
+            after();
+        }
+    }, task);
 
     current = (current + 1) % threads.size();
-    return WorkRequest(std::move(future));
+    return WorkRequest(std::move(future), joined);
 }
 
 }

src/mbgl/util/worker.hpp

@@ -21,9 +21,11 @@ public:
     // work is complete.
     //
     // The return value represents the request to perform the work asynchronously.
-    // Its destructor guarantees that the work is either cancelled and will never
-    // execute, or has finished executing. In other words, the WorkRequest is
-    // guaranteed to outlive any references held by the work function.
+    // Its destructor guarantees that the work function has finished executing, and
+    // that the after function has either finished executing or will not execute.
+    // Together, this means that an object may make a work request with lambdas which
+    // bind references to itself, and if and when those lambdas execute, the references
+    // will still be valid.
     WorkRequest send(Fn work, Fn after);
 
 private:

test/miscellaneous/worker.cpp

@@ -0,0 +1,70 @@
+#include "../fixtures/util.hpp"
+
+#include <mbgl/util/worker.hpp>
+#include <mbgl/util/run_loop.hpp>
+
+using namespace mbgl;
+using namespace mbgl::util;
+
+TEST(Worker, ExecutesWorkAndAfter) {
+    RunLoop loop(uv_default_loop());
+
+    Worker worker(1);
+    WorkRequest request;
+
+    bool didWork = false;
+    bool didAfter = false;
+
+    loop.invoke([&] {
+        request = worker.send([&] {
+            didWork = true;
+        }, [&] {
+            didAfter = true;
+            loop.stop();
+        });
+    });
+
+    uv_run(uv_default_loop(), UV_RUN_DEFAULT);
+    EXPECT_TRUE(didWork);
+    EXPECT_TRUE(didAfter);
+}
+
+TEST(Worker, WorkRequestJoinWaitsForWorkToComplete) {
+    RunLoop loop(uv_default_loop());
+
+    Worker worker(1);
+    bool didWork = false;
+
+    loop.invoke([&] {
+        WorkRequest request = worker.send([&] {
+            usleep(10000);
+            didWork = true;
+        }, [&] {});
+
+        request.join();
+        EXPECT_TRUE(didWork);
+        loop.stop();
+    });
+
+    uv_run(uv_default_loop(), UV_RUN_DEFAULT);
+}
+
+TEST(Worker, WorkRequestJoinCancelsAfter) {
+    RunLoop loop(uv_default_loop());
+
+    Worker worker(1);
+    bool didAfter = false;
+
+    loop.invoke([&] {
+        WorkRequest request = worker.send([&] {
+        }, [&] {
+            didAfter = true;
+        });
+
+        request.join();
+        loop.stop();
+    });
+
+    uv_run(uv_default_loop(), UV_RUN_DEFAULT);
+    EXPECT_FALSE(didAfter);
+}

test/test.gypi

@@ -53,6 +53,7 @@
         'miscellaneous/text_conversions.cpp',
         'miscellaneous/tile.cpp',
         'miscellaneous/variant.cpp',
+        'miscellaneous/worker.cpp',
 
         'storage/storage.hpp',
         'storage/storage.cpp',