stuff

Author: ericm-db

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

@@ -46,15 +46,13 @@ private[sql] class RocksDBStateStoreProvider
   class RocksDBStateStore(lastVersion: Long) extends StateStore {
     /**
      * Trait and classes representing the internal state of the store
-     * 
+     *
      * State transitions:
      * - Initial state: UPDATING
      * - UPDATING -> COMMITTED: After successful commit()
      * - UPDATING -> ABORTED: After abort() or failed commit()
      * - UPDATING -> RELEASED: After release() without committing changes
-     * - COMMITTED -> RELEASED: After release() following a successful commit
-     * - ABORTED -> RELEASED: After release() following an abort
-     * 
+     *
      * The RELEASED state is terminal and indicates that resources have been released
      * without affecting the underlying data (unlike ABORTED which rolls back changes).
      */

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

@@ -120,16 +120,31 @@ trait ReadStateStore {
 
   /**
    * Release resources associated with this state store without rolling back changes.
-   * 
+   *
    * Unlike `abort()` which rolls back uncommitted changes, `release()` simply releases
-   * resources and locks without affecting the state data. This is particularly important
-   * in the read-then-write pattern where:
-   * 
-   * 1. A read-only store is opened to retrieve existing state
-   * 2. The same store is converted to a writable store using `getWriteStore()`
-   * 3. After the write store commits, we need to release resources without rolling back
-   *    the changes that were just committed
-   * 
+   * resources and locks without affecting the state data.
+   *
+   * IMPORTANT: This method is only needed when you have a standalone ReadStateStore that
+   * was NOT converted to a WriteStore. If you use the read-then-write pattern with
+   * `getWriteStore()`, you should NOT call `release()` on the original ReadStateStore.
+   *
+   * Usage scenarios:
+   * 1. When you have a standalone ReadStateStore that you're done with:
+   *    ```
+   *    val readStore = StateStore.getReadOnly(...)
+   *    // use readStore
+   *    readStore.release()  // Correct: release the standalone read store
+   *    ```
+   *
+   * 2. When using the read-then-write pattern:
+   *    ```
+   *    val readStore = StateStore.getReadOnly(...)
+   *    val writeStore = StateStore.getWriteStore(readStore, ...)
+   *    // use writeStore
+   *    writeStore.commit()
+   *    // DO NOT call readStore.release() here - the writeStore handles cleanup
+   *    ```
+   *
    * Implementations should ensure that:
    * 1. Any locks or resources held by this store are released
    * 2. No uncommitted changes are rolled back (unlike `abort()`)

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

@@ -55,7 +55,12 @@ import org.apache.spark.util.SerializableConfiguration
  *   readStore = getReadStore()
  *   writeStore = getWriteStore(readStore) // Reuses the same store connection
  *   writeStore.commit()
- *   // No need to abort/release readStore as it's the same underlying store
+ *  DO NOT call readStore.release() here - the writeStore handles cleanup
+ *
+ * IMPORTANT: When using this read-then-write pattern, you should NOT call release()
+ * on the original ReadStateStore after committing the WriteStore. Since both stores
+ * share the same underlying connection, the WriteStore's commit() will handle all
+ * necessary cleanup.
  */
 trait StateStoreRDDProvider {
   /**

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

@@ -109,16 +109,22 @@ package object state {
 
       val cleanedF = dataRDD.sparkContext.clean(storeReadFn)
       val wrappedF = (store: ReadStateStore, iter: Iterator[T]) => {
-        // Clean up the state store.
+        // Clean up the state store when the task completes successfully.
+        // We use release() instead of abort() because:
+        // 1. For standalone ReadStateStore, we want to release resources without rolling back
+        // 2. For ReadStateStore that will be converted to WriteStore, the WriteStore
+        // will handle cleanup
         val taskContext = TaskContext.get()
         taskContext.addTaskCompletionListener[Unit](_ => {
           store.release()
         })
+        
+        // On task failure, we need to abort to roll back any uncommitted changes.
+        // We use abort() instead of release() because:
+        // 1. We want to roll back any uncommitted changes to maintain consistency
+        // 2. abort() ensures proper cleanup and rollback of uncommitted changes
         taskContext.addTaskFailureListener(new TaskFailureListener {
           override def onTaskFailure(context: TaskContext, error: Throwable): Unit = {
-            // On task failure, we need to abort to roll back any uncommitted changes
-            // We don't call release() here because it would leave the state in an inconsistent state
-            // abort() ensures proper cleanup and rollback of uncommitted changes
             store.abort()
           }
         })