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..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 @@ -108,6 +109,41 @@ class FileSystemStateStore( } } + override def tryExclusiveCreate(key: String, value: Object): Boolean = { + // 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 + } + + try { + val crcPath = new Path(tmpPath.getParent, s".${tmpPath.getName}.crc") + fileContext.delete(crcPath, false) + } catch { + case NonFatal(e) => // Swallow the exception. + } + + claimed + } + 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..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,6 +121,45 @@ class FileSystemStateStoreSpec extends AnyFunSpec with LivyBaseUnitTestSuite { verify(fileContext).delete(pathEq("/.key.tmp.crc"), equal(false)) } + 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] + val tmpPathPattern = "/key\\..*\\.tmp" + when(fileContext.create( + pathMatching(tmpPathPattern), 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() + verify(fileContext).rename(pathMatching(tmpPathPattern), pathEq("/key"), equal(Rename.NONE)) + verify(fileContext).delete(pathMatching("/\\.key\\..*\\.tmp\\.crc"), equal(false)) + } + + it("tryExclusiveCreate should return false and clean up its temp file if the key " + + "already exists") { + val fileContext = mockFileContext("700") + 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") { 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 =>