Add physical data scan tracking to resource groups

Add physical data scan tracking to resource groups. This is an optional feature in resource groups that tracks data scan via quota based tracking.

Author: kekwan

core/trino-main/src/main/java/io/trino/execution/resourcegroups/InternalResourceGroup.java

@@ -101,6 +101,10 @@ public class InternalResourceGroup
     @GuardedBy("root")
     private long cpuQuotaGenerationMillisPerSecond = Long.MAX_VALUE;
     @GuardedBy("root")
+    private long hardPhysicalDataScanLimitBytes = Long.MAX_VALUE;
+    @GuardedBy("root")
+    private long physicalDataScanQuotaGenerationBytesPerSecond = Long.MAX_VALUE;
+    @GuardedBy("root")
     private int schedulingWeight = DEFAULT_WEIGHT;
     @GuardedBy("root")
     private SchedulingPolicy schedulingPolicy = FAIR;
@@ -128,9 +132,9 @@ public class InternalResourceGroup
     private int descendantRunningQueries;
     @GuardedBy("root")
     private int descendantQueuedQueries;
-    // CPU and memory usage is cached because it changes very rapidly while queries are running, and would be expensive to track continuously
+    // CPU, memory and physical data input usage is cached because it changes very rapidly while queries are running, and would be expensive to track continuously
     @GuardedBy("root")
-    private ResourceUsage cachedResourceUsage = new ResourceUsage(0, 0);
+    private ResourceUsage cachedResourceUsage = new ResourceUsage(0, 0, 0);
     @GuardedBy("root")
     private long lastStartMillis;
     private final CounterStat timeBetweenStartsSec = new CounterStat();
@@ -168,9 +172,11 @@ public ResourceGroupInfo getFullInfo()
                     succinctBytes(softMemoryLimitBytes),
                     softConcurrencyLimit,
                     hardConcurrencyLimit,
+                    succinctBytes(hardPhysicalDataScanLimitBytes),
                     maxQueuedQueries,
                     succinctBytes(cachedResourceUsage.getMemoryUsageBytes()),
                     succinctDuration(cachedResourceUsage.getCpuUsageMillis(), MILLISECONDS),
+                    succinctBytes(cachedResourceUsage.getPhysicalInputDataUsageBytes()),
                     getQueuedQueries(),
                     getRunningQueries(),
                     eligibleSubGroups.size(),
@@ -193,9 +199,11 @@ public ResourceGroupInfo getInfo()
                     succinctBytes(softMemoryLimitBytes),
                     softConcurrencyLimit,
                     hardConcurrencyLimit,
+                    succinctBytes(hardPhysicalDataScanLimitBytes),
                     maxQueuedQueries,
                     succinctBytes(cachedResourceUsage.getMemoryUsageBytes()),
                     succinctDuration(cachedResourceUsage.getCpuUsageMillis(), MILLISECONDS),
+                    succinctBytes(cachedResourceUsage.getPhysicalInputDataUsageBytes()),
                     getQueuedQueries(),
                     getRunningQueries(),
                     eligibleSubGroups.size(),
@@ -218,9 +226,11 @@ private ResourceGroupInfo getSummaryInfo()
                     succinctBytes(softMemoryLimitBytes),
                     softConcurrencyLimit,
                     hardConcurrencyLimit,
+                    succinctBytes(hardPhysicalDataScanLimitBytes),
                     maxQueuedQueries,
                     succinctBytes(cachedResourceUsage.getMemoryUsageBytes()),
                     succinctDuration(cachedResourceUsage.getCpuUsageMillis(), MILLISECONDS),
+                    succinctBytes(cachedResourceUsage.getPhysicalInputDataUsageBytes()),
                     getQueuedQueries(),
                     getRunningQueries(),
                     eligibleSubGroups.size(),
@@ -342,6 +352,12 @@ public long getMemoryUsageBytes()
         return getResourceUsageSnapshot().getMemoryUsageBytes();
     }
 
