FEX-Emu--FEX/Source/Tools/LinuxEmulation/LinuxSyscalls/Utils/Threads.cpp

// SPDX-License-Identifier: MIT
#include "LinuxSyscalls/Utils/Threads.h"
#include "LinuxSyscalls/Syscalls.h"

#include <FEXCore/Core/Context.h>
#include <FEXCore/Utils/Threads.h>

#include <FEXCore/fextl/deque.h>

namespace FEX::LinuxEmulation::Threads {
  // Stack pool handling
  struct StackPoolItem {
    void *Ptr;
    size_t Size;
  };

  struct DeadStackPoolItem {
    void *Ptr;
    size_t Size;
    bool ReadyToBeReaped;
  };

  std::mutex DeadStackPoolMutex{};
  std::mutex LiveStackPoolMutex{};

  static fextl::deque<DeadStackPoolItem> DeadStackPool{};
  static fextl::deque<StackPoolItem> LiveStackPool{};

  void *AllocateStackObject() {
    std::lock_guard lk{DeadStackPoolMutex};
    // Keep the first item in the stack pool
    void *Ptr{};

    for (auto it = DeadStackPool.begin(); it != DeadStackPool.end();) {
      auto Ready = std::atomic_ref<bool>(it->ReadyToBeReaped);
      bool ReadyToBeReaped = Ready.load();
      if (Ptr == nullptr && ReadyToBeReaped) {
        Ptr = it->Ptr;
        it = DeadStackPool.erase(it);
        continue;
      }

      if (ReadyToBeReaped) {
        FEXCore::Allocator::munmap(it->Ptr, it->Size);
        it = DeadStackPool.erase(it);
        continue;
      }

      ++it;
    }

    if (Ptr == nullptr) {
      Ptr = FEXCore::Allocator::mmap(nullptr, FEX::LinuxEmulation::Threads::STACK_SIZE, PROT_READ | PROT_WRITE, MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
    }

    return Ptr;
  }

  bool *AddStackToDeadPool(void *Ptr) {
    std::lock_guard lk{DeadStackPoolMutex};
    auto &it = DeadStackPool.emplace_back(DeadStackPoolItem{Ptr, FEX::LinuxEmulation::Threads::STACK_SIZE, false});
    return &it.ReadyToBeReaped;
  }

  void AddStackToLivePool(void *Ptr) {
    std::lock_guard lk{LiveStackPoolMutex};
    LiveStackPool.emplace_back(StackPoolItem{Ptr, FEX::LinuxEmulation::Threads::STACK_SIZE});
  }

  void RemoveStackFromLivePool(void *Ptr) {
    std::lock_guard lk{LiveStackPoolMutex};
    for (auto it = LiveStackPool.begin(); it != LiveStackPool.end(); ++it) {
      if (it->Ptr == Ptr) {
        LiveStackPool.erase(it);
        return;
      }
    }
  }

  [[noreturn]]
  void DeallocateStackObjectAndExit(void *Ptr, int Status) {
    RemoveStackFromLivePool(Ptr);
    auto ReadyToBeReaped = AddStackToDeadPool(Ptr);
    *ReadyToBeReaped = true;

#ifdef _M_ARM_64
  __asm volatile(
    "mov x8, %[SyscallNum];"
    "mov w0, %w[Result];"
    "svc #0;"
    :: [SyscallNum] "i" (SYSCALL_DEF(exit))
    , [Result] "r" (Status)
    : "memory", "x0", "x8");
#else
  __asm volatile(
    "mov %[Result], %%edi;"
    "syscall;"
    :: "a" (SYSCALL_DEF(exit))
    , [Result] "r" (Status)
    : "memory", "rdi");
#endif
    FEX_UNREACHABLE;
  }

#ifdef _M_ARM_64
  __attribute__((naked))
  void StackPivotAndCall(void *Arg, FEXCore::Threads::ThreadFunc Func, uint64_t StackPivot) {
    // x0: Arg
    // x1: Function to call
    // x2: StackPivot
    __asm volatile(R"(
    // Stack pivot.
    mov x3, sp;
    mov sp, x2;

    // Store stack storage location on to current stack
    stp x3, lr, [sp, -16]!;

    // x0 already has argument to pass.
    blr x1

    // Reload stack storage location
    ldp x2, lr, [sp], 16;

    // Stack pivot back
    mov sp, x2;

    ret;
    )"
    ::: "memory");
  }
#else
  __attribute__((naked))
  void StackPivotAndCall(void *Arg, FEXCore::Threads::ThreadFunc Func, uint64_t StackPivot) {
    // rdi: Arg
    // rsi: Function to call
    // rdx: StackPivot
    __asm volatile(R"(
    // Copy original stack in to RSP.
    movq %%rsp, %%rcx;

    // Store original stack on new stack
    pushq %%rcx;

    // Store stack pivot on new stack.
    pushq %%rdx;

    // rdi already contains function argument.
    callq *%%rsi;

    // Restore original stack
    popq %%rsp;

    ret;

    )" ::: "memory");
  }
#endif
  namespace PThreads {
    [[noreturn]]
    void *InitializeThread(void *Ptr);

