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 =>