From 864b77e91acc7064ce6ee6b06e1b497c990c79de Mon Sep 17 00:00:00 2001
From: Gian Merlino <gianmerlino@gmail.com>
Date: Sat, 9 Jul 2022 06:45:38 -0700
Subject: [PATCH] SpillingGrouper: Make DISK_FULL sticky. (#12764)

When we return DISK_FULL to a processing thread, it skips the rest of
the segment and the query is canceled. However, it's possible that the
next segment starts processing before cancellation can kick in. We want
that one, if it occurs, to see DISK_FULL too.
---
 .../query/groupby/epinephelinae/SpillingGrouper.java     | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/SpillingGrouper.java b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/SpillingGrouper.java
index 86365fef12b..dd038acbd18 100644
--- a/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/SpillingGrouper.java
+++ b/processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/SpillingGrouper.java
@@ -79,6 +79,7 @@ public class SpillingGrouper<KeyType> implements Grouper<KeyType>
   private final List<File> dictionaryFiles = new ArrayList<>();
   private final boolean sortHasNonGroupingFields;
 
+  private boolean diskFull = false;
   private boolean spillingAllowed;
 
   public SpillingGrouper(
@@ -171,6 +172,13 @@ public class SpillingGrouper<KeyType> implements Grouper<KeyType>
   @Override
   public AggregateResult aggregate(KeyType key, int keyHash)
   {
+    if (diskFull) {
+      // If the prior return was DISK_FULL, then return it again. When we return DISK_FULL to a processing thread,
+      // it skips the rest of the segment and the query is canceled. However, it's possible that the next segment
+      // starts processing before cancellation can kick in. We want that one, if it occurs, to see DISK_FULL too.
+      return DISK_FULL;
+    }
+
     final AggregateResult result = grouper.aggregate(key, keyHash);
 
     if (result.isOk() || !spillingAllowed || temporaryStorage.maxSize() <= 0) {
@@ -184,6 +192,7 @@ public class SpillingGrouper<KeyType> implements Grouper<KeyType>
         spill();
       }
       catch (TemporaryStorageFullException e) {
+        diskFull = true;
         return DISK_FULL;
       }
       catch (IOException e) {