text: /*
   * Copyright (c) Meta Platforms, Inc. and 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/detail/AtFork.h>
  
  #include <atomic>
  #include <mutex>
  #include <thread>
  
  #include <glog/logging.h>
  
  #include <folly/portability/GMock.h>
  #include <folly/portability/GTest.h>
  
  TEST(ThreadLocal, AtFork) {
    int foo;
    bool forked = false;
    folly::detail::AtFork::registerHandler(
        &foo,
        [&] {
          forked = true;
          return true;
        },
        [] {},
        [] {});
    auto pid = folly::kIsSanitizeThread
        ? folly::detail::AtFork::forkInstrumented(fork)
        : fork();
    if (pid) {
      int status;
      auto pid2 = wait(&status);
      EXPECT_EQ(status, 0);
      EXPECT_EQ(pid, pid2);
    } else {
      exit(0);
    }
    EXPECT_TRUE(forked);
    forked = false;
    folly::detail::AtFork::unregisterHandler(&foo);
    pid = fork();
    if (pid) {
      int status;
      auto pid2 = wait(&status);
      EXPECT_EQ(status, 0);
      EXPECT_EQ(pid, pid2);
    } else {
      exit(0);
    }
    EXPECT_FALSE(forked);
  }
  
  TEST(ThreadLocal, AtForkOrdering) {
    std::atomic<bool> done{false};
    std::atomic<bool> started{false};
    std::mutex a;
    std::mutex b;
    int foo;
    int foo2;
    folly::detail::AtFork::registerHandler(
        &foo,
        [&] { return a.try_lock(); },
        [&] { a.unlock(); },
        [&] { a.unlock(); });
    folly::detail::AtFork::registerHandler(
        &foo2,
        [&] { return b.try_lock(); },
        [&] { b.unlock(); },
        [&] { b.unlock(); });
  
    auto thr = std::thread([&]() {
      std::lock_guard<std::mutex> g(a);
      started = true;
      usleep(100);
      std::lock_guard<std::mutex> g2(b);
    });
    while (!started) {
    }
    auto pid = folly::kIsSanitizeThread
        ? folly::detail::AtFork::forkInstrumented(fork)
        : fork();
    if (pid) {
      int status;
      auto pid2 = wait(&status);
      EXPECT_TRUE(WIFEXITED(status));
      EXPECT_THAT(
          WEXITSTATUS(status),
          folly::kIsSanitizeThread //
              ? testing::AnyOfArray({0, 66})
              : testing::AnyOfArray({0}));
      EXPECT_EQ(pid, pid2);
    } else {
      exit(0);
    }
    done = true;
    thr.join();
  }
  