Optimize IndexCount to terminate slaves early when limit exceeded

thorlcr/activities/indexread/thindexread.cpp

@@ -373,6 +373,8 @@ class CIndexCountActivityMaster : public CIndexReadBase
     typedef CIndexReadBase PARENT;
 
     IHThorIndexCountArg *helper;
+    mptag_t stopTag = TAG_NULL;
+    rowcount_t choosenLimit = RCMAX;
 
     void processKeyedLimit()
     {
@@ -385,10 +387,53 @@ class CIndexCountActivityMaster : public CIndexReadBase
             }
         }
     }
+    rowcount_t aggregateToLimit()
+    {
+        rowcount_t total = 0;
+        unsigned slaves = container.queryJob().querySlaves();
+        unsigned s;
+        bool sentStop = false;
+        ICommunicator &comm = queryJobChannel().queryJobComm();
+
+        for (s=0; s<slaves; s++)
+        {
+            CMessageBuffer msg;
+            rank_t sender;
+            if (!receiveMsg(msg, RANK_ALL, mpTag, &sender))
+                return 0;
+            if (abortSoon)
+                return 0;
+            rowcount_t count;
+            msg.read(count);
+            total += count;
+
+            // If limit exceeded and haven't sent stop signal yet, signal all slaves to stop
+            // This optimization applies whenever there's a choosenLimit set (including IndexExists case where choosenLimit==1)
+            if (!sentStop && choosenLimit != RCMAX && total > choosenLimit && stopTag != TAG_NULL)
+            {
+                sentStop = true;
+                CMessageBuffer stopMsg;
+                stopMsg.append(true); // stop flag
+                for (unsigned i=0; i<slaves; i++)
+                {
+                    comm.send(stopMsg, i+1, stopTag);
+                }
+            }
+        }
+        return total;
+    }
 public:
     CIndexCountActivityMaster(CMasterGraphElement *info) : CIndexReadBase(info)
     {
         helper = (IHThorIndexCountArg *)queryHelper();
+        if (!container.queryLocalOrGrouped())
+            stopTag = container.queryJob().allocateMPTag();
+    }
+    virtual void serializeSlaveData(MemoryBuffer &dst, unsigned slave) override
+    {
+        CIndexReadBase::serializeSlaveData(dst, slave);
+        if (!container.queryLocalOrGrouped())
+            dst.append(stopTag);
     }
     virtual void process() override
     {
@@ -399,12 +444,21 @@ class CIndexCountActivityMaster : public CIndexReadBase
         keyedLimit = (rowcount_t)helper->getKeyedLimit();
         if (keyedLimit != RCMAX)
             processKeyedLimit();
+
+        choosenLimit = helper->getChooseNLimit();
         rowcount_t total = aggregateToLimit();
+
         CMessageBuffer msg;
         msg.append(total);
         ICommunicator &comm = queryJobChannel().queryJobComm();
         verifyex(comm.send(msg, 1, mpTag)); // send to 1st slave only
     }
+    virtual void abort() override
+    {
+        CIndexReadBase::abort();
+        if (stopTag != TAG_NULL)
+            cancelReceiveMsg(RANK_ALL, stopTag);
+    }
 };
 
 CActivityBase *createIndexCountActivityMaster(CMasterGraphElement *info)

thorlcr/activities/indexread/thindexreadslave.cpp

@@ -19,6 +19,7 @@
 #include "jfile.hpp"
 #include "jtime.hpp"
 #include "jsort.hpp"
+#include <atomic>
 
 #include "rtlkey.hpp"
 #include "jhtree.hpp"
@@ -1273,6 +1274,61 @@ class CIndexCountSlaveActivity : public CIndexReadSlaveBase
     rowcount_t preknownTotalCount = 0;
     bool totalCountKnown = false;
     bool done = false;