+    @Managed
+    public long getPhysicalInputDataUsageBytes()
+    {
+        return getResourceUsageSnapshot().getPhysicalInputDataUsageBytes();
+    }
+
     @Managed
     @Override
     public long getSoftMemoryLimitBytes()
@@ -443,6 +459,45 @@ public void setCpuQuotaGenerationMillisPerSecond(long rate)
         }
     }
 
+    @Managed
+    @Override
+    public long getHardPhysicalDataScanLimitBytes()
+    {
+        synchronized (root) {
+            return hardPhysicalDataScanLimitBytes;
+        }
+    }
+
+    @Override
+    public void setHardPhysicalDataScanLimitBytes(long limit)
+    {
+        synchronized (root) {
+            boolean oldCanRun = canRunMore();
+            this.hardPhysicalDataScanLimitBytes = limit;
+            if (canRunMore() != oldCanRun) {
+                updateEligibility();
+            }
+        }
+    }
+
+    @Managed
+    @Override
+    public long getPhysicalDataScanQuotaGenerationBytesPerSecond()
+    {
+        synchronized (root) {
+            return physicalDataScanQuotaGenerationBytesPerSecond;
+        }
+    }
+
+    @Override
+    public void setPhysicalDataScanQuotaGenerationBytesPerSecond(long rate)
+    {
+        checkArgument(rate > 0, "Physical data scan quota generation must be positive");
+        synchronized (root) {
+            physicalDataScanQuotaGenerationBytesPerSecond = rate;
+        }
+    }
+
     @Managed
     @Override
     public int getSoftConcurrencyLimit()
