From 574afd486484e8393d24f362b6308ee2caa3275a Mon Sep 17 00:00:00 2001
From: rajeevrastogi <rajeevrastogi03@gmail.com>
Date: Mon, 28 Dec 2020 17:58:15 +0530
Subject: [PATCH] [I2B55F]: Fix: Consumer gets blocked if producer gets
 cancelled

---
 .../java/io/prestosql/execution/SqlStageExecution.java     | 7 ++-----
 .../java/io/prestosql/execution/StageStateMachine.java     | 5 ++++-
 2 files changed, 6 insertions(+), 6 deletions(-)

diff --git a/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java b/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java
index 6ccc1b83f..124c33ad0 100644
--- a/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java
+++ b/presto-main/src/main/java/io/prestosql/execution/SqlStageExecution.java
@@ -555,9 +555,6 @@ public final class SqlStageExecution
                 }
                 if (finishedTasks.containsAll(allTasks)) {
                     stateMachine.transitionToFinished();
-                    if (isReuseTableScanEnabled(stateMachine.getSession())) {
-                        setReuseTableScanMappingIdStatus(stateMachine);
-                    }
                 }
             }
         }
@@ -568,9 +565,9 @@ public final class SqlStageExecution
     }
 
     //Assuming there will be only one table scan in one stage
-    private static synchronized void setReuseTableScanMappingIdStatus(StageStateMachine state)
+    public static synchronized void setReuseTableScanMappingIdStatus(StageStateMachine state)
     {
-        if (state.getProducerScanNode() == null) {
+        if (!isReuseTableScanEnabled(state.getSession()) || state.getProducerScanNode() == null) {
             return;
         }
 
diff --git a/presto-main/src/main/java/io/prestosql/execution/StageStateMachine.java b/presto-main/src/main/java/io/prestosql/execution/StageStateMachine.java
index ea4c0beb8..a59cb0a0f 100644
--- a/presto-main/src/main/java/io/prestosql/execution/StageStateMachine.java
+++ b/presto-main/src/main/java/io/prestosql/execution/StageStateMachine.java
@@ -203,23 +203,26 @@ public class StageStateMachine
 
     public boolean transitionToFinished()
     {
+        SqlStageExecution.setReuseTableScanMappingIdStatus(this);
         return stageState.setIf(FINISHED, currentState -> !currentState.isDone());
     }
 
     public boolean transitionToCanceled()
     {
+        SqlStageExecution.setReuseTableScanMappingIdStatus(this);
         return stageState.setIf(CANCELED, currentState -> !currentState.isDone());
     }
 
     public boolean transitionToAborted()
     {
+        SqlStageExecution.setReuseTableScanMappingIdStatus(this);
         return stageState.setIf(ABORTED, currentState -> !currentState.isDone());
     }
 
     public boolean transitionToFailed(Throwable throwable)
     {
         requireNonNull(throwable, "throwable is null");
-
+        SqlStageExecution.setReuseTableScanMappingIdStatus(this);
         failureCause.compareAndSet(null, Failures.toFailure(throwable));
         boolean failed = stageState.setIf(FAILED, currentState -> !currentState.isDone());
         if (failed) {