From e4d97d109f6483d168fd5925aa6648a721cef0ad Mon Sep 17 00:00:00 2001
From: Karol Sobczak <karol.sobczak@karolsobczak.com>
Date: Fri, 8 Mar 2019 12:50:12 +0100
Subject: [PATCH] Add OperatorContext#aggregateSystemMemoryContext method

---
 .../main/java/io/prestosql/operator/OperatorContext.java    | 6 ++++++
 .../io/prestosql/memory/context/MemoryTrackingContext.java  | 5 +++++
 2 files changed, 11 insertions(+)

diff --git a/presto-main/src/main/java/io/prestosql/operator/OperatorContext.java b/presto-main/src/main/java/io/prestosql/operator/OperatorContext.java
index eef3d7e338f96..a61380ba93699 100644
--- a/presto-main/src/main/java/io/prestosql/operator/OperatorContext.java
+++ b/presto-main/src/main/java/io/prestosql/operator/OperatorContext.java
@@ -269,6 +269,12 @@ public AggregatedMemoryContext aggregateUserMemoryContext()
         return new InternalAggregatedMemoryContext(operatorMemoryContext.aggregateUserMemoryContext(), memoryFuture, this::updatePeakMemoryReservations, false);
     }
 
+    // caller shouldn't close this context as it's managed by the OperatorContext
+    public AggregatedMemoryContext aggregateSystemMemoryContext()
+    {
+        return new InternalAggregatedMemoryContext(operatorMemoryContext.aggregateSystemMemoryContext(), memoryFuture, this::updatePeakMemoryReservations, false);
+    }
+
     // caller shouldn't close this context as it's managed by the OperatorContext
     public AggregatedMemoryContext aggregateRevocableMemoryContext()
     {
diff --git a/presto-memory-context/src/main/java/io/prestosql/memory/context/MemoryTrackingContext.java b/presto-memory-context/src/main/java/io/prestosql/memory/context/MemoryTrackingContext.java
index 802c262b7fae9..4a02ccb9d7e26 100644
--- a/presto-memory-context/src/main/java/io/prestosql/memory/context/MemoryTrackingContext.java
+++ b/presto-memory-context/src/main/java/io/prestosql/memory/context/MemoryTrackingContext.java
@@ -117,6 +117,11 @@ public AggregatedMemoryContext aggregateRevocableMemoryContext()
         return revocableAggregateMemoryContext;
     }
 
+    public AggregatedMemoryContext aggregateSystemMemoryContext()
+    {
+        return systemAggregateMemoryContext;
+    }
+
     public AggregatedMemoryContext newAggregateSystemMemoryContext()
     {
         return systemAggregateMemoryContext.newAggregatedMemoryContext();