Make AtomicNotificationQueue::size() TSAN-clean

Summary:
taskExecuteCount_ can be incremented from a single thread, but reading it from another thread is still technically a data race. Mark it as atomic to make TSAN happy.
Also make sure we never return < 0 when called from a non-consumer thread (since we only have the synchronization guarantee on the consumer thread).

Differential Revision: D24523284

fbshipit-source-id: 6185125e132d71a83a3f29c909b7e4eab09e6e79

folly/io/async/AtomicNotificationQueue.cpp

@@ -223,9 +223,10 @@ void AtomicNotificationQueue::setMaxReadAtOnce(uint32_t maxAtOnce) {
 }
 
 int32_t AtomicNotificationQueue::size() const {
-  auto queueSize = atomicQueue_.getPushCount() - taskExecuteCount_;
-  DCHECK(queueSize >= 0);
-  return queueSize;
+  auto queueSize = atomicQueue_.getPushCount() -
+      taskExecuteCount_.load(std::memory_order_relaxed);
+  DCHECK(!evb_ || !evb_->isInEventBaseThread() || queueSize >= 0);
+  return queueSize >= 0 ? queueSize : 0;
 }
 
 bool AtomicNotificationQueue::empty() const {
@@ -285,7 +286,11 @@ bool AtomicNotificationQueue::drive() {
         break;
       }
     }
-    ++taskExecuteCount_;
+    // This is faster than fetch_add and is safe because only consumer thread
+    // writes to taskExecuteCount_.
+    taskExecuteCount_.store(
+        taskExecuteCount_.load(std::memory_order_relaxed) + 1,
+        std::memory_order_relaxed);
     std::move(queue_.front()).execute();
     queue_.pop();
   }

folly/io/async/AtomicNotificationQueue.h

@@ -305,7 +305,7 @@ class AtomicNotificationQueue : private EventBase::LoopCallback,
 
   AtomicQueue atomicQueue_;
   Queue queue_;
-  ssize_t taskExecuteCount_{0};
+  std::atomic<ssize_t> taskExecuteCount_{0};
   int32_t maxReadAtOnce_{10};
   int eventfd_{-1};
   int pipeFds_[2]{-1, -1}; // to fallback to on older/non-linux systems