Fix incorrect reporting of task completion

Fixes Amanieu#37

Author: Amanieu

src/scheduler.cpp

@@ -57,16 +57,17 @@ void aligned_free(void* addr) LIBASYNC_NOEXCEPT
 static void generic_wait_handler(task_wait_handle wait_task)
 {
 	// Create an event to wait on
-	task_wait_event thread_event;
+	task_wait_event event;
+	event.init();
 
 	// Create a continuation for the task we are waiting for
-	wait_task.on_finish([&thread_event] {
+	wait_task.on_finish([&event] {
 		// Just signal the thread event
-		thread_event.signal(wait_type::task_finished);
+		event.signal(wait_type::task_finished);
 	});
 
 	// Wait for the event to be set
-	thread_event.wait();
+	event.wait();
 }
 
 #if defined(EMULATE_PTHREAD_THREAD_LOCAL)

src/task_wait_event.h

@@ -32,31 +32,61 @@ enum wait_type {
 
 // OS-supported event object which can be used to wait for either a task to
 // finish or for the scheduler to have more work for the current thread.
+//
+// The event object is lazily initialized to avoid unnecessary API calls.
 class task_wait_event {
-	std::mutex m;
-	std::condition_variable c;
+	std::aligned_storage<sizeof(std::mutex), std::alignment_of<std::mutex>::value>::type m;
+	std::aligned_storage<sizeof(std::condition_variable), std::alignment_of<std::condition_variable>::value>::type c;
 	int event_mask;
+	bool initialized;
+
+	std::mutex& mutex()
+	{
+		return *reinterpret_cast<std::mutex*>(&m);
+	}
+	std::condition_variable& cond()
+	{
+		return *reinterpret_cast<std::condition_variable*>(&c);
+	}
 
 public:
 	task_wait_event()
-		: event_mask(0) {}
+		: event_mask(0), initialized(false) {}
+
+	~task_wait_event()
+	{
+		if (initialized) {
+			mutex().~mutex();
+			cond().~condition_variable();
+		}
+	}
+
+	// Initialize the event, must be done before any other functions are called.
+	void init()
+	{
+		if (!initialized) {
+			new(&m) std::mutex;
+			new(&c) std::condition_variable;
+			initialized = true;
+		}
+	}
 
 	// Wait for an event to occur. Returns the event(s) that occurred. This also
 	// clears any pending events afterwards.
 	int wait()
 	{
-		std::unique_lock<std::mutex> lock(m);
+		std::unique_lock<std::mutex> lock(mutex());
 		while (event_mask == 0)
-			c.wait(lock);
+			cond().wait(lock);
 		int result = event_mask;
 		event_mask = 0;
-		return result != 0;
+		return result;
 	}
 
 	// Check if a specific event is ready
 	bool try_wait(int event)
 	{
-		std::lock_guard<std::mutex> lock(m);
+		std::lock_guard<std::mutex> lock(mutex());
 		int result = event_mask & event;
 		event_mask &= ~event;
 		return result != 0;
@@ -65,10 +95,10 @@ class task_wait_event {
 	// Signal an event and wake up a sleeping thread
 	void signal(int event)
 	{
-		std::unique_lock<std::mutex> lock(m);
+		std::unique_lock<std::mutex> lock(mutex());
 		event_mask |= event;
 		lock.unlock();
-		c.notify_one();
+		cond().notify_one();
 	}
 };
 

src/threadpool_scheduler.cpp

@@ -37,7 +37,6 @@ namespace detail {
 struct LIBASYNC_CACHELINE_ALIGN thread_data_t {
 	work_steal_queue queue;
 	std::minstd_rand rng;
-	task_wait_event event;
 	std::thread handle;
 };
 
@@ -180,11 +179,15 @@ static void thread_task_loop(threadpool_data* impl, std::size_t thread_id, task_
 	// Flag indicating if we have added a continuation to the task
 	bool added_continuation = false;
 
+	// Event to wait on
+	task_wait_event event;
+
 	// Loop while waiting for the task to complete
 	while (true) {
 		// Check if the task has finished. If we have added a continuation, we
-		// need to make sure the event has been signaled.
-		if (wait_task && (added_continuation ? current_thread.event.try_wait(wait_type::task_finished) : wait_task.ready()))
+		// need to make sure the event has been signaled, otherwise the other
+		// thread may try to signal it after we have freed it.
+		if (wait_task && (added_continuation ? event.try_wait(wait_type::task_finished) : wait_task.ready()))
 			return;
 
 		// Try to get a task from the local queue
@@ -220,11 +223,13 @@ static void thread_task_loop(threadpool_data* impl, std::size_t thread_id, task_
 				return;
 			}
 
+			// Initialize the event object
+			event.init();
+
 			// No tasks found, so sleep until something happens.
 			// If a continuation has not been added yet, add it.
 			if (wait_task && !added_continuation) {
 				// Create a continuation for the task we are waiting for
-				task_wait_event& event = current_thread.event;
 				wait_task.on_finish([&event] {
 					// Signal the thread's event
 					event.signal(wait_type::task_finished);
@@ -234,27 +239,29 @@ static void thread_task_loop(threadpool_data* impl, std::size_t thread_id, task_
 
 			// Add our thread to the list of waiting threads
 			size_t num_waiters_val = impl->num_waiters.load(std::memory_order_relaxed);
-			impl->waiters[num_waiters_val] = &current_thread.event;
+			impl->waiters[num_waiters_val] = &event;
 			impl->num_waiters.store(num_waiters_val + 1, std::memory_order_relaxed);
 
 			// Wait for our event to be signaled when a task is scheduled or
 			// the task we are waiting for has completed.
 			locked.unlock();
-			int events = current_thread.event.wait();
+			int events = event.wait();
 			locked.lock();
 
 			// Remove our thread from the list of waiting threads
 			num_waiters_val = impl->num_waiters.load(std::memory_order_relaxed);
 			for (std::size_t i = 0; i < num_waiters_val; i++) {
-				if (impl->waiters[i] == &current_thread.event) {
+				if (impl->waiters[i] == &event) {
 					if (i != num_waiters_val - 1)
 						std::swap(impl->waiters[i], impl->waiters[num_waiters_val - 1]);
 					impl->num_waiters.store(num_waiters_val - 1, std::memory_order_relaxed);
 					break;
 				}
 			}
 
-			// Check again if the task has finished
+			// Check again if the task has finished. We have added a
+			// continuation at this point, so we need to check that the
+			// continuation has finished signaling the event.
 			if (wait_task && (events & wait_type::task_finished))
 				return;
 		}