Core: Fix exception handling in BaseTaskWriter (apache#5683)

* Core: Fix exception handling in BaseTaskWriter.

* Fix state check.

aws/src/main/java/org/apache/iceberg/aws/s3/S3OutputStream.java

@@ -105,7 +105,6 @@ class S3OutputStream extends PositionOutputStream {
 
   private long pos = 0;
   private boolean closed = false;
-  private Throwable closeFailureException;
 
   @SuppressWarnings("StaticAssignmentInConstructor")
   S3OutputStream(S3Client s3, S3URI location, AwsProperties awsProperties, MetricsContext metrics)
@@ -258,15 +257,6 @@ private void newStream() throws IOException {
 
   @Override
   public void close() throws IOException {
-
-    // A failed s3 close removes state that is required for a successful close.
-    // Any future close on this stream should fail.
-    if (closeFailureException != null) {
-      throw new IOException(
-          "Attempted to close an S3 output stream that failed to close earlier",
-          closeFailureException);
-    }
-
     if (closed) {
       return;
     }
@@ -277,9 +267,6 @@ public void close() throws IOException {
     try {
       stream.close();
       completeUploads();
-    } catch (Exception e) {
-      closeFailureException = e;
-      throw e;
     } finally {
       cleanUpStagingFiles();
     }

aws/src/test/java/org/apache/iceberg/aws/s3/TestS3OutputStream.java

@@ -183,7 +183,7 @@ public void testWriteWithChecksumEnabled() {
   }
 
   @Test
-  public void testCloseFailureShouldPersistOnFutureClose() throws IOException {
+  public void testDoubleClose() throws IOException {
     IllegalStateException mockException =
         new IllegalStateException("mock failure to completeUploads on close");
     Mockito.doThrow(mockException)
@@ -195,9 +195,7 @@ public void testCloseFailureShouldPersistOnFutureClose() throws IOException {
         .isInstanceOf(mockException.getClass())
         .hasMessageContaining(mockException.getMessage());
 
-    Assertions.assertThatThrownBy(stream::close)
-        .isInstanceOf(IOException.class)
-        .hasCause(mockException);
+    Assertions.assertThatNoException().isThrownBy(stream::close);
   }
 
   private void writeTest() {

core/src/main/java/org/apache/iceberg/io/BaseTaskWriter.java

@@ -52,6 +52,7 @@ public abstract class BaseTaskWriter<T> implements TaskWriter<T> {
   private final OutputFileFactory fileFactory;
   private final FileIO io;
   private final long targetFileSize;
+  private Throwable failure;
 
   protected BaseTaskWriter(
       PartitionSpec spec,
@@ -72,6 +73,12 @@ protected PartitionSpec spec() {
     return spec;
   }
 
+  protected void setFailure(Throwable throwable) {
+    if (failure == null) {
+      this.failure = throwable;
+    }
+  }
+
   @Override
   public void abort() throws IOException {
     close();
@@ -88,6 +95,8 @@ public void abort() throws IOException {
   public WriteResult complete() throws IOException {
     close();
 
+    Preconditions.checkState(failure == null, "Cannot return results from failed writer", failure);
+
     return WriteResult.builder()
         .addDataFiles(completedDataFiles)
         .addDeleteFiles(completedDeleteFiles)
@@ -180,28 +189,43 @@ public void deleteKey(T key) throws IOException {
 
     @Override
     public void close() throws IOException {
-      // Close data writer and add completed data files.
-      if (dataWriter != null) {
-        dataWriter.close();
-        dataWriter = null;
-      }
+      try {
+        // Close data writer and add completed data files.
+        if (dataWriter != null) {
+          try {
+            dataWriter.close();
+          } finally {
+            dataWriter = null;
+          }
+        }
 
-      // Close eq-delete writer and add completed equality-delete files.
-      if (eqDeleteWriter != null) {
-        eqDeleteWriter.close();
-        eqDeleteWriter = null;
-      }
+        // Close eq-delete writer and add completed equality-delete files.
+        if (eqDeleteWriter != null) {
+          try {
+            eqDeleteWriter.close();
+          } finally {
+            eqDeleteWriter = null;
+          }
+        }
 
-      if (insertedRowMap != null) {
-        insertedRowMap.clear();
-        insertedRowMap = null;
-      }
+        if (insertedRowMap != null) {
+          insertedRowMap.clear();
+          insertedRowMap = null;
+        }
 
-      // Add the completed pos-delete files.
-      if (posDeleteWriter != null) {
-        completedDeleteFiles.addAll(posDeleteWriter.complete());
-        referencedDataFiles.addAll(posDeleteWriter.referencedDataFiles());
-        posDeleteWriter = null;
+        // Add the completed pos-delete files.
+        if (posDeleteWriter != null) {
+          try {
+            // complete will call close
+            completedDeleteFiles.addAll(posDeleteWriter.complete());
+            referencedDataFiles.addAll(posDeleteWriter.referencedDataFiles());
+          } finally {
+            posDeleteWriter = null;
+          }
+        }
+      } catch (IOException | RuntimeException e) {
+        setFailure(e);
+        throw e;
       }
     }
   }
@@ -286,22 +310,29 @@ private boolean shouldRollToNewFile() {
 
     private void closeCurrent() throws IOException {
       if (currentWriter != null) {
-        currentWriter.close();
-
-        if (currentRows == 0L) {
-          try {
-            io.deleteFile(currentFile.encryptingOutputFile());
-          } catch (UncheckedIOException e) {
-            // the file may not have been created, and it isn't worth failing the job to clean up,
-            // skip deleting
+        try {
+          currentWriter.close();
+
+          if (currentRows == 0L) {
+            try {
+              io.deleteFile(currentFile.encryptingOutputFile());
+            } catch (UncheckedIOException e) {
+              // the file may not have been created, and it isn't worth failing the job to clean up,
+              // skip deleting
+            }
+          } else {
+            complete(currentWriter);
           }
-        } else {
-          complete(currentWriter);
-        }
 
-        this.currentFile = null;
-        this.currentWriter = null;
-        this.currentRows = 0;
+        } catch (IOException | RuntimeException e) {
+          setFailure(e);
+          throw e;
+
+        } finally {
+          this.currentFile = null;
+          this.currentWriter = null;
+          this.currentRows = 0;
+        }
       }
     }
 

core/src/main/java/org/apache/iceberg/io/SortedPosDeleteWriter.java

@@ -52,6 +52,7 @@ class SortedPosDeleteWriter<T> implements FileWriter<PositionDelete<T>, DeleteWr
 
   private int records = 0;
   private boolean closed = false;
+  private Throwable failure;
 
   SortedPosDeleteWriter(
       FileAppenderFactory<T> appenderFactory,
@@ -74,6 +75,12 @@ class SortedPosDeleteWriter<T> implements FileWriter<PositionDelete<T>, DeleteWr
     this(appenderFactory, fileFactory, format, partition, DEFAULT_RECORDS_NUM_THRESHOLD);
   }
 
+  protected void setFailure(Throwable throwable) {
+    if (failure == null) {
+      this.failure = throwable;
+    }
+  }
+
   @Override
   public long length() {
     throw new UnsupportedOperationException(
@@ -109,6 +116,8 @@ public void delete(CharSequence path, long pos, T row) {
   public List<DeleteFile> complete() throws IOException {
     close();
 
+    Preconditions.checkState(failure == null, "Cannot return results from failed writer", failure);
+
     return completedFiles;
   }
 
@@ -119,8 +128,8 @@ public CharSequenceSet referencedDataFiles() {
   @Override
   public void close() throws IOException {
     if (!closed) {
-      flushDeletes();
       this.closed = true;
+      flushDeletes();
     }
   }
 
@@ -161,6 +170,7 @@ private void flushDeletes() {
         positions.forEach(posRow -> closeableWriter.delete(path, posRow.pos(), posRow.row()));
       }
     } catch (IOException e) {
+      setFailure(e);
       throw new UncheckedIOException(
           "Failed to write the sorted path/pos pairs to pos-delete file: "
               + outputFile.encryptingOutputFile().location(),