Avoid disarming the list if arm is called multiple times on an empty …

…list

Summary: Avoid disarming the list if arm is called multiple times on an empty list

Reviewed By: ot

Differential Revision: D52076109

fbshipit-source-id: 3b0441cbcfe892b7b4410b1f3a96712cd9cd0a91

folly/AtomicIntrusiveLinkedList.h

@@ -123,27 +123,23 @@ class AtomicIntrusiveLinkedList {
    * Atomically replaces the head with kQueueArmedTag if it is nullptr
    * or it replaces the head with nullptr otherwise
    * Calling arm() for a queue where the head is already kQueueArmedTag
-   * will cause the head to become nullptr
+   * will return nullptr without changing the head value
    * @return previous head value
    */
 
   T* arm() {
-    auto oldHead = head_.load(std::memory_order_relaxed);
-    if (!oldHead &&
-        head_.compare_exchange_strong(
-            oldHead,
-            reinterpret_cast<T*>(kQueueArmedTag),
-            std::memory_order_release,
-            std::memory_order_relaxed)) {
-      return nullptr;
-    }
-
-    oldHead = head_.exchange(nullptr, std::memory_order_acquire);
-    if (oldHead && reinterpret_cast<intptr_t>(oldHead) != kQueueArmedTag) {
-      return oldHead;
-    }
-
-    return nullptr;
+    auto* kTagPtr = reinterpret_cast<T*>(kQueueArmedTag);
+    T* oldHead = head_.load(std::memory_order_relaxed);
+    T* newHead;
+    do {
+      newHead = (oldHead == nullptr || oldHead == kTagPtr) ? kTagPtr : nullptr;
+    } while (!head_.compare_exchange_strong(
+        oldHead,
+        newHead,
+        std::memory_order_acq_rel,
+        std::memory_order_relaxed));
+
+    return (oldHead == nullptr || oldHead == kTagPtr) ? nullptr : oldHead;
   }
 
   /**

folly/test/AtomicLinkedListTest.cpp

@@ -177,6 +177,20 @@ TEST(AtomicIntrusiveLinkedList, Arm) {
     EXPECT_TRUE(list.empty());
     EXPECT_EQ(nullptr, list.unsafeHead());
   }
+  // arm twice
+  {
+    EXPECT_EQ(nullptr, list.arm());
+    EXPECT_EQ(nullptr, list.arm());
+    EXPECT_TRUE(list.insertHeadArm(&a));
+    EXPECT_EQ(&a, list.unsafeHead());
+    auto* ret = list.arm();
+    // unlink
+    while (ret) {
+      auto* next = ret->next();
+      ret->next() = nullptr;
+      ret = next;
+    }
+  }
 
   // do not arm first
   {