1717
1818#include < algorithm>
1919#include < atomic>
20+ #include < condition_variable>
2021#include < cstdint>
2122#include < limits>
2223#include < cstdio>
2930#include < thread>
3031#include < vector>
3132
33+ #if defined(__unix__) || defined(__APPLE__)
34+ #include < pthread.h>
35+ #endif
36+
3237#define PK_ABI_VERSION 3
3338#define PK_PARITY_REVISION 5
3439#define PK_MAX_CARS 4
@@ -307,6 +312,11 @@ struct PkBatchResult {
307312// suite asserts this stops growing with the number of batches rather than trusting the reuse.
308313std::atomic<int64_t > g_scratch_constructions{0 };
309314
315+ // Counts worker threads started for the life of the library, for pk_worker_thread_count(). Thread
316+ // creation measured 7.5us each and a benchmark plan used to start 201,545 of them - 1.5s of a 26s
317+ // plan, which is what the pool exists to stop paying. The suite asserts this stops growing.
318+ std::atomic<int64_t > g_worker_threads_started{0 };
319+
310320struct PkScratch {
311321 std::vector<int32_t > charge_member, export_member;
312322 std::vector<int32_t > clipped_start, clipped_end;
@@ -1342,6 +1352,170 @@ static void run_batch_job(const ContextStore *c, const PkBatchJob &job, PkBatchR
13421352 out.status = pk_run_one (c, &scenario, &out.result , job.soc_range_start_step , job.soc_range_end_step , &out.soc_range_min , &out.soc_range_max , scratch);
13431353}
13441354
1355+
1356+ // A pool of worker threads that outlives every batch and every context.
1357+ //
1358+ // pk_run_batch used to build and join a fresh std::vector<std::thread> per call. Measured over one
1359+ // benchmark plan that was 201,545 thread creations costing 1.5s - against 3.1s of simulation, so the
1360+ // creation cost was cancelling most of what the parallelism won and threading measured 1.6% overall.
1361+ // Parking workers on a condition variable replaces each creation with a wake.
1362+ //
1363+ // Workers park until `generation` changes, then take lane `index` of a stride partition. The calling
1364+ // thread takes lane 0 itself, so a six-job batch wakes five workers rather than six and the caller
1365+ // does useful work instead of blocking. Nobody holds the mutex while simulating; run_batch_job only
1366+ // reads the shared const context and writes its own result slot, which is what made the previous
1367+ // per-call threading safe and is unchanged here.
1368+ class PkThreadPool {
1369+ public:
1370+ // Start workers until at least `want` exist. Grows, never shrinks. Returns how many there are.
1371+ int32_t reserve_workers (int32_t want)
1372+ {
1373+ std::lock_guard<std::mutex> lock (mutex);
1374+ while (static_cast <int32_t >(workers.size ()) < want) {
1375+ const int32_t lane = static_cast <int32_t >(workers.size ()) + 1 ; // lane 0 is the caller
1376+ std::unique_ptr<Worker> worker (new Worker ());
1377+ try {
1378+ worker->thread = std::thread (&PkThreadPool::worker_loop, this , lane, worker.get ());
1379+ } catch (const std::system_error &) {
1380+ // Out of thread resources - keep what we have and let the caller use fewer lanes.
1381+ break ;
1382+ }
1383+ workers.push_back (std::move (worker));
1384+ g_worker_threads_started.fetch_add (1 , std::memory_order_relaxed);
1385+ }
1386+ return static_cast <int32_t >(workers.size ());
1387+ }
1388+
1389+ // Run jobs[0..n_jobs) across `use` lanes, the calling thread taking lane 0. `use` must be at
1390+ // least 1 and no more than worker_count() + 1.
1391+ void run (const ContextStore *c, const PkBatchJob *jobs, int32_t n_jobs, PkBatchResult *results, int32_t use)
1392+ {
1393+ {
1394+ std::lock_guard<std::mutex> lock (mutex);
1395+ cur_ctx = c;
1396+ cur_jobs = jobs;
1397+ cur_results = results;
1398+ cur_n_jobs = n_jobs;
1399+ cur_use = use;
1400+ outstanding = use - 1 ;
1401+ // Each worker has its own flag and its own condition variable, so exactly the lanes this
1402+ // batch needs are woken. Broadcasting instead would wake every worker in the pool - with
1403+ // a median batch of six jobs against sixteen workers that was ~250,000 wakeups per plan
1404+ // spent entirely on threads that immediately went back to sleep.
1405+ for (int32_t lane = 1 ; lane < use; lane++) {
1406+ workers[lane - 1 ]->has_work = true ;
1407+ }
1408+ }
1409+ for (int32_t lane = 1 ; lane < use; lane++) {
1410+ workers[lane - 1 ]->ready .notify_one ();
1411+ }
1412+
1413+ PkScratch &scratch = thread_scratch ();
1414+ for (int32_t i = 0 ; i < n_jobs; i += use) {
1415+ run_batch_job (c, jobs[i], results[i], scratch);
1416+ }
1417+
1418+ std::unique_lock<std::mutex> lock (mutex);
1419+ work_done.wait (lock, [this ] { return outstanding == 0 ; });
1420+ }
1421+
1422+ private:
1423+ // One parked worker. The flag is written under the pool mutex before its condition variable is
1424+ // signalled, so a wake cannot be lost even if the worker has not re-parked yet.
1425+ struct Worker {
1426+ std::condition_variable ready;
1427+ bool has_work = false ;
1428+ std::thread thread;
1429+ };
1430+
1431+ void worker_loop (int32_t lane, Worker *self)
1432+ {
1433+ std::unique_lock<std::mutex> lock (mutex);
1434+ for (;;) {
1435+ self->ready .wait (lock, [self] { return self->has_work ; });
1436+ self->has_work = false ;
1437+ const ContextStore *c = cur_ctx;
1438+ const PkBatchJob *jobs = cur_jobs;
1439+ PkBatchResult *results = cur_results;
1440+ const int32_t n_jobs = cur_n_jobs;
1441+ const int32_t use = cur_use;
1442+ lock.unlock ();
1443+
1444+ PkScratch &scratch = thread_scratch ();
1445+ for (int32_t i = lane; i < n_jobs; i += use) {
1446+ run_batch_job (c, jobs[i], results[i], scratch);
1447+ }
1448+
1449+ lock.lock ();
1450+ if (--outstanding == 0 ) {
1451+ work_done.notify_one ();
1452+ }
1453+ }
1454+ }
1455+
1456+ std::mutex mutex;
1457+ std::condition_variable work_done; // the caller waits here
1458+ const ContextStore *cur_ctx = nullptr ;
1459+ const PkBatchJob *cur_jobs = nullptr ;
1460+ PkBatchResult *cur_results = nullptr ;
1461+ int32_t cur_n_jobs = 0 ;
1462+ int32_t cur_use = 1 ;
1463+ int32_t outstanding = 0 ;
1464+ std::vector<std::unique_ptr<Worker>> workers;
1465+ };
1466+
1467+ // The pool is created on first threaded use and deliberately never destroyed: a static destructor
1468+ // running while CPython tears the process down risks a deadlock for no benefit, and the OS reclaims
1469+ // the threads at exit.
1470+ std::atomic<PkThreadPool *> g_pool{nullptr };
1471+ std::mutex g_pool_create_mutex;
1472+ // Held for a whole batch. Python can enter pk_run_batch from two threads at once because ctypes
1473+ // releases the GIL, and one shared pool cannot serve two batches concurrently - its published job
1474+ // state would race. Nothing in predbat does this today; serialising is the safe answer if it ever does.
1475+ std::mutex g_dispatch_mutex;
1476+
1477+ #if defined(__unix__) || defined(__APPLE__)
1478+ // Threads do not survive fork(), so a child would inherit a pool whose workers do not exist and hang
1479+ // waiting for them. The prepare/parent pair keep the dispatch mutex consistent across the fork, and
1480+ // the child drops the pool so its next batch builds a fresh one. Nothing here allocates or takes a
1481+ // new lock, which is what the child handler is allowed to do; the dead pool leaks in the child.
1482+ void pk_atfork_prepare ()
1483+ {
1484+ g_dispatch_mutex.lock ();
1485+ }
1486+
1487+ void pk_atfork_parent ()
1488+ {
1489+ g_dispatch_mutex.unlock ();
1490+ }
1491+
1492+ void pk_atfork_child ()
1493+ {
1494+ g_dispatch_mutex.unlock ();
1495+ g_pool.store (nullptr , std::memory_order_relaxed);
1496+ }
1497+ #endif
1498+
1499+ // Fetch the pool, creating it on first use.
1500+ static PkThreadPool *pool_singleton ()
1501+ {
1502+ PkThreadPool *pool = g_pool.load (std::memory_order_acquire);
1503+ if (pool) {
1504+ return pool;
1505+ }
1506+ std::lock_guard<std::mutex> lock (g_pool_create_mutex);
1507+ pool = g_pool.load (std::memory_order_relaxed);
1508+ if (!pool) {
1509+ #if defined(__unix__) || defined(__APPLE__)
1510+ static std::once_flag atfork_once;
1511+ std::call_once (atfork_once, [] { pthread_atfork (pk_atfork_prepare, pk_atfork_parent, pk_atfork_child); });
1512+ #endif
1513+ pool = new PkThreadPool ();
1514+ g_pool.store (pool, std::memory_order_release);
1515+ }
1516+ return pool;
1517+ }
1518+
13451519// Run n_jobs scenarios against one context in a single call.
13461520//
13471521// The context lookup and the Python/C boundary are paid once for a whole fan-out rather than once
@@ -1363,20 +1537,20 @@ int32_t pk_run_batch(int64_t handle, const PkBatchJob *jobs, int32_t n_jobs, PkB
13631537 // Scenarios are independent - each reads the shared const context and writes only its own result
13641538 // slot - so they are split across threads by stride.
13651539 if (n_threads > 1 && n_jobs > 1 ) {
1366- const int32_t use = n_threads < n_jobs ? n_threads : n_jobs;
1367- std::vector<std::thread> pool;
1368- pool.reserve (use);
1369- for (int32_t t = 0 ; t < use; t++) {
1370- pool.emplace_back ([c, jobs, results, n_jobs, use, t]() {
1371- for (int32_t i = t; i < n_jobs; i += use) {
1372- run_batch_job (c, jobs[i], results[i], thread_scratch ());
1373- }
1374- });
1540+ std::lock_guard<std::mutex> dispatch (g_dispatch_mutex);
1541+ int32_t use = n_threads < n_jobs ? n_threads : n_jobs;
1542+ PkThreadPool *pool = pool_singleton ();
1543+ // Clamp to the lanes that actually exist: reserve_workers stops early if the system refuses a
1544+ // thread, and a batch must never wait on a worker that was never started.
1545+ const int32_t lanes = pool->reserve_workers (use - 1 ) + 1 ;
1546+ if (use > lanes) {
1547+ use = lanes;
13751548 }
1376- for (auto &th : pool) {
1377- th.join ();
1549+ if (use > 1 ) {
1550+ pool->run (c, jobs, n_jobs, results, use);
1551+ return 0 ;
13781552 }
1379- return 0 ;
1553+ // No workers available - fall through and run the batch inline.
13801554 }
13811555 for (int32_t i = 0 ; i < n_jobs; i++) {
13821556 run_batch_job (c, jobs[i], results[i], thread_scratch ());
@@ -1396,6 +1570,17 @@ int64_t pk_scratch_construct_count(void)
13961570}
13971571
13981572
1573+ // Test hook: how many worker threads have been started since the library was loaded.
1574+ //
1575+ // Batches used to build and join a fresh set of threads on every call. The suite asserts this count
1576+ // stops growing once the pool is warm, because a thread creation costs several microseconds against
1577+ // a median batch of six jobs - it was cancelling the parallelism it was there to provide.
1578+ int64_t pk_worker_thread_count (void )
1579+ {
1580+ return g_worker_threads_started.load (std::memory_order_relaxed);
1581+ }
1582+
1583+
13991584// Test hook: sweep SoC densely and confirm the precomputed bucket boundaries give exactly the same
14001585// percent as the round_py path they replace. Returns the number of disagreements (0 = equivalent).
14011586int32_t pk_verify_soc_percent_table (double soc_max, int32_t samples)
0 commit comments