From 5080e1ae539e812a8bd96c0db39dbbd8a49f6224 Mon Sep 17 00:00:00 2001
From: Alexey Antipovsky <alexey.vladimirovich@mariadb.com>
Date: Wed, 16 Dec 2020 18:21:35 +0300
Subject: [PATCH] MCOL-4031 More accurate memory usage counting while sorting

---
 dbcon/joblist/limitedorderby.cpp    | 37 +++++++++++++++++++++++++----
 dbcon/joblist/limitedorderby.h      |  6 +++--
 utils/rowgroup/rowgroup.h           | 12 +++++++---
 utils/windowfunction/idborderby.cpp |  2 +-
 4 files changed, 46 insertions(+), 11 deletions(-)

diff --git a/dbcon/joblist/limitedorderby.cpp b/dbcon/joblist/limitedorderby.cpp
index f4cf096b2..82cc73957 100644
--- a/dbcon/joblist/limitedorderby.cpp
+++ b/dbcon/joblist/limitedorderby.cpp
@@ -42,9 +42,10 @@ using namespace ordering;
 namespace joblist
 {
 
+const uint64_t LimitedOrderBy::fMaxUncommited = 102400; // 100KiB - make it configurable?
 
 // LimitedOrderBy class implementation
-LimitedOrderBy::LimitedOrderBy() : fStart(0), fCount(-1)
+LimitedOrderBy::LimitedOrderBy() : fStart(0), fCount(-1), fUncommitedMemory(0)
 {
     fRule.fIdbCompare = this;
 }
@@ -122,6 +123,20 @@ void LimitedOrderBy::processRow(const rowgroup::Row& row)
         OrderByRow newRow(fRow0, fRule);
         fOrderByQueue.push(newRow);
 
+        uint64_t memSizeInc = sizeof(newRow);
+        fUncommitedMemory += memSizeInc;
+        if (fUncommitedMemory >= fMaxUncommited)
+        {
+            fMemSize += fUncommitedMemory;
+            if (!fRm->getMemory(fUncommitedMemory, fSessionMemLimit))
+            {
+                cerr << IDBErrorInfo::instance()->errorMsg(fErrorCode) << " @"
+                     << __FILE__ << ":" << __LINE__;
+                throw IDBExcept(fErrorCode);
+            }
+            fUncommitedMemory = 0;
+        }
+
         // add to the distinct map
         if (fDistinct)
             fDistinctMap->insert(fRow0.getPointer());
@@ -132,7 +147,7 @@ void LimitedOrderBy::processRow(const rowgroup::Row& row)
         if (fRowGroup.getRowCount() >= fRowsPerRG)
         {
             fDataQueue.push(fData);
-            uint64_t newSize = fRowsPerRG * fRowGroup.getRowSize();
+            uint64_t newSize = fRowGroup.getSizeWithStrings() - fRowGroup.getHeaderSize();
             fMemSize += newSize;
 
             if (!fRm->getMemory(newSize, fSessionMemLimit))
@@ -157,7 +172,7 @@ void LimitedOrderBy::processRow(const rowgroup::Row& row)
 
         if (fDistinct)
         {
-			fDistinctMap->erase(fOrderByQueue.top().fData);
+            fDistinctMap->erase(fOrderByQueue.top().fData);
             fDistinctMap->insert(row1.getPointer());
         }
 
@@ -173,6 +188,18 @@ void LimitedOrderBy::processRow(const rowgroup::Row& row)
  */ 
 void LimitedOrderBy::finalize()
 {
+    if (fUncommitedMemory > 0)
+    {
+        fMemSize += fUncommitedMemory;
+        if (!fRm->getMemory(fUncommitedMemory, fSessionMemLimit))
+        {
+            cerr << IDBErrorInfo::instance()->errorMsg(fErrorCode) << " @"
+                 << __FILE__ << ":" << __LINE__;
+            throw IDBExcept(fErrorCode);
+        }
+        fUncommitedMemory = 0;
+    }
+
     queue<RGData> tempQueue;
     if (fRowGroup.getRowCount() > 0)
         fDataQueue.push(fData);
@@ -181,7 +208,7 @@ void LimitedOrderBy::finalize()
     {
         // *DRRTUY Very memory intensive. CS needs to account active
         // memory only and release memory if needed.
-        uint64_t memSizeInc = fRowsPerRG * fRowGroup.getRowSize();
+        uint64_t memSizeInc = fRowGroup.getSizeWithStrings() - fRowGroup.getHeaderSize();
         fMemSize += memSizeInc;
 
         if (!fRm->getMemory(memSizeInc, fSessionMemLimit))
@@ -212,7 +239,7 @@ void LimitedOrderBy::finalize()
         }
 
         list<RGData>::iterator tempListIter = tempRGDataList.begin();
-        
+
         i = 0;
         uint32_t rSize = fRow0.getSize();
         uint64_t preLastRowNumb = fRowsPerRG - 1;
diff --git a/dbcon/joblist/limitedorderby.h b/dbcon/joblist/limitedorderby.h
index 62517277b..d54acf56e 100644
--- a/dbcon/joblist/limitedorderby.h
+++ b/dbcon/joblist/limitedorderby.h
@@ -61,8 +61,10 @@ public:
     void finalize();
 
 protected:
-    uint64_t                            fStart;
-    uint64_t                            fCount;
+    uint64_t              fStart;
+    uint64_t              fCount;
+    uint64_t              fUncommitedMemory;
+    static const uint64_t fMaxUncommited;
 };
 
 
diff --git a/utils/rowgroup/rowgroup.h b/utils/rowgroup/rowgroup.h
index 950cc6511..630c98c0e 100644
--- a/utils/rowgroup/rowgroup.h
+++ b/utils/rowgroup/rowgroup.h
@@ -1435,6 +1435,7 @@ public:
 
     // this returns the size of the row data with the string table
     inline uint64_t getSizeWithStrings() const;
+    inline uint64_t getSizeWithStrings(uint64_t n) const;
 
     // sets the row count to 0 and the baseRid to something
     // effectively initializing whatever chunk of memory
@@ -1708,12 +1709,17 @@ inline uint32_t RowGroup::getRowSizeWithStrings() const
     return oldOffsets[columnCount];
 }
 
-inline uint64_t RowGroup::getSizeWithStrings() const
+inline uint64_t RowGroup::getSizeWithStrings(uint64_t n) const
 {
     if (strings == NULL)
-        return getDataSize();
+        return getDataSize(n);
     else
-        return getDataSize() + strings->getSize();
+        return getDataSize(n) + strings->getSize();
+}
+
+inline uint64_t RowGroup::getSizeWithStrings() const
+{
+    return getSizeWithStrings(getRowCount());
 }
 
 inline bool RowGroup::isCharType(uint32_t colIndex) const
diff --git a/utils/windowfunction/idborderby.cpp b/utils/windowfunction/idborderby.cpp
index f87a43708..efec1b47a 100644
--- a/utils/windowfunction/idborderby.cpp
+++ b/utils/windowfunction/idborderby.cpp
@@ -772,7 +772,7 @@ void IdbOrderBy::initialize(const RowGroup& rg)
     // initialize rows
     IdbCompare::initialize(rg);
 
-    uint64_t newSize = fRowsPerRG * rg.getRowSize();
+    uint64_t newSize = rg.getSizeWithStrings(fRowsPerRG);
     fMemSize += newSize;
 
     if (!fRm->getMemory(newSize, fSessionMemLimit))