Improve exception handling of stream readers using resource management via 'using',

Author: yruslan

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/processor/CobolProcessor.scala

@@ -21,8 +21,9 @@ import za.co.absa.cobrix.cobol.processor.impl.{CobolProcessorInPlace, CobolProce
 import za.co.absa.cobrix.cobol.reader.parameters.{CobolParametersParser, Parameters, ReaderParameters}
 import za.co.absa.cobrix.cobol.reader.schema.CobolSchema
 import za.co.absa.cobrix.cobol.reader.stream.{FSStream, SimpleStream}
+import za.co.absa.cobrix.cobol.utils.UsingUtils
 
-import java.io.{BufferedInputStream, BufferedOutputStream, FileOutputStream, OutputStream}
+import java.io.{BufferedOutputStream, FileOutputStream, OutputStream}
 import scala.collection.mutable
 
 
@@ -159,43 +160,12 @@ object CobolProcessor {
         case CobolProcessingStrategy.ToVariableLength => new CobolProcessorToRdw(readerParameters, copybook, copybookContents, options)
       }
 
-      val ifs = new FSStream(fileToProcess)
-      val ofs = new BufferedOutputStream(new FileOutputStream(outputFile))
-
-      var originalException: Throwable = null
-
-      val recordCount = try {
-        processor.process(ifs, ofs)(rawRecordProcessor)
-      } catch {
-        case ex: Throwable =>
-          originalException = ex
-          0L
-      } finally {
-        try {
-          ifs.close()
-        } catch {
-          case e: Throwable =>
-            if (originalException != null) {
-              originalException.addSuppressed(e)
-            } else {
-              originalException = e
-            }
-        }
-
-        try {
-          ofs.close()
-        } catch {
-          case e: Throwable =>
-            if (originalException != null) {
-              originalException.addSuppressed(e)
-            } else {
-              originalException = e
-            }
+      val recordCount = UsingUtils.using(new FSStream(fileToProcess)) { ifs =>
+        UsingUtils.using(new BufferedOutputStream(new FileOutputStream(outputFile))) { ofs =>
+          processor.process(ifs, ofs)(rawRecordProcessor)
         }
       }
 
-      if (originalException != null) throw originalException
-
       recordCount
     }
   }

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/reader/stream/SimpleStream.scala

@@ -18,7 +18,7 @@ package za.co.absa.cobrix.cobol.reader.stream
 
 
 /** This trait defines a simple abstraction for processing variable length record data. */
-trait SimpleStream {
+trait SimpleStream extends AutoCloseable {
   def size: Long
 
   def totalSize: Long
@@ -35,6 +35,4 @@ trait SimpleStream {
   def copyStream(): SimpleStream
 
   @throws(classOf[Exception]) def next(numberOfBytes: Int): Array[Byte]
-
-  @throws(classOf[Exception]) def close(): Unit
 }

cobol-parser/src/main/scala/za/co/absa/cobrix/cobol/utils/UsingUtils.scala

@@ -0,0 +1,62 @@
+/*
+ * Copyright 2018 ABSA Group Limited
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package za.co.absa.cobrix.cobol.utils
+
+import scala.util.control.NonFatal
+
+object UsingUtils {
+  /**
+    * Executes the given action with a resource that implements the AutoCloseable interface, ensuring
+    * proper closure of the resource. Any exception that occurs during the action or resource closure
+    * is handled appropriately, with suppressed exceptions added where relevant. Null resources are not supported.
+    *
+    * @param resource a lazily evaluated resource that implements AutoCloseable
+    * @param action   a function to be executed using the provided resource
+    * @tparam T the type of the resource, which must extend AutoCloseable
+    * @throws Throwable if either the action or resource closure fails. If both fail, the action's exception
+    *                   is thrown with the closure's exception added as suppressed
+    */
+  def using[T <: AutoCloseable,U](resource: => T)(action: T => U): U = {
+    var thrownException: Option[Throwable] = None
+    var suppressedException: Option[Throwable] = None
+    val openedResource = resource
+
+    val result = try {
+      Option(action(openedResource))
+    } catch {
+      case NonFatal(ex) =>
+        thrownException = Option(ex)
+        None
+    } finally
+      if (openedResource != null) {
+        try
+          openedResource.close()
+        catch {
+          case NonFatal(ex) => suppressedException = Option(ex)
+        }
+      }
+
+    (thrownException, suppressedException) match {
+      case (Some(thrown), Some(suppressed)) =>
+        thrown.addSuppressed(suppressed)
+        throw thrown
+      case (Some(thrown), None)     => throw thrown
+      case (None, Some(suppressed)) => throw suppressed
+      case (None, None)             => result.getOrElse(throw new IllegalArgumentException("Action returned null"))
+    }
+  }
+}

spark-cobol/src/main/scala/za/co/absa/cobrix/spark/cobol/SparkCobolProcessor.scala

@@ -28,6 +28,7 @@ import za.co.absa.cobrix.cobol.reader.index.entry.SparseIndexEntry
 import za.co.absa.cobrix.cobol.reader.parameters.CobolParametersParser.PARAM_GENERATE_RECORD_ID
 import za.co.absa.cobrix.cobol.reader.parameters.{CobolParameters, CobolParametersParser, Parameters}
 import za.co.absa.cobrix.cobol.reader.schema.CobolSchema
+import za.co.absa.cobrix.cobol.utils.UsingUtils
 import za.co.absa.cobrix.spark.cobol.reader.VarLenReader
 import za.co.absa.cobrix.spark.cobol.source.index.IndexBuilder
 import za.co.absa.cobrix.spark.cobol.source.parameters.LocalityParameters
@@ -223,7 +224,6 @@ object SparkCobolProcessor {
         indexes.flatMap(indexEntry => {
           val filePathName = filesMap(indexEntry.fileId)
           val path = new Path(filePathName)
-          val fileSystem = path.getFileSystem(sconf.value)
           val fileName = path.getName
           val numOfBytes = if (indexEntry.offsetTo > 0L) indexEntry.offsetTo - indexEntry.offsetFrom else 0L
           val numOfBytesMsg = if (numOfBytes > 0) s"${numOfBytes / Constants.megabyte} MB" else "until the end"
@@ -237,9 +237,7 @@ object SparkCobolProcessor {
 
       case _ =>
         spark.sparkContext.parallelize(listOfFiles).flatMap { inputFile =>
-          val hadoopConfig = sconf.value
           log.info(s"Going to process data from $inputFile")
-          val inputFs = new Path(inputFile).getFileSystem(hadoopConfig)
           val ifs = new FileStreamer(inputFile, sconf.value)
 
           CobolProcessorBase.getRecordExtractor(readerParameters, copybookContents, ifs, None)
@@ -265,47 +263,16 @@ object SparkCobolProcessor {
 
       Future {
         val hadoopConfig = sconf.value
-        val inputFs = new Path(inputFIle).getFileSystem(hadoopConfig)
-        val ifs = new FileStreamer(inputFIle, sconf.value)
+
         val outputFile = new Path(outputPath, fileName)
         val outputFs = outputFile.getFileSystem(hadoopConfig)
-        val ofs = new BufferedOutputStream(outputFs.create(outputFile, true))
-
-        var originalException: Throwable = null
-
-        val recordCount = try {
-          cobolProcessor.process(ifs, ofs)(rawRecordProcessor)
-        } catch {
-          case ex: Throwable =>
-            originalException = ex
-            0L
-        } finally {
-          // Ugly code to ensure no exceptions escape unnoticed.
-          try {
-            ifs.close()
-          } catch {
-            case e: Throwable =>
-              if (originalException != null) {
-                originalException.addSuppressed(e)
-              } else {
-                originalException = e
-              }
-          }
 
-          try {
-            ofs.close()
-          } catch {
-            case e: Throwable =>
-              if (originalException != null) {
-                originalException.addSuppressed(e)
-              } else {
-                originalException = e
-              }
+        val recordCount = UsingUtils.using(new FileStreamer(inputFIle, sconf.value)) { ifs =>
+          UsingUtils.using(new BufferedOutputStream(outputFs.create(outputFile, true))) { ofs =>
+            cobolProcessor.process(ifs, ofs)(rawRecordProcessor)
           }
         }
 
-        if (originalException != null) throw originalException
-
         log.info(s"Writing to $outputFile succeeded!")
         recordCount
       }