Merge pull request #37748: Where a Throwable is being caught, take care to rethrow OutOfMemoryError

Author: scwhittle

runners/core-java/src/main/java/org/apache/beam/runners/core/SimpleDoFnRunner.java

@@ -176,9 +176,8 @@ public void startBundle() {
     // This can contain user code. Wrap it in case it throws an exception.
     try {
       invoker.invokeStartBundle(new DoFnStartBundleArgumentProvider());
-    } catch (Throwable t) {
-      // Exception in user code.
-      throw wrapUserCodeException(t);
+    } catch (Exception e) {
+      throw wrapUserCodeException(e);
     }
   }
 
@@ -215,8 +214,8 @@ private void invokeProcessElement(WindowedValue<InputT> elem) {
     // This can contain user code. Wrap it in case it throws an exception.
     try {
       invoker.invokeProcessElement(new DoFnProcessContext(elem));
-    } catch (Exception ex) {
-      throw wrapUserCodeException(ex);
+    } catch (Exception e) {
+      throw wrapUserCodeException(e);
     }
   }
 
@@ -225,9 +224,8 @@ public void finishBundle() {
     // This can contain user code. Wrap it in case it throws an exception.
     try {
       invoker.invokeFinishBundle(new DoFnFinishBundleArgumentProvider());
-    } catch (Throwable t) {
-      // Exception in user code.
-      throw wrapUserCodeException(t);
+    } catch (Exception e) {
+      throw wrapUserCodeException(e);
     }
   }
 
@@ -237,8 +235,8 @@ public <KeyT> void onWindowExpiration(BoundedWindow window, Instant timestamp, K
         new OnWindowExpirationArgumentProvider<>(window, timestamp, key));
   }
 
-  private RuntimeException wrapUserCodeException(Throwable t) {
-    throw UserCodeException.wrapIf(!isSystemDoFn(), t);
+  private RuntimeException wrapUserCodeException(Exception e) {
+    throw UserCodeException.wrapIf(!isSystemDoFn(), e);
   }
 
   private boolean isSystemDoFn() {

runners/direct-java/src/main/java/org/apache/beam/runners/direct/DirectRunner.java

@@ -223,10 +223,9 @@ public DirectPipelineResult run(Pipeline pipeline) {
           result.waitUntilFinish();
         } catch (UserCodeException userException) {
           throw new PipelineExecutionException(userException.getCause());
+        } catch (RuntimeException | OutOfMemoryError e) {
+          throw e;
         } catch (Throwable t) {
-          if (t instanceof RuntimeException) {
-            throw (RuntimeException) t;
-          }
           throw new RuntimeException(t);
         }
       }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/BatchDataflowWorker.java

@@ -276,11 +276,11 @@ boolean doWork(WorkItem workItem, WorkItemStatusClient workItemStatusClient) thr
       executeWork(worker, progressUpdater);
       workItemStatusClient.reportSuccess();
       return true;
-
+    } catch (OutOfMemoryError oom) {
+      throw oom;
     } catch (Throwable e) {
       workItemStatusClient.reportError(e);
       return false;
-
     } finally {
       if (worker != null) {
         try {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowFnRunner.java

@@ -92,11 +92,9 @@ private void invokeProcessElement(WindowedValue<InputT> elem) {
     // This can contain user code. Wrap it in case it throws an exception.
     try {
       fn.processElement(elem.getValue(), options, stepContext, sideInputReader, outputManager);
+    } catch (RuntimeException ex) {
+      throw ex;
     } catch (Exception ex) {
-      if (ex instanceof RuntimeException) {
-        throw (RuntimeException) ex;
-      }
-
       throw new RuntimeException(ex);
     }
   }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/GroupAlsoByWindowsParDoFn.java

@@ -172,6 +172,8 @@ private DoFnRunner<InputT, KV<K, Iterable<V>>> createRunner() {
         output -> {
           try {
             receiver.process(output);
+          } catch (OutOfMemoryError oom) {
+            throw oom;
           } catch (Throwable t) {
             throw new RuntimeException(t);
           }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/WorkerCustomSources.java

@@ -680,6 +680,8 @@ static ApproximateReportedProgress getReaderProgress(BoundedSource.BoundedReader
         if (fractionConsumed != null) {
           progress.setFractionConsumed(fractionConsumed);
         }
+      } catch (OutOfMemoryError oom) {
+        throw oom;
       } catch (Throwable t) {
         LOG.warn("Error estimating fraction consumed from reader {}", reader, t);
       }
@@ -689,6 +691,8 @@ static ApproximateReportedProgress getReaderProgress(BoundedSource.BoundedReader
         if (parallelism != null) {
           progress.setConsumedParallelism(parallelism);
         }
+      } catch (OutOfMemoryError oom) {
+        throw oom;
       } catch (Throwable t) {
         LOG.warn("Error estimating consumed parallelism from reader {}", reader, t);
       }
@@ -698,6 +702,8 @@ static ApproximateReportedProgress getReaderProgress(BoundedSource.BoundedReader
         if (parallelism != null) {
           progress.setRemainingParallelism(parallelism);
         }
+      } catch (OutOfMemoryError oom) {
+        throw oom;
       } catch (Throwable t) {
         LOG.warn("Error estimating remaining parallelism from reader {}", reader, t);
       }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/logging/DataflowWorkerLoggingHandler.java

@@ -377,7 +377,7 @@ private void reportFailure(String message, Exception e, int code) {
       if (manager != null) {
         manager.error(message, e, code);
       }
-    } catch (Throwable t) {
+    } catch (Exception ex) {
       // Failed to report logging failure. No meaningful action left.
     }
   }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/JfrInterop.java

@@ -59,6 +59,8 @@ class JfrInterop {
       recordingClose = jfrClass.getMethod("close");
 
       jfrConfig = configClass.getMethod("getConfiguration", String.class).invoke(null, "profile");
+    } catch (OutOfMemoryError oom) {
+      throw oom;
     } catch (Throwable t) {
       throw new RuntimeException(t);
     }

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/MemoryMonitor.java

@@ -529,7 +529,7 @@ private void updateData(long now, long lastTimeWokeUp) {
 
     // Clearing this list should "release" some memory that will be needed to dump the heap.
     // We could try to reallocate it again if we later notice memory pressure has subsided,
-    // but that is risk. Further, leaving this released may help with the memory pressure.
+    // but that is risky. Further, leaving this released may help with the memory pressure.
     reservedForDumpingHeap = null;
 
     try {

runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/util/common/worker/BatchingShuffleEntryReader.java

@@ -129,7 +129,7 @@ private void fillEntries() {
         nextStartPosition = batch.nextStartPosition;
         entries = batch.entries.listIterator();
         currentBatch = batch;
-      } catch (RuntimeException e) {
+      } catch (RuntimeException | OutOfMemoryError e) {
         throw e;
       } catch (Throwable t) {
         throw new RuntimeException(t);