EDK: wait sets: Implement WaitSetDispatcher::WaitSet{AddRemove}Impl().

mojo/edk/system/wait_set_dispatcher.cc

Index: mojo/edk/system/wait_set_dispatcher.cc
diff --git a/mojo/edk/system/wait_set_dispatcher.cc b/mojo/edk/system/wait_set_dispatcher.cc
index 68090641c8f4d180ba9fa14db04dc8497f960de7..3107bb6b8acdf9725ed53071dbaa1932dde6ea10 100644
--- a/mojo/edk/system/wait_set_dispatcher.cc
+++ b/mojo/edk/system/wait_set_dispatcher.cc
@@ -4,7 +4,10 @@
 
 #include "mojo/edk/system/wait_set_dispatcher.h"
 
+#include <utility>
+
 #include "base/logging.h"
+#include "mojo/edk/system/configuration.h"
 #include "mojo/edk/system/options_validation.h"
 
 using mojo::util::MutexLocker;
@@ -13,6 +16,11 @@ using mojo::util::RefPtr;
 namespace mojo {
 namespace system {
 
+WaitSetDispatcher::Entry::Entry(MojoHandleSignals signals, uint64_t cookie)
+    : signals(signals), cookie(cookie) {}
+
+WaitSetDispatcher::Entry::~Entry() {}
+
 // static
 constexpr MojoHandleRights WaitSetDispatcher::kDefaultHandleRights;
 
@@ -94,6 +102,37 @@ WaitSetDispatcher::WaitSetDispatcher() {}
 
 WaitSetDispatcher::~WaitSetDispatcher() {}
 
+void WaitSetDispatcher::CloseImplNoLock() {
+  mutex().AssertHeld();
+
+  CookieToEntryMap entries;
+  std::swap(entries_, entries);
+  possibly_triggered_head_ = nullptr;
+  possibly_triggered_tail_ = nullptr;
+  possibly_triggered_count_ = 0u;
+
+  // We want to remove the awakables outside the lock, so we have to unlock
+  // |mutex()|. Note that while unlocked, |Awake()| may get called.
+  // TODO(vtl): This is pretty terrible, but changing it would require pretty
+  // invasive changes in many other places. We really count on |Dispatcher| not
+  // doing anything interesting after calling |CloseImplNoLock()|, and since
+  // |CloseImplNoLock()| is allowed to do nothing all the lock invariants are
+  // satisfied.
+  DCHECK(is_closed_no_lock());
+  mutex().Unlock();
+
+  for (auto& p : entries) {
+    const auto& entry = p.second;
+    if (entry->dispatcher)
+      entry->dispatcher->RemoveAwakableWithContext(this, entry->cookie,
+                                                   nullptr);
+  }
+
+  // The caller of |CloseImplNoLock()| expects |mutex()| to be locked, so we
+  // have to re-lock it (even though it really should do nothing afterwards).
+  mutex().Lock();
+}
+
 RefPtr<Dispatcher>
 WaitSetDispatcher::CreateEquivalentDispatcherAndCloseImplNoLock(
     MessagePipe* /*message_pipe*/,
@@ -105,30 +144,107 @@ WaitSetDispatcher::CreateEquivalentDispatcherAndCloseImplNoLock(
 
 MojoResult WaitSetDispatcher::WaitSetAddImpl(
     UserPointer<const MojoWaitSetAddOptions> options,
-    Handle&& handle,
+    RefPtr<Dispatcher>&& dispatcher,
     MojoHandleSignals signals,
     uint64_t cookie) {
-  MutexLocker locker(&mutex());
-  if (is_closed_no_lock())
+  Entry* entry = nullptr;
+  {
+    MutexLocker locker(&mutex());
+    if (is_closed_no_lock())
+      return MOJO_RESULT_INVALID_ARGUMENT;
+    if (is_busy_)
+      return MOJO_RESULT_BUSY;
+    MojoWaitSetAddOptions validated_options;
+    MojoResult result = ValidateWaitSetAddOptions(options, &validated_options);
+    if (result != MOJO_RESULT_OK)
+      return result;
+    if (entries_.find(cookie) != entries_.end())
+      return MOJO_RESULT_ALREADY_EXISTS;
+    if (entries_.size() >= GetConfiguration().max_wait_set_num_entries)
+      return MOJO_RESULT_RESOURCE_EXHAUSTED;
+    // Note: We'll have to set the entry's dispatcher later.
+    entry = new Entry(signals, cookie);
+    entries_[cookie] = std::unique_ptr<Entry>(entry);
+
+    is_busy_ = true;
+  }
+
+  HandleSignalsState signals_state;
+  MojoResult result = dispatcher->AddAwakableUnconditional(
+      this, signals, cookie, &signals_state);
+
+  // Can't use |MutexLocker|, since we need to do some work outside the lock
+  // in some code paths.
+  mutex().Lock();
+  DCHECK(is_busy_);
+  is_busy_ = false;
+
+  // Note: We may have been closed while |mutex()| was unlocked, so we have to
+  // check again!
+  if (is_closed_no_lock()) {
+    // Warning: In this case, |entry| has been invalidated, since it was owned
+    // by |entries_|.
+    DCHECK(entries_.empty());
+    mutex().Unlock();
+    if (result == MOJO_RESULT_OK || result == MOJO_RESULT_ALREADY_EXISTS) {
+      // We have to remove ourself from the target dispatcher's awakable list.
+      dispatcher->RemoveAwakableWithContext(this, cookie, nullptr);
+    }
     return MOJO_RESULT_INVALID_ARGUMENT;
-  MojoWaitSetAddOptions validated_options;
-  MojoResult result = ValidateWaitSetAddOptions(options, &validated_options);
-  if (result != MOJO_RESULT_OK)
+  }
+
+  DCHECK(entries_.find(cookie) != entries_.end());
+  DCHECK_EQ(entries_[cookie].get(), entry);
+
+  if (result == MOJO_RESULT_ALREADY_EXISTS) {
+    // It was added, but the wait condition is already satisfied.
+    AddPossiblyTriggeredNoLock(entry, Entry::TriggerState::POSSIBLY_SATISFIED);
+  } else if (result == MOJO_RESULT_FAILED_PRECONDITION) {
+    // The condition is never-satisfiable. Leave a zombie entry (i.e., leave
+    // |dispatcher| null).
+    mutex().Unlock();
+    return MOJO_RESULT_OK;
+  } else if (result != MOJO_RESULT_OK) {
+    size_t num_erased = entries_.erase(cookie);
+    DCHECK_EQ(num_erased, 1u);
+    mutex().Unlock();
     return result;
+  }
 
-  // TODO(vtl)
-  NOTIMPLEMENTED();
-  return MOJO_RESULT_UNIMPLEMENTED;
+  // Update the entry to actually have the dispatcher.
+  entry->dispatcher = std::move(dispatcher);
+
+  mutex().Unlock();
+  return MOJO_RESULT_OK;
 }
 
 MojoResult WaitSetDispatcher::WaitSetRemoveImpl(uint64_t cookie) {
-  MutexLocker locker(&mutex());
-  if (is_closed_no_lock())
-    return MOJO_RESULT_INVALID_ARGUMENT;
-
-  // TODO(vtl)
-  NOTIMPLEMENTED();
-  return MOJO_RESULT_UNIMPLEMENTED;
+  RefPtr<Dispatcher> dispatcher;
+  {
+    MutexLocker locker(&mutex());
+    if (is_closed_no_lock())
+      return MOJO_RESULT_INVALID_ARGUMENT;
+    if (is_busy_)
+      return MOJO_RESULT_BUSY;
+    auto it = entries_.find(cookie);
+    if (it == entries_.end())
+      return MOJO_RESULT_NOT_FOUND;
+
+    Entry* entry = it->second.get();
+    // We'll remove ourself from the target dispatcher's awakable list outside
+    // the lock.
+    dispatcher = std::move(entry->dispatcher);
+
+    if (entry->trigger_state != Entry::TriggerState::NOT_TRIGGERED)
+      RemovePossiblyTriggeredNoLock(entry);
+
+    // Note: This invalidates |entry|.
+    entries_.erase(it);
+  }
+
+  if (dispatcher)
+    dispatcher->RemoveAwakableWithContext(this, cookie, nullptr);
+  return MOJO_RESULT_OK;
 }
 
 MojoResult WaitSetDispatcher::WaitSetWaitImpl(
@@ -145,5 +261,117 @@ MojoResult WaitSetDispatcher::WaitSetWaitImpl(
   return MOJO_RESULT_UNIMPLEMENTED;
 }
 
+bool WaitSetDispatcher::Awake(MojoResult result, uint64_t context) {
+  MutexLocker locker(&mutex());
+
+  if (is_closed_no_lock()) {
+    // See |CloseImplNoLock()|: This case may occur while we're unlocked in
+    // |CloseImplNoLock()| (after that, we will have been removed from all the
+    // awakable lists, so |Awake()| should no longer be called). We may as well
+    // return false here, which will automatically remove ourselves from the
+    // awakable list (|CloseImplNoLock()| will call
+    // |RemoveAwakableWithContext()| anyway, but that's OK).
+    return false;
+  }
+
+  auto it = entries_.find(context);
+  DCHECK(it != entries_.end());
+  const auto& entry = it->second;
+  switch (result) {
+    case MOJO_RESULT_OK:
+      if (entry->trigger_state == Entry::TriggerState::NOT_TRIGGERED) {
+        AddPossiblyTriggeredNoLock(entry.get(),
+                                   Entry::TriggerState::POSSIBLY_SATISFIED);
+      }
+      return true;
+    case MOJO_RESULT_CANCELLED:
+      if (entry->trigger_state == Entry::TriggerState::NOT_TRIGGERED) {
+        AddPossiblyTriggeredNoLock(entry.get(), Entry::TriggerState::CLOSED);
+      } else {
+        // We should only ever get at most one "closed".
+        DCHECK_NE(static_cast<int>(entry->trigger_state),
+                  static_cast<int>(Entry::TriggerState::CLOSED));
+        entry->trigger_state = Entry::TriggerState::CLOSED;
+      }
+      entry->dispatcher = nullptr;
+      return false;
+    case MOJO_RESULT_FAILED_PRECONDITION:
+      // Never satisfiable.
+      if (entry->trigger_state == Entry::TriggerState::NOT_TRIGGERED) {
+        AddPossiblyTriggeredNoLock(entry.get(),
+                                   Entry::TriggerState::NEVER_SATISFIABLE);
+      } else {
+        if (entry->trigger_state == Entry::TriggerState::POSSIBLY_SATISFIED) {
+          entry->trigger_state = Entry::TriggerState::NEVER_SATISFIABLE;
+        } else {
+          // It's possible to get repeated "never satisfiable" triggers, but we
+          // shouldn't get anything after "closed".
+          DCHECK_NE(static_cast<int>(entry->trigger_state),
+                    static_cast<int>(Entry::TriggerState::CLOSED));
+        }
+      }
+      // Due to some action on some other thread, it may become satisfiable
+      // again, so continue to be awoken.
+      return true;
+    default:
+      NOTREACHED();
+      break;
+  }
+  return false;
+}
+
+void WaitSetDispatcher::AddPossiblyTriggeredNoLock(
+    Entry* entry,
+    Entry::TriggerState new_trigger_state) {
+  DCHECK_EQ(static_cast<int>(entry->trigger_state),
+            static_cast<int>(Entry::TriggerState::NOT_TRIGGERED));
+  DCHECK(!entry->possibly_triggered_previous);
+  DCHECK(!entry->possibly_triggered_next);
+  DCHECK_NE(static_cast<int>(new_trigger_state),
+            static_cast<int>(Entry::TriggerState::NOT_TRIGGERED));
+
+  entry->trigger_state = new_trigger_state;
+  possibly_triggered_count_++;
+
+  if (!possibly_triggered_tail_) {
+    DCHECK(!possibly_triggered_head_);
+    possibly_triggered_head_ = entry;
+    possibly_triggered_tail_ = entry;
+    return;
+  }
+
+  Entry* old_tail = possibly_triggered_tail_;
+  entry->possibly_triggered_previous = old_tail;
+  DCHECK(!old_tail->possibly_triggered_next);
+  old_tail->possibly_triggered_next = entry;
+  possibly_triggered_tail_ = entry;
+}
+
+void WaitSetDispatcher::RemovePossiblyTriggeredNoLock(Entry* entry) {
+  DCHECK_NE(static_cast<int>(entry->trigger_state),
+            static_cast<int>(Entry::TriggerState::NOT_TRIGGERED));
+  entry->trigger_state = Entry::TriggerState::NOT_TRIGGERED;
+  possibly_triggered_count_--;
+
+  if (!entry->possibly_triggered_previous) {
+    DCHECK_EQ(entry, possibly_triggered_head_);
+    possibly_triggered_head_ = entry->possibly_triggered_next;
+  } else {
+    entry->possibly_triggered_previous->possibly_triggered_next =
+        entry->possibly_triggered_next;
+  }
+
+  if (!entry->possibly_triggered_next) {
+    DCHECK_EQ(entry, possibly_triggered_tail_);
+    possibly_triggered_tail_ = entry->possibly_triggered_previous;
+  } else {
+    entry->possibly_triggered_next->possibly_triggered_previous =
+        entry->possibly_triggered_previous;
+  }
+
+  entry->possibly_triggered_previous = nullptr;
+  entry->possibly_triggered_next = nullptr;
+}
+
 }  // namespace system
 }  // namespace mojo