removing clera

Author: ericm-db

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

@@ -488,6 +488,7 @@ private[sql] class RocksDBStateStoreProvider
         existingStore match {
           case Some(stateStore: RocksDBStateStore) =>
             // Reuse existing store for getWriteStore case
+            StateStoreThreadLocalTracker.setUsedForWriteStore(true)
             stateStore
           case Some(_) =>
             throw new IllegalArgumentException("Existing store must be a RocksDBStateStore")

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

@@ -37,10 +37,13 @@ object StateStoreThreadLocalTracker {
   def setStore(store: ReadStateStore): Unit = readStore.set(store)
 
   def getStore: Option[ReadStateStore] = {
-    usedForWriteStore.set(true)
     Option(readStore.get())
   }
 
+  def setUsedForWriteStore(used: Boolean): Unit = {
+    usedForWriteStore.set(used)
+  }
+
   def isUsedForWriteStore: Boolean = usedForWriteStore.get()
 
   def clearStore(): Unit = readStore.remove()

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

@@ -113,12 +113,10 @@ package object state {
         val ctxt = TaskContext.get()
         ctxt.addTaskCompletionListener[Unit](_ => {
           if (!StateStoreThreadLocalTracker.isUsedForWriteStore) store.release()
-          StateStoreThreadLocalTracker.clearStore()
         })
         ctxt.addTaskFailureListener(new TaskFailureListener {
           override def onTaskFailure(context: TaskContext, error: Throwable): Unit =
             if (!StateStoreThreadLocalTracker.isUsedForWriteStore) store.abort()
-          StateStoreThreadLocalTracker.clearStore()
         })
         cleanedF(store, iter)
       }