Commit 8ccd13f0 authored by Yedidya Feldblum's avatar Yedidya Feldblum Committed by Facebook GitHub Bot

Use relaxed_atomic in ShutdownSocketSet

Summary:
[Folly] Use `relaxed_atomic` in `ShutdownSocketSet`.

It currently uses a slab of `std::atomic` states, every operation on which uses `std::memory_order_relaxed`. Which is what `folly::relaxed_atomic` is for.

Reviewed By: ot, magedm

Differential Revision: D6290453

fbshipit-source-id: a6b4df115b0de429dd544a97fdf97c088459a6d9
parent 1a0ac04b
...@@ -73,8 +73,8 @@ static NetworkSocket at(size_t p) { ...@@ -73,8 +73,8 @@ static NetworkSocket at(size_t p) {
ShutdownSocketSet::ShutdownSocketSet(size_t capacity) ShutdownSocketSet::ShutdownSocketSet(size_t capacity)
: capacity_(cap(capacity)), : capacity_(cap(capacity)),
data_(static_cast<std::atomic<uint8_t>*>( data_(static_cast<relaxed_atomic<uint8_t>*>(
folly::checkedCalloc(capacity_, sizeof(std::atomic<uint8_t>)))), folly::checkedCalloc(capacity_, sizeof(relaxed_atomic<uint8_t>)))),
nullFile_("/dev/null", O_RDWR) {} nullFile_("/dev/null", O_RDWR) {}
void ShutdownSocketSet::add(NetworkSocket fd) { void ShutdownSocketSet::add(NetworkSocket fd) {
...@@ -87,8 +87,7 @@ void ShutdownSocketSet::add(NetworkSocket fd) { ...@@ -87,8 +87,7 @@ void ShutdownSocketSet::add(NetworkSocket fd) {
auto& sref = data_[p]; auto& sref = data_[p];
uint8_t prevState = FREE; uint8_t prevState = FREE;
CHECK(sref.compare_exchange_strong( CHECK(sref.compare_exchange_strong(prevState, IN_USE))
prevState, IN_USE, std::memory_order_relaxed))
<< "Invalid prev state for fd " << fd << ": " << int(prevState); << "Invalid prev state for fd " << fd << ": " << int(prevState);
} }
...@@ -102,19 +101,18 @@ void ShutdownSocketSet::remove(NetworkSocket fd) { ...@@ -102,19 +101,18 @@ void ShutdownSocketSet::remove(NetworkSocket fd) {
auto& sref = data_[p]; auto& sref = data_[p];
uint8_t prevState = 0; uint8_t prevState = 0;
prevState = sref.load(std::memory_order_relaxed); prevState = sref.load();
do { do {
switch (prevState) { switch (prevState) {
case IN_SHUTDOWN: case IN_SHUTDOWN:
std::this_thread::sleep_for(std::chrono::milliseconds(1)); std::this_thread::sleep_for(std::chrono::milliseconds(1));
prevState = sref.load(std::memory_order_relaxed); prevState = sref.load();
continue; continue;
case FREE: case FREE:
LOG(FATAL) << "Invalid prev state for fd " << fd << ": " LOG(FATAL) << "Invalid prev state for fd " << fd << ": "
<< int(prevState); << int(prevState);
} }
} while ( } while (!sref.compare_exchange_weak(prevState, FREE));
!sref.compare_exchange_weak(prevState, FREE, std::memory_order_relaxed));
} }
int ShutdownSocketSet::close(NetworkSocket fd) { int ShutdownSocketSet::close(NetworkSocket fd) {
...@@ -125,7 +123,7 @@ int ShutdownSocketSet::close(NetworkSocket fd) { ...@@ -125,7 +123,7 @@ int ShutdownSocketSet::close(NetworkSocket fd) {
} }
auto& sref = data_[p]; auto& sref = data_[p];
uint8_t prevState = sref.load(std::memory_order_relaxed); uint8_t prevState = sref.load();
uint8_t newState = 0; uint8_t newState = 0;
do { do {
...@@ -141,8 +139,7 @@ int ShutdownSocketSet::close(NetworkSocket fd) { ...@@ -141,8 +139,7 @@ int ShutdownSocketSet::close(NetworkSocket fd) {
LOG(FATAL) << "Invalid prev state for fd " << fd << ": " LOG(FATAL) << "Invalid prev state for fd " << fd << ": "
<< int(prevState); << int(prevState);
} }
} while (!sref.compare_exchange_strong( } while (!sref.compare_exchange_strong(prevState, newState));
prevState, newState, std::memory_order_relaxed));
return newState == FREE ? folly::closeNoInt(fd) : 0; return newState == FREE ? folly::closeNoInt(fd) : 0;
} }
...@@ -157,16 +154,14 @@ void ShutdownSocketSet::shutdown(NetworkSocket fd, bool abortive) { ...@@ -157,16 +154,14 @@ void ShutdownSocketSet::shutdown(NetworkSocket fd, bool abortive) {
auto& sref = data_[p]; auto& sref = data_[p];
uint8_t prevState = IN_USE; uint8_t prevState = IN_USE;
if (!sref.compare_exchange_strong( if (!sref.compare_exchange_strong(prevState, IN_SHUTDOWN)) {
prevState, IN_SHUTDOWN, std::memory_order_relaxed)) {
return; return;
} }
doShutdown(fd, abortive); doShutdown(fd, abortive);
prevState = IN_SHUTDOWN; prevState = IN_SHUTDOWN;
if (sref.compare_exchange_strong( if (sref.compare_exchange_strong(prevState, SHUT_DOWN)) {
prevState, SHUT_DOWN, std::memory_order_relaxed)) {
return; return;
} }
...@@ -175,15 +170,14 @@ void ShutdownSocketSet::shutdown(NetworkSocket fd, bool abortive) { ...@@ -175,15 +170,14 @@ void ShutdownSocketSet::shutdown(NetworkSocket fd, bool abortive) {
folly::closeNoInt(fd); // ignore errors, nothing to do folly::closeNoInt(fd); // ignore errors, nothing to do
CHECK( CHECK(sref.compare_exchange_strong(prevState, FREE))
sref.compare_exchange_strong(prevState, FREE, std::memory_order_relaxed))
<< "Invalid prev state for fd " << fd << ": " << int(prevState); << "Invalid prev state for fd " << fd << ": " << int(prevState);
} }
void ShutdownSocketSet::shutdownAll(bool abortive) { void ShutdownSocketSet::shutdownAll(bool abortive) {
for (size_t p = 0; p < capacity_; ++p) { for (size_t p = 0; p < capacity_; ++p) {
auto& sref = data_[p]; auto& sref = data_[p];
if (sref.load(std::memory_order_relaxed) == IN_USE) { if (sref.load() == IN_USE) {
shutdown(at(p), abortive); shutdown(at(p), abortive);
} }
} }
......
...@@ -16,12 +16,12 @@ ...@@ -16,12 +16,12 @@
#pragma once #pragma once
#include <atomic>
#include <cstdlib> #include <cstdlib>
#include <memory> #include <memory>
#include <folly/File.h> #include <folly/File.h>
#include <folly/net/NetworkSocket.h> #include <folly/net/NetworkSocket.h>
#include <folly/synchronization/RelaxedAtomic.h>
namespace folly { namespace folly {
...@@ -139,7 +139,7 @@ class ShutdownSocketSet { ...@@ -139,7 +139,7 @@ class ShutdownSocketSet {
}; };
size_t const capacity_; size_t const capacity_;
std::unique_ptr<std::atomic<uint8_t>[], Free> data_; std::unique_ptr<relaxed_atomic<uint8_t>[], Free> data_;
folly::File nullFile_; folly::File nullFile_;
}; };
......
Markdown is supported
0%
or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment