assign based on invoker pool size

Author: Brendan Doyle

core/invoker/src/main/scala/org/apache/openwhisk/core/invoker/InstanceIdAssigner.scala

@@ -41,7 +41,26 @@ private[invoker] class InstanceIdAssigner(connectionString: String)(implicit log
 
     val rootPath = "/invokers/idAssignment/mapping"
     val myIdPath = s"$rootPath/$name"
-    val assignedId = if (overwriteId.isEmpty) {
+    val assignedId = overwriteId.map(newId => {
+      val invokers = zkClient.getChildren.forPath(rootPath).asScala
+
+      if (invokers.size < newId)
+        throw new IllegalArgumentException(s"invokerReg: cannot assign $newId to $name: not enough invokers")
+
+      //check if the invokerId already exists for another unique name and delete if it does
+      invokers
+        .map(uniqueName => {
+          val idPath = s"$rootPath/$uniqueName"
+          (idPath, BigInt(zkClient.getData.forPath(idPath)).intValue)
+        })
+        .find(_._2 == newId)
+        .map(id => zkClient.delete().forPath(id._1))
+
+      zkClient.create().orSetData().forPath(myIdPath, BigInt(newId).toByteArray)
+
+      logger.info(this, s"invokerReg: invoker $name was assigned invokerId $newId")
+      newId
+    }).getOrElse({
       Option(zkClient.checkExists().forPath(myIdPath)) match {
         case None =>
           // path doesn't exist -> no previous mapping for this invoker
@@ -51,8 +70,11 @@ private[invoker] class InstanceIdAssigner(connectionString: String)(implicit log
 
           def assignId(): Int = {
             val current = idCounter.getVersionedValue()
-            if (idCounter.trySetCount(current, current.getValue() + 1)) {
-              current.getValue()
+            val numInvokers = Option(zkClient.checkExists().forPath(rootPath))
+              .map(_ => zkClient.getChildren.forPath(rootPath).size())
+              .getOrElse(0)
+            if (idCounter.trySetCount(current, numInvokers + 1)) {
+              numInvokers
             } else {
               assignId()
             }
@@ -66,33 +88,12 @@ private[invoker] class InstanceIdAssigner(connectionString: String)(implicit log
 
         case Some(_) =>
           // path already exists -> there is a previous mapping for this invoker we should use
-          val rawOldId = zkClient.getData().forPath(myIdPath)
+          val rawOldId = zkClient.getData.forPath(myIdPath)
           val oldId = BigInt(rawOldId).intValue
           logger.info(this, s"invokerReg: invoker $name was assigned its previous invokerId $oldId")
           oldId
       }
-    } else {
-      val newId = overwriteId.get
-
-      //check if the invokerId already exists for another unique name
-      val instanceIdExists = zkClient
-        .getChildren()
-        .forPath(rootPath)
-        .asScala
-        .map(uniqueName => {
-          val idPath = s"$rootPath/$uniqueName"
-          BigInt(zkClient.getData().forPath(idPath)).intValue
-        })
-        .find(_ == newId)
-
-      if (instanceIdExists.nonEmpty) {
-        throw new IllegalArgumentException(s"invokerReg: an invoker with id $newId already exists in zookeeper")
-      }
-
-      zkClient.create().orSetData().forPath(myIdPath, BigInt(newId).toByteArray)
-      logger.info(this, s"invokerReg: invoker $name was assigned invokerId $newId")
-      newId
-    }
+    })
 
     zkClient.close()
     assignedId

tests/src/test/scala/org/apache/openwhisk/core/invoker/test/InstanceIdAssignerTests.scala

@@ -53,13 +53,20 @@ class InstanceIdAssignerTests extends FlatSpec with Matchers with StreamLogging
   it should "attempt to overwrite id for unique name if overwrite set" in {
     val assigner = new InstanceIdAssigner(zkServer.getConnectString)
     assigner.setAndGetId("foo") shouldBe 0
-    assigner.setAndGetId("foo", Some(2)) shouldBe 2
+    assigner.setAndGetId("bar", Some(0))
   }
 
-  it should "fail to overwrite an id for unique name that already exists" in {
+  it should "overwrite an id for unique name that already exists and reset overwritten id" in {
     val assigner = new InstanceIdAssigner(zkServer.getConnectString)
     assigner.setAndGetId("foo") shouldBe 0
-    assigner.setAndGetId("bar") shouldBe 1
-    assertThrows[IllegalArgumentException](assigner.setAndGetId("bar", Some(1)))
+    assigner.setAndGetId("bar", Some(0)) shouldBe 0
+    assigner.setAndGetId("foo") shouldBe 1
+    assigner.setAndGetId("cat") shouldBe 2
+  }
+
+  it should "fail to overwrite an id too large for the invoker pool size" in {
+    val assigner = new InstanceIdAssigner(zkServer.getConnectString)
+    assigner.setAndGetId("foo") shouldBe 0
+    assertThrows[IllegalArgumentException](assigner.setAndGetId("bar", Some(2)))
   }
 }