@@ -731,7 +786,7 @@ private void startInBackground(ManagedQueryExecution query)
     {
         checkState(Thread.holdsLock(root), "Must hold lock to start a query");
         synchronized (root) {
-            runningQueries.put(query, new ResourceUsage(0, 0));
+            runningQueries.put(query, new ResourceUsage(0, 0, 0));
             InternalResourceGroup group = this;
             group.getStartedQueries().update(1);
             while (group.parent.isPresent()) {
@@ -757,11 +812,11 @@ public void updateGroupsAndProcessQueuedQueries()
         }
     }
 
-    public void generateCpuQuota(long elapsedSeconds)
+    public void generateQuotas(long elapsedSeconds)
     {
         synchronized (root) {
             if (elapsedSeconds > 0) {
-                internalGenerateCpuQuota(elapsedSeconds);
+                internalGenerateQuotas(elapsedSeconds);
             }
         }
     }
@@ -784,12 +839,12 @@ private void queryFinished(ManagedQueryExecution query)
 
             // The query is present in runningQueries
             if (lastUsage != null) {
-                // CPU is measured cumulatively (i.e. total CPU used until this moment by the query). Memory is measured
-                // instantaneously (how much memory the query is using at this moment). At query completion, memory usage
-                // drops to zero.
+                // CPU and data scan are measured cumulatively (i.e. total CPU & physical input data used until this moment by the query).
+                // Memory is measured instantaneously (how much memory the query is using at this moment). At query completion, memory usage drops to zero.
                 ResourceUsage finalUsage = new ResourceUsage(
                         query.getTotalCpuTime().toMillis(),
-                        0L);
+                        0L,
+                        query.getBasicQueryInfo().getQueryStats().getPhysicalInputDataSize().toBytes());
                 ResourceUsage delta = finalUsage.subtract(lastUsage);
 
                 runningQueries.remove(query);
@@ -827,15 +882,16 @@ private ResourceUsage updateResourceUsageAndGetDelta()
     {
         checkState(Thread.holdsLock(root), "Must hold lock to refresh stats");
         synchronized (root) {
-            ResourceUsage groupUsageDelta = new ResourceUsage(0, 0);
+            ResourceUsage groupUsageDelta = new ResourceUsage(0, 0, 0);
 
             for (Map.Entry<ManagedQueryExecution, ResourceUsage> entry : runningQueries.entrySet()) {
                 ManagedQueryExecution query = entry.getKey();
                 ResourceUsage oldResourceUsage = entry.getValue();
 
                 ResourceUsage newResourceUsage = new ResourceUsage(
                         query.getTotalCpuTime().toMillis(),
-                        query.getTotalMemoryReservation().toBytes());
+                        query.getTotalMemoryReservation().toBytes(),
+                        query.getBasicQueryInfo().getQueryStats().getPhysicalInputDataSize().toBytes());
 
                 // Compute delta and update usage
                 ResourceUsage queryUsageDelta = newResourceUsage.subtract(oldResourceUsage);
@@ -850,7 +906,7 @@ private ResourceUsage updateResourceUsageAndGetDelta()
                 groupUsageDelta = groupUsageDelta.add(subGroupUsageDelta);
                 cachedResourceUsage = cachedResourceUsage.add(subGroupUsageDelta);
 
-                if (!subGroupUsageDelta.equals(new ResourceUsage(0, 0))) {
+                if (!subGroupUsageDelta.equals(new ResourceUsage(0, 0, 0))) {
                     subGroup.updateEligibility();
                 }
             }
@@ -859,32 +915,40 @@ private ResourceUsage updateResourceUsageAndGetDelta()
         }
     }
 
-    private void internalGenerateCpuQuota(long elapsedSeconds)
+    private void internalGenerateQuotas(long elapsedSeconds)
     {
-        checkState(Thread.holdsLock(root), "Must hold lock to generate cpu quota");
+        checkState(Thread.holdsLock(root), "Must hold lock to generate quotas");
         synchronized (root) {
-            long newQuota = saturatedMultiply(elapsedSeconds, cpuQuotaGenerationMillisPerSecond);
-
-            long oldUsageMillis = cachedResourceUsage.getCpuUsageMillis();
-            long newCpuUsageMillis = saturatedSubtract(oldUsageMillis, newQuota);
+            long oldCpuUsageMillis = cachedResourceUsage.getCpuUsageMillis();
+            long oldDataScanBytes = cachedResourceUsage.getPhysicalInputDataUsageBytes();
 
-            if (newCpuUsageMillis < 0 || newCpuUsageMillis == Long.MAX_VALUE) {
-                newCpuUsageMillis = 0;
-            }
-
-            cachedResourceUsage = new ResourceUsage(newCpuUsageMillis, cachedResourceUsage.getMemoryUsageBytes());
+            long newCpuUsageMillis = computeNewUsage(oldCpuUsageMillis, elapsedSeconds, cpuQuotaGenerationMillisPerSecond);
+            long newDataScanBytes = computeNewUsage(oldDataScanBytes, elapsedSeconds, physicalDataScanQuotaGenerationBytesPerSecond);
+            cachedResourceUsage = new ResourceUsage(newCpuUsageMillis, cachedResourceUsage.getMemoryUsageBytes(), newDataScanBytes);
 
-            if ((newCpuUsageMillis < hardCpuLimitMillis && oldUsageMillis >= hardCpuLimitMillis) ||
-                    (newCpuUsageMillis < softCpuLimitMillis && oldUsageMillis >= softCpuLimitMillis)) {
+            if ((newCpuUsageMillis < hardCpuLimitMillis && oldCpuUsageMillis >= hardCpuLimitMillis) ||
+                    (newCpuUsageMillis < softCpuLimitMillis && oldCpuUsageMillis >= softCpuLimitMillis) ||
+                    (newDataScanBytes < hardPhysicalDataScanLimitBytes && oldDataScanBytes >= hardPhysicalDataScanLimitBytes)) {
                 updateEligibility();
             }
 
             for (InternalResourceGroup group : subGroups.values()) {
-                group.internalGenerateCpuQuota(elapsedSeconds);
+                group.internalGenerateQuotas(elapsedSeconds);
             }
         }
     }
 
+    private static long computeNewUsage(long currentUsage, long elapsedSeconds, long generationRate)
+    {
+        long quotaToRegenerate = saturatedMultiply(elapsedSeconds, generationRate);
+        long newUsage = saturatedSubtract(currentUsage, quotaToRegenerate);
+
+        if (newUsage < 0 || newUsage == Long.MAX_VALUE) {
+            newUsage = 0;
+        }
+        return newUsage;
+    }
+
     private boolean internalStartNext()
     {
         checkState(Thread.holdsLock(root), "Must hold lock to find next query");
@@ -996,8 +1060,8 @@ private boolean canRunMore()
         synchronized (root) {
             long cpuUsageMillis = cachedResourceUsage.getCpuUsageMillis();
             long memoryUsageBytes = cachedResourceUsage.getMemoryUsageBytes();
-
-            if ((cpuUsageMillis >= hardCpuLimitMillis) || (memoryUsageBytes > softMemoryLimitBytes)) {
+            long physicalInputDataUsageBytes = cachedResourceUsage.getPhysicalInputDataUsageBytes();
+            if ((cpuUsageMillis >= hardCpuLimitMillis) || (memoryUsageBytes > softMemoryLimitBytes) || (physicalInputDataUsageBytes >= hardPhysicalDataScanLimitBytes)) {
                 return false;
             }
 

core/trino-main/src/main/java/io/trino/execution/resourcegroups/InternalResourceGroupManager.java

@@ -79,7 +79,7 @@ public final class InternalResourceGroupManager<C>
     private final ResourceGroupConfigurationManager<?> legacyManager;
     private final MBeanExporter exporter;
     private final AtomicBoolean started = new AtomicBoolean();
-    private final AtomicLong lastCpuQuotaGenerationNanos = new AtomicLong(System.nanoTime());
+    private final AtomicLong lastQuotaGenerationNanos = new AtomicLong(System.nanoTime());
     private final Map<String, ResourceGroupConfigurationManagerFactory> configurationManagerFactories = new ConcurrentHashMap<>();
     private final SecretsResolver secretsResolver;
 
@@ -202,19 +202,19 @@ public void start()
     private void refreshAndStartQueries()
     {
         long nanoTime = System.nanoTime();
-        long elapsedSeconds = NANOSECONDS.toSeconds(nanoTime - lastCpuQuotaGenerationNanos.get());
+        long elapsedSeconds = NANOSECONDS.toSeconds(nanoTime - lastQuotaGenerationNanos.get());
         if (elapsedSeconds > 0) {
-            // Only advance our clock on second boundaries to avoid calling generateCpuQuota() too frequently, and because it would be a no-op for zero seconds.
-            lastCpuQuotaGenerationNanos.addAndGet(elapsedSeconds * 1_000_000_000L);
+            // Only advance our clock on second boundaries to avoid calling generateQuotas() too frequently, and because it would be a no-op for zero seconds.
+            lastQuotaGenerationNanos.addAndGet(elapsedSeconds * 1_000_000_000L);
         }
         else if (elapsedSeconds < 0) {
             // nano time has overflowed
-            lastCpuQuotaGenerationNanos.set(nanoTime);
+            lastQuotaGenerationNanos.set(nanoTime);
         }
         for (InternalResourceGroup group : rootGroups) {
             try {
                 if (elapsedSeconds > 0) {
-                    group.generateCpuQuota(elapsedSeconds);
+                    group.generateQuotas(elapsedSeconds);
                 }
             }
             catch (RuntimeException e) {

core/trino-main/src/main/java/io/trino/execution/resourcegroups/ResourceUsage.java

@@ -25,25 +25,29 @@ final class ResourceUsage
 {
     private final long cpuUsageMillis;
     private final long memoryUsageBytes;
+    private final long physicalInputDataUsageBytes;
 
-    public ResourceUsage(long cpuUsageMillis, long memoryUsageBytes)
+    public ResourceUsage(long cpuUsageMillis, long memoryUsageBytes, long physicalInputDataUsageBytes)
     {
         this.cpuUsageMillis = cpuUsageMillis;
         this.memoryUsageBytes = memoryUsageBytes;
+        this.physicalInputDataUsageBytes = physicalInputDataUsageBytes;
     }
 
     public ResourceUsage add(ResourceUsage other)
     {
         long newCpuUsageMillis = saturatedAdd(this.cpuUsageMillis, other.cpuUsageMillis);
         long newMemoryUsageBytes = saturatedAdd(this.memoryUsageBytes, other.memoryUsageBytes);
-        return new ResourceUsage(newCpuUsageMillis, newMemoryUsageBytes);
+        long newPhysicalInputDataUsageBytes = saturatedAdd(this.physicalInputDataUsageBytes, other.physicalInputDataUsageBytes);
+        return new ResourceUsage(newCpuUsageMillis, newMemoryUsageBytes, newPhysicalInputDataUsageBytes);
     }
 
     public ResourceUsage subtract(ResourceUsage other)
     {
         long newCpuUsageMillis = saturatedSubtract(this.cpuUsageMillis, other.cpuUsageMillis);
         long newMemoryUsageBytes = saturatedSubtract(this.memoryUsageBytes, other.memoryUsageBytes);
-        return new ResourceUsage(newCpuUsageMillis, newMemoryUsageBytes);
+        long newPhysicalInputDataUsageBytes = saturatedSubtract(this.physicalInputDataUsageBytes, other.physicalInputDataUsageBytes);
+        return new ResourceUsage(newCpuUsageMillis, newMemoryUsageBytes, newPhysicalInputDataUsageBytes);
     }
 
     public long getCpuUsageMillis()
@@ -56,6 +60,11 @@ public long getMemoryUsageBytes()
         return memoryUsageBytes;
     }
 
+    public long getPhysicalInputDataUsageBytes()
+    {
+        return physicalInputDataUsageBytes;
+    }
+
     @Override
     public boolean equals(Object other)
     {
@@ -68,12 +77,13 @@ public boolean equals(Object other)
 
         ResourceUsage otherUsage = (ResourceUsage) other;
         return cpuUsageMillis == otherUsage.cpuUsageMillis
-                && memoryUsageBytes == otherUsage.memoryUsageBytes;
+                && memoryUsageBytes == otherUsage.memoryUsageBytes
+                && physicalInputDataUsageBytes == otherUsage.physicalInputDataUsageBytes;
     }
 
     @Override
     public int hashCode()
     {
-        return Objects.hash(cpuUsageMillis, memoryUsageBytes);
+        return Objects.hash(cpuUsageMillis, memoryUsageBytes, physicalInputDataUsageBytes);
     }
 }

core/trino-main/src/main/java/io/trino/server/ResourceGroupInfo.java

@@ -37,9 +37,11 @@ public record ResourceGroupInfo(
         DataSize softMemoryLimit,
         int softConcurrencyLimit,
         int hardConcurrencyLimit,
+        DataSize hardPhysicalDataScanLimit,
         int maxQueuedQueries,
         DataSize memoryUsage,
         Duration cpuUsage,
+        DataSize physicalInputDataUsage,
         int numQueuedQueries,
         int numRunningQueries,
         int numEligibleSubGroups,
@@ -54,7 +56,8 @@ public record ResourceGroupInfo(
         requireNonNull(softMemoryLimit, "softMemoryLimit is null");
         requireNonNull(memoryUsage, "memoryUsage is null");
         requireNonNull(cpuUsage, "cpuUsage is null");
-
+        requireNonNull(hardPhysicalDataScanLimit, "hardPhysicalDataScanLimit is null");
+        requireNonNull(physicalInputDataUsage, "physicalInputDataUsage is null");
         subGroups = subGroups.map(ImmutableList::copyOf);
         runningQueries = runningQueries.map(ImmutableList::copyOf);
     }