Restore threadpool num threads query behavior (#20474)

Differential Revision: D109521669

Pull Request resolved: #20474

Author: GregoryComer

extension/threadpool/threadpool.cpp

@@ -52,13 +52,14 @@ void child_atfork() {
 #endif
 
 ThreadPool::ThreadPool(size_t thread_count)
-    : threadpool_(pthreadpool_create(thread_count), pthreadpool_destroy) {}
+    : threadpool_(pthreadpool_create(thread_count), pthreadpool_destroy),
+      thread_count_(
+          threadpool_ ? pthreadpool_get_threads_count(threadpool_.get()) : 0) {}
 
 size_t ThreadPool::get_thread_count() const {
   std::lock_guard<std::mutex> lock{mutex_};
 
-  ET_CHECK_MSG(threadpool_.get(), "Invalid threadpool!");
-  return pthreadpool_get_threads_count(threadpool_.get());
+  return thread_count_;
 }
 
 bool ThreadPool::_unsafe_reset_threadpool(uint32_t new_thread_count) {
@@ -72,13 +73,16 @@ bool ThreadPool::_unsafe_reset_threadpool(uint32_t new_thread_count) {
   std::lock_guard<std::mutex> lock{mutex_};
 
   threadpool_.reset(pthreadpool_create(new_thread_count));
+  thread_count_ =
+      threadpool_ ? pthreadpool_get_threads_count(threadpool_.get()) : 0;
   return true;
 }
 
 void ThreadPool::_unsafe_destroy_threadpool() {
   std::lock_guard<std::mutex> lock{mutex_};
   ET_LOG(Info, "Destroying threadpool.");
   threadpool_.reset();
+  thread_count_ = 0;
 }
 
 void ThreadPool::run(

extension/threadpool/threadpool.h

@@ -50,6 +50,10 @@ class ThreadPool final {
   ThreadPool(ThreadPool&&) = delete;
   ThreadPool& operator=(ThreadPool&&) = delete;
 
+  /**
+   * Returns the number of threads in the threadpool. This number does not
+   * reflect temporary lowering via UseNThreadsThreadPoolGuard.
+   */
   size_t get_thread_count() const;
 
   /**
@@ -97,6 +101,8 @@ class ThreadPool final {
   // which case this mutex will be useful. Otherwise remove it.
   mutable std::mutex mutex_;
   std::unique_ptr<pthreadpool, decltype(&pthreadpool_destroy)> threadpool_;
+  // The number of threads in the threadpool.
+  size_t thread_count_;
 };
 
 /**