More improvements

Author: tdas

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/CheckpointFileManager.scala

@@ -243,8 +243,8 @@ class FileSystemBasedCheckpointFileManager(path: Path, hadoopConf: Configuration
     }
 
     try {
-      if (!fs.rename(srcPath, dstPath)) {
-        if (fs.exists(dstPath) && !overwriteIfPossible) {
+      if (!fs.rename(srcPath, dstPath) && !overwriteIfPossible) {
+        if (fs.exists(dstPath)) {
           // Some implementations of FileSystem may not throw FileAlreadyExistsException but
           // only return false if file already exists. Explicitly throw the error.
           // Note that this is definitely not atomic, so this is only a best-effort attempt

sql/core/src/main/scala/org/apache/spark/sql/execution/streaming/state/HDFSBackedStateStoreProvider.scala

@@ -148,7 +148,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
       // the other used for read+write. We don't want the read-only to delete state files.
       if (state == UPDATING) {
         state = ABORTED
-        cancelIfPossible(compressedStream, deltaFileStream)
+        cancelDeltaFile(compressedStream, deltaFileStream)
       } else {
         state = ABORTED
       }
@@ -403,7 +403,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
       output.close()
     } catch {
       case e: Throwable =>
-        cancelIfPossible(compressedStream = output, rawStream = rawOutput)
+        cancelDeltaFile(compressedStream = output, rawStream = rawOutput)
         throw e
     }
     logInfo(s"Written snapshot file for version $version of $this at $targetFile")
@@ -415,7 +415,7 @@ private[state] class HDFSBackedStateStoreProvider extends StateStoreProvider wit
    * @param compressedStream the compressed stream.
    * @param rawStream the underlying stream which needs to be cancelled.
    */
-  private def cancelIfPossible(
+  private def cancelDeltaFile(
       compressedStream: DataOutputStream,
       rawStream: CancellableFSDataOutputStream): Unit = {
     try {

sql/core/src/test/scala/org/apache/spark/sql/execution/streaming/CheckpointFileManagerSuite.scala

@@ -54,16 +54,25 @@ abstract class CheckpointFileManagerTests extends SparkFunSuite {
       assert(fm.list(basePath, acceptAllFilter).exists(_.getPath.getName == "dir"))
       assert(fm.list(basePath, rejectAllFilter).length === 0)
 
-      // Create atomic and exists
-      val path = new Path(s"$dir/file")
+      // Create atomic without overwrite
+      var path = new Path(s"$dir/file")
       assert(!fm.exists(path))
       fm.createAtomic(path, overwriteIfPossible = false).cancel()
       assert(!fm.exists(path))
       fm.createAtomic(path, overwriteIfPossible = false).close()
       assert(fm.exists(path))
       intercept[IOException] {
+        // should throw exception since file exists and overwrite is false
         fm.createAtomic(path, overwriteIfPossible = false).close()
       }
+
+      // Create atomic with overwrite if possible
+      path = new Path(s"$dir/file2")
+      assert(!fm.exists(path))
+      fm.createAtomic(path, overwriteIfPossible = true).cancel()
+      assert(!fm.exists(path))
+      fm.createAtomic(path, overwriteIfPossible = true).close()
+      assert(fm.exists(path))
       fm.createAtomic(path, overwriteIfPossible = true).close()  // should not throw exception
 
       // Open and delete