From 0f9f6d484b11a57511cb2475b2291dc4dc9c1d47 Mon Sep 17 00:00:00 2001 From: soumyadeep-roy Date: Wed, 9 Sep 2026 12:09:56 +0530 Subject: [PATCH 1/2] Add StateStore.tryExclusiveCreate() atomic create-if-absent primitive Adds a create-if-absent primitive to the StateStore abstraction so callers can atomically persist a value only when no value is already stored at that key, instead of always overwriting via set(). - ZooKeeperStateStore/ZooKeeperManager: relies on ZooKeeper's create() failing with NodeExistsException, avoiding a racy separate exists check. - FileSystemStateStore: uses FileContext.create() with CREATE (no OVERWRITE), which fails atomically with FileAlreadyExistsException. - BlackholeStateStore: always returns true (no-op store, nothing to conflict with). - SessionStore.trySave() exposes the primitive at the session level. --- .../server/recovery/BlackholeStateStore.scala | 3 ++ .../recovery/FileSystemStateStore.scala | 16 ++++++++++ .../livy/server/recovery/SessionStore.scala | 9 ++++++ .../livy/server/recovery/StateStore.scala | 9 ++++++ .../server/recovery/ZooKeeperManager.scala | 15 +++++++++- .../server/recovery/ZooKeeperStateStore.scala | 4 +++ .../recovery/BlackholeStateStoreSpec.scala | 4 +++ .../recovery/FileSystemStateStoreSpec.scala | 25 ++++++++++++++++ .../server/recovery/SessionStoreSpec.scala | 21 +++++++++++++ .../recovery/ZooKeeperStateStoreSpec.scala | 30 +++++++++++++++++++ 10 files changed, 135 insertions(+), 1 deletion(-) diff --git a/server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala index df9a712a8..957ddde97 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/BlackholeStateStore.scala +++ b/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] diff --git a/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala index 6fee7f0e2..77cf10229 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala +++ b/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 => diff --git a/server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala index 04292957c..4d1f0084f 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/SessionStore.scala +++ b/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)) } diff --git a/server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala index 2e454db07..b26722c93 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/StateStore.scala +++ b/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. diff --git a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala b/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala index 9ad86bdc8..29fa25644 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperManager.scala +++ b/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 diff --git a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala index ceb2258d0..75b896344 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/ZooKeeperStateStore.scala +++ b/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)) } diff --git a/server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala index 014652270..6ff8caa37 100644 --- a/server/src/test/scala/org/apache/livy/server/recovery/BlackholeStateStoreSpec.scala +++ b/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 diff --git a/server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala index cbedb1e6e..0f5898a9c 100644 --- a/server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala +++ b/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 {} diff --git a/server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala index 923f70981..b038cd31d 100644 --- a/server/src/test/scala/org/apache/livy/server/recovery/SessionStoreSpec.scala +++ b/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)), diff --git a/server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala index 544d6c86d..90e195b6c 100644 --- a/server/src/test/scala/org/apache/livy/server/recovery/ZooKeeperStateStoreSpec.scala +++ b/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 => From 784dbb085b7d85ec5a3aa02c58ceb695eace6cae Mon Sep 17 00:00:00 2001 From: soumyadeep-roy Date: Fri, 18 Sep 2026 20:56:52 +0530 Subject: [PATCH 2/2] [LIVY-1078] Make tryExclusiveCreate crash-safe via write-then-rename tryExclusiveCreate() previously wrote directly to the destination path and relied on CreateFlag.CREATE (no OVERWRITE) to fail atomically when the key already existed. If livy-server crashed mid-write, the destination would be left partially written yet "claimed" forever, since every future exclusive-create attempt would see the existing file and back off. Write to a uniquely-named temp file first, then perform the exclusive claim via a non-overwriting atomic rename (mirrors set()'s existing write-then-rename shape). The destination is only ever visible once it is complete, and the rename itself still fails atomically with FileAlreadyExistsException if another writer wins the race, so no separate (and racy) exists-check is needed. Clean up the orphaned temp file when the rename loses the race. --- .../recovery/FileSystemStateStore.scala | 40 ++++++++++++++----- .../recovery/FileSystemStateStoreSpec.scala | 30 ++++++++++++-- 2 files changed, 56 insertions(+), 14 deletions(-) diff --git a/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala b/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala index 77cf10229..9b5a41943 100644 --- a/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala +++ b/server/src/main/scala/org/apache/livy/server/recovery/FileSystemStateStore.scala @@ -20,6 +20,7 @@ package org.apache.livy.server.recovery import java.io.{FileNotFoundException, IOException} import java.net.URI import java.util +import java.util.UUID import java.util.concurrent.TimeUnit import scala.reflect.ClassTag @@ -109,19 +110,38 @@ 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() + // Write to a uniquely-named temp file first, then atomically rename it into place + // without overwrite -- same write-then-rename shape as set(), so a crash mid-write + // never leaves a partially-written file at the destination. Omitting Rename.OVERWRITE + // still makes the claim exclusive: the rename fails atomically with + // FileAlreadyExistsException if another writer already claimed the key in the + // meantime, so no separate (and racy) exists-check is needed. + val tmpPath = absPath(s"$key.${UUID.randomUUID()}.tmp") + val createFlag = util.EnumSet.of(CreateFlag.CREATE, CreateFlag.OVERWRITE) + + usingResource(fileContext.create(tmpPath, createFlag, CreateOpts.createParent())) { tmpFile => + tmpFile.write(serializeToBytes(value)) + tmpFile.close() + } + + val claimed = + try { + fileContext.rename(tmpPath, absPath(key), Rename.NONE) + true + } catch { + case _: FileAlreadyExistsException => + fileContext.delete(tmpPath, false) + false } - true + + try { + val crcPath = new Path(tmpPath.getParent, s".${tmpPath.getName}.crc") + fileContext.delete(crcPath, false) } catch { - case _: FileAlreadyExistsException => false + case NonFatal(e) => // Swallow the exception. } + + claimed } override def get[T: ClassTag](key: String): Option[T] = { diff --git a/server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala b/server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala index 0f5898a9c..45e3cb209 100644 --- a/server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala +++ b/server/src/test/scala/org/apache/livy/server/recovery/FileSystemStateStoreSpec.scala @@ -47,6 +47,14 @@ class FileSystemStateStoreSpec extends AnyFunSpec with LivyBaseUnitTestSuite { override def describeTo(d: Description): Unit = { matcher.describeTo(d) } }) + def pathMatching(pattern: String): Path = argThat(new ArgumentMatcher[Path] { + override def matches(path: Any): Boolean = path.toString.matches(pattern) + + override def describeTo(d: Description): Unit = { + d.appendText(s"path matching $pattern") + } + }) + def makeConf(): LivyConf = { val conf = new LivyConf() conf.set(LivyConf.RECOVERY_STATE_STORE_URL, "file://tmp/") @@ -113,10 +121,13 @@ 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") { + it("tryExclusiveCreate should write to a temp file, then rename it into place " + + "without overwrite, and return true if the key doesn't exist") { val fileContext = mockFileContext("700") val outputStream = mock[FSDataOutputStream] - when(fileContext.create(pathEq("/key"), any[util.EnumSet[CreateFlag]], any[CreateOpts])) + val tmpPathPattern = "/key\\..*\\.tmp" + when(fileContext.create( + pathMatching(tmpPathPattern), any[util.EnumSet[CreateFlag]], any[CreateOpts])) .thenReturn(outputStream) val stateStore = new FileSystemStateStore(makeConf(), Some(fileContext)) @@ -126,16 +137,27 @@ class FileSystemStateStoreSpec extends AnyFunSpec with LivyBaseUnitTestSuite { created shouldBe true verify(outputStream).write(""""value"""".getBytes) verify(outputStream, atLeastOnce).close() + verify(fileContext).rename(pathMatching(tmpPathPattern), pathEq("/key"), equal(Rename.NONE)) + verify(fileContext).delete(pathMatching("/\\.key\\..*\\.tmp\\.crc"), equal(false)) } - it("tryExclusiveCreate should return false if the key already exists") { + it("tryExclusiveCreate should return false and clean up its temp file if the key " + + "already exists") { val fileContext = mockFileContext("700") - when(fileContext.create(pathEq("/key"), any[util.EnumSet[CreateFlag]], any[CreateOpts])) + val outputStream = mock[FSDataOutputStream] + val tmpPathPattern = "/key\\..*\\.tmp" + when(fileContext.create( + pathMatching(tmpPathPattern), any[util.EnumSet[CreateFlag]], any[CreateOpts])) + .thenReturn(outputStream) + when(fileContext.rename(pathMatching(tmpPathPattern), pathEq("/key"), equal(Rename.NONE))) .thenThrow(new FileAlreadyExistsException("Unit test")) val stateStore = new FileSystemStateStore(makeConf(), Some(fileContext)) stateStore.tryExclusiveCreate("key", "value") shouldBe false + + // The temp file must not be left behind once we know we lost the claim race. + verify(fileContext).delete(pathMatching(tmpPathPattern), equal(false)) } it("get should read file") {