[LIVY-1078] Add StateStore.tryExclusiveCreate() abstraction across state-store impls

server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala

@@ -28,6 +28,9 @@ import org.apache.livy.LivyConf
 class BlackholeStateStore(livyConf: LivyConf) extends StateStore(livyConf) {
   def set(key: String, value: Object): Unit = {}
 
+  // Recovery is disabled, so there's no persisted state to conflict with.
+  def tryExclusiveCreate(key: String, value: Object): Boolean = true
+
   def get[T: ClassTag](key: String): Option[T] = None
 
   def getChildren(key: String): Seq[String] = List.empty[String]

server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala

@@ -108,6 +108,22 @@ class FileSystemStateStore(
     }
   }
 
+  override def tryExclusiveCreate(key: String, value: Object): Boolean = {
+    // CREATE without OVERWRITE fails atomically with FileAlreadyExistsException if the
+    // destination already exists, so no separate exists-check (which would be racy) is needed.
+    val createFlag = util.EnumSet.of(CreateFlag.CREATE)
+    try {
+      usingResource(fileContext.create(absPath(key), createFlag, CreateOpts.createParent())) {
+        newFile =>
+          newFile.write(serializeToBytes(value))
+          newFile.close()
+      }
+      true
+    } catch {
+      case _: FileAlreadyExistsException => false
+    }
+  }
+
   override def get[T: ClassTag](key: String): Option[T] = {
     try {
       usingResource(fileContext.open(absPath(key))) { is =>

server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala

@@ -46,6 +46,15 @@ class SessionStore(
     store.set(sessionPath(sessionType, m.id), m)
   }
 
+  /**
+   * Persist a session to the session state store only if no session is already stored
+   * at that path.
+   * @return true if the session was persisted, false if a session with this id already exists.
+   */
+  def trySave(sessionType: String, m: RecoveryMetadata): Boolean = {
+    store.tryExclusiveCreate(sessionPath(sessionType, m.id), m)
+  }
+
   def saveNextSessionId(sessionType: String, id: Int): Unit = {
     store.set(sessionManagerPath(sessionType), SessionManagerState(id))
   }

server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala

@@ -58,6 +58,15 @@ abstract class StateStore(livyConf: LivyConf) extends JsonMapper {
    */
   def get[T: ClassTag](key: String): Option[T]
 
+  /**
+   * Atomically create a key-value pair in this state store only if the key doesn't already
+   * exist. Unlike [[set]], this never overwrites an existing value.
+   * @return true if the key was created, false if the key already exists.
+   * @throws Exception Throw when persisting the state store fails for a reason other than
+   *                    the key already existing.
+   */
+  def tryExclusiveCreate(key: String, value: Object): Boolean
+
   /**
    * Treat keys in this state store as a directory tree and
    * return names of the direct children of the key.

server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala

@@ -25,7 +25,7 @@ import org.apache.curator.framework.CuratorFramework
 import org.apache.curator.framework.CuratorFrameworkFactory
 import org.apache.curator.framework.state.{ConnectionState, ConnectionStateListener}
 import org.apache.curator.retry.RetryNTimes
-import org.apache.zookeeper.KeeperException.NoNodeException
+import org.apache.zookeeper.KeeperException.{NodeExistsException, NoNodeException}
 import org.apache.zookeeper.client.ZKClientConfig
 
 import org.apache.livy.LivyConf
@@ -157,6 +157,19 @@ class ZooKeeperManager(
     }
   }
 
+  // Atomically create the znode only if it doesn't already exist. Relies on ZooKeeper's
+  // create() failing with NodeExistsException rather than a separate exists-check, which
+  // would be racy against concurrent creators.
+  def tryCreate(key: String, value: Object): Boolean = {
+    val data = serializeToBytes(value)
+    try {
+      curatorClient.create().creatingParentsIfNeeded().forPath(key, data)
+      true
+    } catch {
+      case _: NodeExistsException => false
+    }
+  }
+
   def get[T: ClassTag](key: String): Option[T] = {
     if (curatorClient.checkExists().forPath(key) == null) {
       None

server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala

@@ -36,6 +36,10 @@ class ZooKeeperStateStore(
     zkManager.get(prefixKey(key))
   }
 
+  override def tryExclusiveCreate(key: String, value: Object): Boolean = {
+    zkManager.tryCreate(prefixKey(key), value)
+  }
+
   override def getChildren(key: String): Seq[String] = {
     zkManager.getChildren(prefixKey(key))
   }

server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala

@@ -31,6 +31,10 @@ class BlackholeStateStoreSpec extends AnyFunSpec with LivyBaseUnitTestSuite {
       stateStore.set("", 1.asInstanceOf[Object])
     }
 
+    it("tryExclusiveCreate should return true and not throw") {
+      stateStore.tryExclusiveCreate("", 1.asInstanceOf[Object]) shouldBe true
+    }
+
     it("get should return None") {
       val v = stateStore.get[Object]("")
       v shouldBe None

server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala

@@ -113,6 +113,31 @@ class FileSystemStateStoreSpec extends AnyFunSpec with LivyBaseUnitTestSuite {
       verify(fileContext).delete(pathEq("/.key.tmp.crc"), equal(false))
     }
 
+    it("tryExclusiveCreate should write file and return true if key doesn't exist") {
+      val fileContext = mockFileContext("700")
+      val outputStream = mock[FSDataOutputStream]
+      when(fileContext.create(pathEq("/key"), any[util.EnumSet[CreateFlag]], any[CreateOpts]))
+        .thenReturn(outputStream)
+
+      val stateStore = new FileSystemStateStore(makeConf(), Some(fileContext))
+
+      val created = stateStore.tryExclusiveCreate("key", "value")
+
+      created shouldBe true
+      verify(outputStream).write(""""value"""".getBytes)
+      verify(outputStream, atLeastOnce).close()
+    }
+
+    it("tryExclusiveCreate should return false if the key already exists") {
+      val fileContext = mockFileContext("700")
+      when(fileContext.create(pathEq("/key"), any[util.EnumSet[CreateFlag]], any[CreateOpts]))
+        .thenThrow(new FileAlreadyExistsException("Unit test"))
+
+      val stateStore = new FileSystemStateStore(makeConf(), Some(fileContext))
+
+      stateStore.tryExclusiveCreate("key", "value") shouldBe false
+    }
+
     it("get should read file") {
       val fileContext = mockFileContext("700")
       abstract class MockInputStream extends InputStream with Seekable with PositionedReadable {}

server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala

@@ -45,6 +45,27 @@ class SessionStoreSpec extends AnyFunSpec with LivyBaseUnitTestSuite {
       verify(stateStore).set(s"$sessionPath/99", m)
     }
 
+    it("should exclusively create session state when trying to save a session") {
+      val stateStore = mock[StateStore]
+      val sessionStore = new SessionStore(conf, stateStore)
+
+      val m = TestRecoveryMetadata(99)
+      when(stateStore.tryExclusiveCreate(s"$sessionPath/99", m)).thenReturn(true)
+
+      sessionStore.trySave(sessionType, m) shouldBe true
+      verify(stateStore).tryExclusiveCreate(s"$sessionPath/99", m)
+    }
+
+    it("should return false from trySave if the session already exists") {
+      val stateStore = mock[StateStore]
+      val sessionStore = new SessionStore(conf, stateStore)
+
+      val m = TestRecoveryMetadata(99)
+      when(stateStore.tryExclusiveCreate(s"$sessionPath/99", m)).thenReturn(false)
+
+      sessionStore.trySave(sessionType, m) shouldBe false
+    }
+
     it("should return existing sessions") {
       val validMetadata = Map(
         "0" -> Some(TestRecoveryMetadata(0)),

server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala

@@ -23,6 +23,7 @@ import org.apache.curator.framework.CuratorFramework
 import org.apache.curator.framework.api._
 import org.apache.curator.framework.listen.Listenable
 import org.apache.curator.framework.state.{ConnectionState, ConnectionStateListener}
+import org.apache.zookeeper.KeeperException.NodeExistsException
 import org.apache.zookeeper.data.Stat
 import org.mockito.ArgumentCaptor
 import org.mockito.Mockito._
@@ -107,6 +108,35 @@ class ZooKeeperStateStoreSpec extends AnyFunSpec with LivyBaseUnitTestSuite {
       }
     }
 
+    it("tryExclusiveCreate should create key and return true if it doesn't exist") {
+      withMock { f =>
+        val createBuilder = mock[CreateBuilder]
+        when(f.curatorClient.create()).thenReturn(createBuilder)
+        val p = mock[ProtectACLCreateModeStatPathAndBytesable[String]]
+        when(createBuilder.creatingParentsIfNeeded()).thenReturn(p)
+
+        val created = f.stateStore.tryExclusiveCreate("key", 1.asInstanceOf[Object])
+
+        created shouldBe true
+        verify(p).forPath(prefixedKey, Array[Byte](49))
+      }
+    }
+
+    it("tryExclusiveCreate should return false if the key already exists") {
+      withMock { f =>
+        val createBuilder = mock[CreateBuilder]
+        when(f.curatorClient.create()).thenReturn(createBuilder)
+        val p = mock[ProtectACLCreateModeStatPathAndBytesable[String]]
+        when(createBuilder.creatingParentsIfNeeded()).thenReturn(p)
+        when(p.forPath(prefixedKey, Array[Byte](49)))
+          .thenThrow(new NodeExistsException(prefixedKey))
+
+        val created = f.stateStore.tryExclusiveCreate("key", 1.asInstanceOf[Object])
+
+        created shouldBe false
+      }
+    }
+
     it("get should retrieve retry policy configs") {
       conf.set(LivyConf.ZK_RETRY_POLICY, "11,77")
         withMock { f =>