    class PThread final : public FEXCore::Threads::Thread {
      public:
      PThread(FEXCore::Threads::ThreadFunc Func, void *Arg)
        : UserFunc {Func}
        , UserArg {Arg} {
        pthread_attr_t Attr{};
        Stack = AllocateStackObject();
        // pthreads allocates its dtv region behind our back and there is nothing we can do about it.
        FEXCore::Allocator::YesIKnowImNotSupposedToUseTheGlibcAllocator glibc;
        AddStackToLivePool(Stack);
        pthread_attr_init(&Attr);
        // Allocate a minimum size stack through pthreads, then stack pivot to FEX's allocated stack.
        // This is required due to a race condition with pthread's DTV/TLS regions when a stack is reused before pthreads deletes that thread's
        // DTV/TLS regions.
        // This can be seen as a crash when running Steam fairly easily, but is very confusing when debugging.
        // The cause of this race condition is from glibc associating a DTV/TLS region with a stack region until the kernel clears the
        // `set_tid_address` address construct. If the stack is reused before the address is set to zero, then glibc won't initialize the new thread's
        // DTV/TLS region, resulting in TLS usage crashing.
        pthread_attr_setstacksize(&Attr, PTHREAD_STACK_MIN);
        pthread_create(&Thread, &Attr, InitializeThread, this);
        pthread_attr_destroy(&Attr);
      }

      bool joinable() override {
        pthread_attr_t Attr{};
        if (pthread_getattr_np(Thread, &Attr) == 0) {
          int AttachState{};
          if (pthread_attr_getdetachstate(&Attr, &AttachState) == 0) {
            if (AttachState == PTHREAD_CREATE_JOINABLE) {
              return true;
            }
          }
        }
        return false;
      }

      bool join(void **ret) override {
        return pthread_join(Thread, ret) == 0;
      }

      bool detach() override {
        return pthread_detach(Thread) == 0;
      }

      bool IsSelf() override {
        auto self = pthread_self();
        return self == Thread;
      }

      FEXCore::Threads::ThreadFunc GetUserFunc() const {
        return UserFunc;
      }

      void *GetUserArg() const {
        return UserArg;
      }

      void *GetPivotStack() const {
        return Stack;
      }

      private:
      pthread_t Thread;
      FEXCore::Threads::ThreadFunc UserFunc;
      void *UserArg;
      void *Stack{};
    };

    [[noreturn]]
    void *InitializeThread(void *Ptr) {
      void *StackBase{};
      {
        PThread *Thread{reinterpret_cast<PThread*>(Ptr)};
        StackBase = Thread->GetPivotStack();

        // Run the user function.
        // `Thread` object is dead after this function returns.
        StackPivotAndCall(Thread->GetUserArg(), Thread->GetUserFunc(), reinterpret_cast<uint64_t>(StackBase) + FEX::LinuxEmulation::Threads::STACK_SIZE);
      }

      // TLS/DTV teardown is something FEX can't control. Disable glibc checking when we leave a pthread.
      FEXCore::Allocator::YesIKnowImNotSupposedToUseTheGlibcAllocator::HardDisable();

      DeallocateStackObjectAndExit(StackBase, 0);

      FEX_UNREACHABLE;
    }

    fextl::unique_ptr<FEXCore::Threads::Thread> CreateThread_PThread(
      FEXCore::Threads::ThreadFunc Func,
      void* Arg) {
      return fextl::make_unique<PThread>(Func, Arg);
    }

    void CleanupAfterFork_PThread() {
      // We don't need to pull the mutex here
      // After a fork we are the only thread running
      // Just need to make sure not to delete our own stack
      uintptr_t StackLocation = reinterpret_cast<uintptr_t>(alloca(0));

      auto ClearStackPool = [&](auto &StackPool) {
        for (auto it = StackPool.begin(); it != StackPool.end(); ) {
          auto &Item = *it;
          uintptr_t ItemStack = reinterpret_cast<uintptr_t>(Item.Ptr);
          if (ItemStack <= StackLocation && (ItemStack + Item.Size) > StackLocation) {
            // This is our stack item, skip it
            ++it;
          }
          else {
            // Untracked stack. Clean it up
            FEXCore::Allocator::munmap(Item.Ptr, Item.Size);
            it = StackPool.erase(it);
          }
        }
      };

      // Clear both dead stacks and live stacks
      ClearStackPool(DeadStackPool);
      ClearStackPool(LiveStackPool);

      LogMan::Throw::AFmt((DeadStackPool.size() + LiveStackPool.size()) <= 1,
                          "After fork we should only have zero or one tracked stacks!");
    }
  };

  void SetupThreadHandlers() {
    FEXCore::Threads::Pointers Ptrs = {
      .CreateThread = PThreads::CreateThread_PThread,
      .CleanupAfterFork = PThreads::CleanupAfterFork_PThread,
    };

    FEXCore::Threads::Thread::SetInternalPointers(Ptrs);
  }

  void Shutdown() {
    std::lock_guard lk{DeadStackPoolMutex};
    std::lock_guard lk2{LiveStackPoolMutex};
    // Erase all the dead stack pools
    for (auto &Item : DeadStackPool) {
      FEXCore::Allocator::munmap(Item.Ptr, Item.Size);
    }

    // Now clean up any that are considered to still be live
    // We are in shutdown phase, everything in the process is dead
    for (auto &Item : LiveStackPool) {
      FEXCore::Allocator::munmap(Item.Ptr, Item.Size);
    }

    DeadStackPool.clear();
    LiveStackPool.clear();
  }
}