Add folly::get_cached_pid()

Summary:
The function folly::get_cached_pid() saves and reuses the result of calling getpid().
At fork, the value is updated for the child process.

It is thread-safe. The first call is async-signal-unsafe but subsequent calls are async-signal-safe.

Reviewed By: yfeldblum, davidtgoldblatt

Differential Revision: D23433702

fbshipit-source-id: a026300206c1726fc1c2c6a588828f2e6559c4e2

folly/system/Pid.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/system/Pid.h>
+
+#include <folly/Likely.h>
+#include <folly/detail/AtFork.h>
+
+#include <glog/logging.h>
+
+#include <atomic>
+
+namespace folly {
+
+namespace {
+
+enum class State : uint8_t { INVALID, LOCKED, VALID };
+
+class PidState {
+  std::atomic<State> value_{State::INVALID};
+
+ public:
+  FOLLY_ALWAYS_INLINE State load() noexcept {
+    return value_.load(std::memory_order_acquire);
+  }
+
+  void store(State state) noexcept {
+    value_.store(state, std::memory_order_release);
+  }
+
+  bool cas(State& expected, State newstate) noexcept {
+    return value_.compare_exchange_strong(
+        expected,
+        newstate,
+        std::memory_order_relaxed,
+        std::memory_order_relaxed);
+  }
+}; // PidState
+
+class PidCache {
+  PidState state_;
+  pid_t pid_;
+
+ public:
+  FOLLY_ALWAYS_INLINE pid_t get() {
+    DCHECK(!valid() || pid_ == getpid());
+    return valid() ? pid_ : init();
+  }
+
+ private:
+  bool valid() {
+    return state_.load() == State::VALID;
+  }
+
+  FOLLY_COLD pid_t init() {
+    pid_t pid = getpid();
+    auto s = state_.load();
+    if (s == State::INVALID && state_.cas(s, State::LOCKED)) {
+      folly::detail::AtFork::registerHandler(
+          this, [] { return true; }, [] {}, [this] { pid_ = getpid(); });
+      pid_ = pid;
+      state_.store(State::VALID);
+    }
+    return pid;
+  }
+}; // PidCache
+
+static PidCache cache_;
+
+} // namespace
+
+pid_t get_cached_pid() {
+  return cache_.get();
+}
+
+} // namespace folly

folly/system/Pid.h

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <folly/portability/Unistd.h>
+
+namespace folly {
+
+/**
+ * Calls getpid() and returns the returned value, with a thread-safe
+ * cache in front. The cache is updated in the child after fork().
+ *
+ * Thread-safety: MT-safe.
+ *
+ * Async-signal-safety: The first call is AS-unsafe but subsequent
+ * calls are AS-safe.
+ */
+pid_t get_cached_pid();
+
+} // namespace folly

folly/system/test/PidTest.cpp

+/*
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <folly/system/Pid.h>
+
+#include <folly/portability/GTest.h>
+
+#include <glog/logging.h>
+
+TEST(PidTest, basic) {
+  pid_t mypid = getpid();
+  EXPECT_EQ(mypid, folly::get_cached_pid());
+  EXPECT_EQ(getpid(), folly::get_cached_pid());
+}
+
+TEST(PidTest, fork) {
+  pid_t mypid = getpid();
+  EXPECT_EQ(mypid, folly::get_cached_pid());
+  auto pid = fork();
+  if (pid > 0) {
+    // parent
+    EXPECT_EQ(mypid, folly::get_cached_pid());
+    EXPECT_NE(pid, mypid);
+  } else if (pid == 0) {
+    // child
+    EXPECT_NE(mypid, folly::get_cached_pid()); // mypid still has parent's pid
+    mypid = getpid(); // get child's pid
+    EXPECT_EQ(mypid, folly::get_cached_pid());
+    std::exit(0); // Do not print gtest results
+  } else {
+    PLOG(FATAL) << "Failed to fork()";
+  }
+}