[New Scheduler] Add container counter

core/scheduler/src/main/scala/org/apache/openwhisk/core/scheduler/queue/ContainerCounter.scala

@@ -0,0 +1,121 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.openwhisk.core.scheduler.queue
+
+import java.util.concurrent.atomic.AtomicInteger
+
+import akka.actor.{Actor, ActorRef, ActorSystem, Props}
+import org.apache.openwhisk.common.Logging
+import org.apache.openwhisk.core.etcd.EtcdClient
+import org.apache.openwhisk.core.etcd.EtcdKV.ContainerKeys
+import org.apache.openwhisk.core.service.{DeleteEvent, PutEvent, UnwatchEndpoint, WatchEndpoint, WatchEndpointOperation}
+
+import scala.collection.concurrent.TrieMap
+import scala.concurrent.{ExecutionContext, Future}
+
+class ContainerCounter(invocationNamespace: String, etcdClient: EtcdClient, watcherService: ActorRef)(
+  implicit val actorSystem: ActorSystem,
+  ec: ExecutionContext,
+  logging: Logging) {
+  private[queue] var existingContainerNumByNamespace: Int = 0
+  private[queue] var inProgressContainerNumByNamespace: Int = 0
+  private[queue] val references = new AtomicInteger(0)
+  private val watcherName = s"container-counter-$invocationNamespace"
+
+  private val inProgressContainerPrefixKeyByNamespace =
+    ContainerKeys.inProgressContainerPrefixByNamespace(invocationNamespace)
+  private val existingContainerPrefixKeyByNamespace =
+    ContainerKeys.existingContainersPrefixByNamespace(invocationNamespace)
+
+  private val watchedKeys = Seq(inProgressContainerPrefixKeyByNamespace, existingContainerPrefixKeyByNamespace)
+
+  private val watcher =
+    actorSystem.actorOf(Props(new Actor {
+      private var countingKeys = Set.empty[String]
+      private var waitingForCountKeys = Set.empty[String]
+
+      override def receive: Receive = {
+        case operation: WatchEndpointOperation if operation.isPrefix =>
+          if (countingKeys
+                .contains(operation.watchKey))
+            waitingForCountKeys += operation.watchKey
+          else {
+            countingKeys += operation.watchKey
+            refreshContainerCount(operation.watchKey)
+          }
+
+        case ReadyToGetCount(key) =>
+          if (waitingForCountKeys.contains(key)) {
+            waitingForCountKeys -= key
+            refreshContainerCount(key)
+          } else
+            countingKeys -= key
+      }
+    }))
+
+  private def refreshContainerCount(key: String): Future[Unit] = {
+    etcdClient
+      .getCount(key)
+      .map { count =>
+        key match {
+          case `inProgressContainerPrefixKeyByNamespace` => inProgressContainerNumByNamespace = count.toInt
+          case `existingContainerPrefixKeyByNamespace`   => existingContainerNumByNamespace = count.toInt
+        }
+        watcher ! ReadyToGetCount(key)
+      }
+      .recover {
+        case t: Throwable =>
+          logging.error(
+            this,
+            s"failed to get the number of existing containers for ${invocationNamespace} due to ${t}.")
+          watcher ! ReadyToGetCount(key)
+      }
+  }
+
+  def increaseReference(): ContainerCounter = {
+    if (references.incrementAndGet() == 1) {
+      watchedKeys.foreach { key =>
+        watcherService.tell(WatchEndpoint(key, "", true, watcherName, Set(PutEvent, DeleteEvent)), watcher)
+      }
+
+    }
+    this
+  }
+
+  def close(): Unit = {
+    if (references.decrementAndGet() == 0) {
+      watchedKeys.foreach { key =>
+        watcherService ! UnwatchEndpoint(key, true, watcherName)
+      }
+      NamespaceContainerCount.instances.remove(invocationNamespace)
+    }
+  }
+}
+
+object NamespaceContainerCount {
+  private[queue] val instances = TrieMap[String, ContainerCounter]()
+  def apply(namespace: String, etcdClient: EtcdClient, watcherService: ActorRef)(implicit actorSystem: ActorSystem,
+                                                                                 ec: ExecutionContext,
+                                                                                 logging: Logging): ContainerCounter = {
+    instances
+      .getOrElseUpdate(namespace, new ContainerCounter(namespace, etcdClient, watcherService))
+      .increaseReference()
+  }
+}
+
+case class ReadyToGetCount(key: String)