+    mptag_t stopTag = TAG_NULL;
+    std::atomic<bool> stopped{false};
+
+    class CStopHandler : public CSimpleInterface, implements IThreaded
+    {
+        CIndexCountSlaveActivity &activity;
+        CThreaded threaded;
+        bool running = false;
+    public:
+        CStopHandler(CIndexCountSlaveActivity &_activity) 
+            : activity(_activity), threaded("CIndexCountSlaveActivity::CStopHandler")
+        {
+        }
+        ~CStopHandler()
+        {
+            stop();
+        }
+        void start()
+        {
+            if (!running && activity.stopTag != TAG_NULL)
+            {
+                running = true;
+                threaded.init(this, false);
+            }
+        }
+        void stop()
+        {
+            if (running)
+            {
+                running = false;
+                activity.container.queryJobChannel().queryJobComm().cancel(0, activity.stopTag);
+                threaded.join();
+            }
+        }
+        virtual void threadmain() override
+        {
+            CMessageBuffer msg;
+            while (running)
+            {
+                if (activity.container.queryJobChannel().queryJobComm().recv(msg, 0, activity.stopTag, nullptr))
+                {
+                    bool stopFlag;
+                    msg.read(stopFlag);
+                    if (stopFlag)
+                    {
+                        activity.stopped.store(true, std::memory_order_release);
+                        break;
+                    }
+                    msg.clear();
+                }
+                else
+                    break; // recv failed or was cancelled
+            }
+        }
+    } stopHandler;
 
     bool checkKeyedLimit()
     {
@@ -1285,12 +1341,23 @@ class CIndexCountSlaveActivity : public CIndexReadSlaveBase
         }
         return true;
     }
+    bool checkStopped()
+    {
+        return stopped.load(std::memory_order_acquire);
+    }
 public:
-    CIndexCountSlaveActivity(CGraphElementBase *_container) : CIndexReadSlaveBase(_container)
+    CIndexCountSlaveActivity(CGraphElementBase *_container) 
+        : CIndexReadSlaveBase(_container), stopHandler(*this)
     {
         helper = static_cast <IHThorIndexCountArg *> (container.queryHelper());
         appendOutputLinked(this);
     }
+    virtual void init(MemoryBuffer &data, MemoryBuffer &slaveData) override
+    {
+        PARENT::init(data, slaveData);
+        if (!container.queryLocalOrGrouped())
+            data.read(stopTag);
+    }
     virtual void prepareManager(IKeyManager *manager) override
     {
         PARENT::prepareManager(manager);
@@ -1323,6 +1390,8 @@ class CIndexCountSlaveActivity : public CIndexReadSlaveBase
             preknownTotalCount = 0;
         }
         done = false;
+        stopped.store(false, std::memory_order_release);
+        stopHandler.start();
     }
 
 // IRowStream
@@ -1371,11 +1440,17 @@ class CIndexCountSlaveActivity : public CIndexReadSlaveBase
                                 callback.finishedRow();
                             if ((totalCount > choosenLimit))
                                 break;
+                            // Check if master signaled us to stop
+                            if (checkStopped())
+                                break;
                         }
                         if (keyManager)
                             resetManager(keyManager);
                         if ((totalCount > choosenLimit))
                             break;
+                        // Check if master signaled us to stop
+                        if (checkStopped())
+                            break;
                     }
                     if (_currentManager)
                         prepareManager(_currentManager);
@@ -1423,6 +1498,7 @@ class CIndexCountSlaveActivity : public CIndexReadSlaveBase
     }
     virtual void stop() override
     {
+        stopHandler.stop();
         if (RCMAX != keyedLimit) // NB: will not be true if nextRow() has handled
         {
             keyedLimitCount = sendGetCount(keyedProcessed);
@@ -1433,6 +1509,7 @@ class CIndexCountSlaveActivity : public CIndexReadSlaveBase
     }
     virtual void abort() override
     {
+        stopHandler.stop();
         CIndexReadSlaveBase::abort();
         cancelReceiveMsg(0, mpTag);
     }