|
| 1 | +package org.apache.openwhisk.core.loadBalancer |
| 2 | + |
| 3 | +import akka.actor.{ActorRef, ActorSystem, Props} |
| 4 | +import akka.stream.ActorMaterializer |
| 5 | +import org.apache.openwhisk.common.{Logging, TransactionId} |
| 6 | +import org.apache.openwhisk.core.{ConfigKeys, WhiskConfig} |
| 7 | +import org.apache.openwhisk.core.WhiskConfig._ |
| 8 | +import org.apache.openwhisk.core.connector.{ActivationMessage, MessagingProvider} |
| 9 | +import org.apache.openwhisk.core.entity._ |
| 10 | +import org.apache.openwhisk.spi.SpiLoader |
| 11 | +import pureconfig.loadConfigOrThrow |
| 12 | +import spray.json._ |
| 13 | +import pureconfig._ |
| 14 | +import pureconfig.generic.auto._ |
| 15 | + |
| 16 | +import scala.concurrent.Future |
| 17 | + |
| 18 | +class MuxBalancer(config: WhiskConfig, |
| 19 | + feedFactory: FeedFactory, |
| 20 | + controllerInstance: ControllerInstanceId, |
| 21 | + implicit val messagingProvider: MessagingProvider = SpiLoader.get[MessagingProvider], |
| 22 | + override val lbConfig: ShardingContainerPoolBalancerConfig = |
| 23 | + loadConfigOrThrow[ShardingContainerPoolBalancerConfig](ConfigKeys.loadbalancer))( |
| 24 | + implicit actorSystem: ActorSystem, |
| 25 | + logging: Logging, |
| 26 | + materializer: ActorMaterializer) |
| 27 | + extends CommonLoadBalancer(config, feedFactory, controllerInstance) { |
| 28 | + |
| 29 | + private val defaultLoadBalancer = |
| 30 | + getClass[LoadBalancerProvider](lbConfig.strategy.default).instance(config, controllerInstance) |
| 31 | + private val customLoadBalancerMap: Map[String, LoadBalancer] = |
| 32 | + lbConfig.strategy.custom.foldLeft(Map.empty[String, LoadBalancer]) { |
| 33 | + case (result, (name, strategyConfig)) => |
| 34 | + result + (name -> getClass[LoadBalancerProvider](strategyConfig.className).instance(config, controllerInstance)) |
| 35 | + } |
| 36 | + |
| 37 | + /** |
| 38 | + * Instantiates an object of the given type. |
| 39 | + * |
| 40 | + * Similar to SpiLoader.get, with the difference that the constructed class does not need to be declared as Spi. |
| 41 | + * Thus there could be multiple classes implementing same interface constructed at the same time |
| 42 | + * |
| 43 | + * @param name the name of the class |
| 44 | + * @tparam A expected type to return |
| 45 | + * @return instance of the class |
| 46 | + */ |
| 47 | + private def getClass[A](name: String): A = { |
| 48 | + val clazz = Class.forName(name + "$") |
| 49 | + val classInst = clazz.getField("MODULE$").get(clazz).asInstanceOf[A] |
| 50 | + classInst |
| 51 | + } |
| 52 | + |
| 53 | + override def invokerHealth(): Future[IndexedSeq[InvokerHealth]] = Future.successful(IndexedSeq.empty[InvokerHealth]) |
| 54 | + override protected def releaseInvoker(invoker: InvokerInstanceId, entry: ActivationEntry) = { |
| 55 | + // Currently do nothing |
| 56 | + } |
| 57 | + override protected val invokerPool: ActorRef = actorSystem.actorOf(Props.empty) |
| 58 | + |
| 59 | + /** |
| 60 | + * Publish a message to the loadbalancer |
| 61 | + * |
| 62 | + * Select the LoadBalancer based on the annotation, if available, otherwise use the default one |
| 63 | + **/ |
| 64 | + override def publish(action: ExecutableWhiskActionMetaData, msg: ActivationMessage)( |
| 65 | + implicit transid: TransactionId): Future[Future[Either[ActivationId, WhiskActivation]]] = { |
| 66 | + action.annotations.get("activationStrategy") match { |
| 67 | + case None => |
| 68 | + defaultLoadBalancer.publish(action, msg) |
| 69 | + case Some(JsString(value)) => { |
| 70 | + if (customLoadBalancerMap.contains(value)) { |
| 71 | + customLoadBalancerMap(value).publish(action, msg) |
| 72 | + } else { |
| 73 | + defaultLoadBalancer.publish(action, msg) |
| 74 | + } |
| 75 | + } |
| 76 | + case Some(_) => defaultLoadBalancer.publish(action, msg) |
| 77 | + } |
| 78 | + } |
| 79 | +} |
| 80 | + |
| 81 | +object MuxBalancer extends LoadBalancerProvider { |
| 82 | + |
| 83 | + override def instance(whiskConfig: WhiskConfig, instance: ControllerInstanceId)( |
| 84 | + implicit actorSystem: ActorSystem, |
| 85 | + logging: Logging, |
| 86 | + materializer: ActorMaterializer): LoadBalancer = { |
| 87 | + |
| 88 | + new MuxBalancer(whiskConfig, createFeedFactory(whiskConfig, instance), instance) |
| 89 | + } |
| 90 | + |
| 91 | + def requiredProperties = |
| 92 | + ExecManifest.requiredProperties ++ |
| 93 | + wskApiHost |
| 94 | +} |
0 commit comments