From 5956efafdd8e2afdba20f79f491d2f25c831b7c2 Mon Sep 17 00:00:00 2001
From: Tomasz Sobczyk <tomasz.sobczyk1997@gmail.com>
Date: Fri, 17 Sep 2021 18:58:30 +0200
Subject: [PATCH] Hard-kill search in generate_training_data when the node
 count is 3x over the limit.

---
 src/search.cpp | 22 ++++++++++++++++++++++
 src/thread.cpp |  2 +-
 src/thread.h   |  1 +
 3 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/src/search.cpp b/src/search.cpp
index a811deda..5fa2a5ac 100644
--- a/src/search.cpp
+++ b/src/search.cpp
@@ -577,6 +577,7 @@ namespace {
     {
         // Step 2. Check for aborted search and immediate draw
         if (   Threads.stop.load(std::memory_order_relaxed)
+            || (thisThread->maxNodes && thisThread->nodes.load(std::memory_order_relaxed) >= thisThread->maxNodes)
             || pos.is_draw(ss->ply)
             || ss->ply >= MAX_PLY)
             return (ss->ply >= MAX_PLY && !ss->inCheck) ? evaluate(pos)
@@ -1965,6 +1966,7 @@ namespace Search
 
       // Zero initialization of the number of search nodes
       th->nodes = 0;
+      th->maxNodes = 0;
 
       // Clear all history types. This initialization takes a little time, and
       // the accuracy of the search is rather low, so the good and bad are
@@ -2061,6 +2063,13 @@ namespace Search
 
   ValueAndPV search(Position& pos, int depth_, size_t multiPV /* = 1 */, uint64_t nodesLimit /* = 0 */)
   {
+    // Sometimes a depth takes extreme amount of time (in the order of x1000 or more)
+    // than the previous depth, which can cause the search bounded by nodes to go for a long time.
+    // Because of that we add an additional limit that's 10x higher and is checked within
+    // the search function and can kill the search regardless of whether the full depth
+    // has been completed or not.
+    constexpr uint64_t nodesOversearchFactor = 10;
+
     std::vector<Move> pvs;
 
     Depth depth = depth_;
@@ -2101,6 +2110,11 @@ namespace Search
     Value delta = -VALUE_INFINITE;
     Value bestValue = -VALUE_INFINITE;
 
+    if (nodesLimit != 0)
+    {
+      th->maxNodes = nodesLimit * nodesOversearchFactor;
+    }
+
     while ((rootDepth += 1) <= depth
       // exit this loop even if the node limit is exceeded
       // The number of search nodes is passed in the argument of this function.
@@ -2141,6 +2155,11 @@ namespace Search
           Depth adjustedDepth = std::max(1, rootDepth);
           bestValue = Stockfish::search<Root>(pos, ss, alpha, beta, adjustedDepth, false);
 
+          if (th->maxNodes && th->nodes.load(std::memory_order_relaxed) >= th->maxNodes)
+          {
+            break;
+          }
+
           stable_sort(rootMoves.begin() + pvIdx, rootMoves.end());
           //my_stable_sort(pos.this_thread()->thread_id(),&rootMoves[0] + pvIdx, rootMoves.size() - pvIdx);
 
@@ -2187,6 +2206,8 @@ namespace Search
     // Considering multiPV, the score of rootMoves[0] is returned as bestValue.
     bestValue = rootMoves[0].score;
 
+    th->maxNodes = 0;
+
     return ValueAndPV(bestValue, pvs);
   }
 
@@ -2764,6 +2785,7 @@ namespace Search
         th->nmpMinPly = th->bestMoveChanges = th->failedHighCnt = 0;
         th->ttHitAverage = TtHitAverageWindow * TtHitAverageResolution / 2;
         th->nodes = 0;
+        th->maxNodes = 0;
       }
 
 
diff --git a/src/thread.cpp b/src/thread.cpp
index 5051708d..5bead318 100644
--- a/src/thread.cpp
+++ b/src/thread.cpp
@@ -34,7 +34,7 @@ ThreadPool Threads; // Global object
 /// Thread constructor launches the thread and waits until it goes to sleep
 /// in idle_loop(). Note that 'searching' and 'exit' should be already set.
 
-Thread::Thread(size_t n) : idx(n), stdThread(&Thread::idle_loop, this) {
+Thread::Thread(size_t n) : idx(n), stdThread(&Thread::idle_loop, this), maxNodes(0) {
 
   wait_for_search_finished();
   wait_for_worker_finished();
diff --git a/src/thread.h b/src/thread.h
index c0218577..782e2035 100644
--- a/src/thread.h
+++ b/src/thread.h
@@ -92,6 +92,7 @@ public:
   int selDepth, nmpMinPly;
   Color nmpColor;
   std::atomic<uint64_t> nodes, tbHits, bestMoveChanges;
+  uint64_t maxNodes;
 
   Position rootPos;
   StateInfo rootState;