akka.actor.ActorSelection Scala Examples

package cmwell.ctrl.tasks

import akka.actor.ActorSelection
import akka.actor.FSM.Failure
import akka.util.Timeout
import cmwell.ctrl.checkers.{ComponentState, GreenStatus, YellowStatus}
import cmwell.ctrl.commands._
import cmwell.ctrl.hc._
import cmwell.ctrl.server.CommandActor
import com.typesafe.scalalogging.LazyLogging
import k.grid.Grid

import scala.concurrent.{Future, Promise}
import scala.concurrent.duration._
import akka.pattern.ask

import scala.util.Success
import scala.concurrent.ExecutionContext.Implicits.global


case class AddNode(node: String) extends Task with LazyLogging {
  implicit val timeout = Timeout(15.seconds)

  private def startElasticsearch(cmd: ActorSelection, prom: Promise[Unit]): Unit = {
    logger.info(s"Starting Elasticsearch on node $node")
    cmd ! StartElasticsearch
    Grid.system.scheduler.scheduleOnce(60.seconds) {
      (HealthActor.ref ? GetElasticsearchDetailedStatus).mapTo[ElasticsearchGridStatus].map { f =>
        f.getStatesMap.get(node) match {
          case Some(s) =>
            if (s.getColor == GreenStatus || s.getColor == YellowStatus) prom.success(())
            else startElasticsearch(cmd, prom)
          case None => startElasticsearch(cmd, prom)
        }
      }
    }
  }

  private def startCassandra(cmd: ActorSelection, prom: Promise[Unit]): Unit = {
    logger.info(s"Starting Cassandra on node $node")
    cmd ! StartCassandra
    Grid.system.scheduler.scheduleOnce(60.seconds) {
      (HealthActor.ref ? GetCassandraDetailedStatus).mapTo[CassandraGridStatus].map { f =>
        f.getStatesMap.get(node) match {
          case Some(s) =>
            if (s.getColor == GreenStatus) prom.success(())
            else startCassandra(cmd, prom)
          case None => startCassandra(cmd, prom)
        }
      }
    }
  }

  override def exec: Future[TaskResult] = {
    val cmd = CommandActor.select(node)
    val esPromise = Promise[Unit]
    val casPromise = Promise[Unit]

    startElasticsearch(cmd, esPromise)
    startCassandra(cmd, casPromise)

    val esCancelable = cancel(esPromise, 24.hours)
    val casCancelable = cancel(casPromise, 24.hours)
    val esFuture = esPromise.future
    val casFuture = casPromise.future

    // cancel the cancelables when the future succeeded
    esFuture.foreach(x => esCancelable.cancel())
    casFuture.foreach(x => casCancelable.cancel())

    val fut = for {
      esStarted <- esFuture
      casStarted <- casFuture
    } yield {
      logger.info("Starting CM-WELL components")
      cmd ! StartKafka
      cmd ! StartBg
      cmd ! StartWebserver
      cmd ! StartCw
      cmd ! StartDc
    }

    fut.map(r => TaskSuccessful).recover { case err: Throwable => TaskFailed }
  }
} 

package io.coral.api.security

import akka.actor.ActorSelection
import akka.event.slf4j.Logger
import spray.http.HttpRequest
import spray.routing.directives.AuthMagnet
import io.coral.api.{Runtime, CoralConfig}
import scala.concurrent.ExecutionContext

object AuthenticatorChooser {
	var config: Option[CoralConfig] = None
}


class AuthenticatorChooser() {
	// When the application is started, it is not defined
	var authenticationMode: Option[String] = None

	def authenticator(cassandra: ActorSelection,
					  authenticatorActor: ActorSelection)
					 (implicit ec: ExecutionContext): AuthMagnet[AuthInfo] = {
		// This is only set the first time the API service is run
		if (!authenticationMode.isDefined) {
			setAuthenticationMode(AuthenticatorChooser.config)
		}

		authenticationMode match {
			case Some("accept-all") =>
				AcceptAllAuthenticator.acceptAllAuthenticator(AuthenticatorChooser.config.get)
			case Some("coral") =>
				CoralAuthenticator.coralAuthenticator(cassandra, authenticatorActor)
			case Some("ldap") =>
				LDAPAuthenticator.ldapAuthenticator(authenticatorActor)
			case other =>
				throw new Exception(s"Invalid authentication mode provided: $other")
		}
	}

	def setAuthenticationMode(config: Option[CoralConfig]) {
		authenticationMode = config match {
			case None => None
			case Some(c) => c.coral.authentication.mode match {
				case "accept-all" => Some("accept-all")
				case "coral" => Some("coral")
				case "ldap" => Some("ldap")
				case _ =>
					Logger(getClass.getName).error("[FATAL ERROR]: Invalid authentication mode provided. Exiting.")
					System.exit(1)
					None
			}
		}
	}

	def coralAuthorizer(request: HttpRequest, authInfo: AuthInfo, runtime: Runtime)(
		implicit ec: ExecutionContext): Boolean = {
			authInfo.hasPermission(request, Some(runtime))
	}

	def coralAuthorizer(request: HttpRequest, authInfo: AuthInfo)(implicit ec: ExecutionContext): Boolean = {
		authInfo.hasPermission(request, None)
	}
} 

package org.apache.toree.kernel.protocol.v5.kernel

import akka.actor.{ActorSelection, ActorSystem, Props}
import akka.testkit.{TestKit, TestProbe}
import org.apache.toree.kernel.protocol.v5.MessageType
import org.scalatest.{FunSpecLike, Matchers}
import test.utils.TestProbeProxyActor
import test.utils.MaxAkkaTestTimeout

class SimpleActorLoaderSpec extends TestKit(
  ActorSystem(
    "SimpleActorLoaderSpecSystem",
    None,
    Some(org.apache.toree.Main.getClass.getClassLoader)
  )
)
  with FunSpecLike with Matchers
{
  describe("SimpleActorLoader") {
    //val system = ActorSystem("SimpleActorLoaderSystem")
    val testMessage: String = "Hello Message"

    describe("#load( MessageType )") {
      it("should load a MessageType Actor"){
        //  Create a new test probe to verify our selection works
        val messageTypeProbe: TestProbe = new TestProbe(system)

        //  Add an actor to the system to send a message to
        system.actorOf(
          Props(classOf[TestProbeProxyActor], messageTypeProbe),
          name = MessageType.Outgoing.ExecuteInput.toString
        )

        //  Create the ActorLoader with our test system
        val actorLoader: SimpleActorLoader = SimpleActorLoader(system)

        //  Get the actor and send it a message
        val loadedMessageActor: ActorSelection =
          actorLoader.load(MessageType.Outgoing.ExecuteInput)

        loadedMessageActor ! testMessage

        //  Assert the probe received the message
        messageTypeProbe.expectMsg(MaxAkkaTestTimeout, testMessage)
      }
    }

  }
} 

package org.apache.toree.kernel.protocol.v5.kernel.socket

import java.nio.charset.Charset

import akka.actor.{Props, ActorSelection, ActorRef, ActorSystem}
import akka.testkit.{TestProbe, ImplicitSender, TestKit}
import akka.util.ByteString
import org.apache.toree.communication.ZMQMessage
import org.apache.toree.kernel.protocol.v5.kernel.Utilities._
import org.apache.toree.kernel.protocol.v5Test._
import org.apache.toree.kernel.protocol.v5.{KernelMessage, SystemActorType}
import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
import com.typesafe.config.ConfigFactory
import org.scalatest.mock.MockitoSugar
import org.scalatest.{Matchers, FunSpecLike}
import org.mockito.Mockito._
import org.mockito.Matchers._
import test.utils.MaxAkkaTestTimeout

object StdinSpec {
  val config ="""
    akka {
      loglevel = "WARNING"
    }"""
}

class StdinSpec extends TestKit(ActorSystem(
  "StdinActorSpec",
  ConfigFactory.parseString(StdinSpec.config),
  org.apache.toree.Main.getClass.getClassLoader
)) with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {
  describe("Stdin") {
    val socketFactory = mock[SocketFactory]
    val actorLoader = mock[ActorLoader]
    val socketProbe : TestProbe = TestProbe()
    when(socketFactory.Stdin(any(classOf[ActorSystem]), any(classOf[ActorRef]))).thenReturn(socketProbe.ref)

    val relayProbe : TestProbe = TestProbe()
    val relaySelection : ActorSelection = system.actorSelection(relayProbe.ref.path)
    when(actorLoader.load(SystemActorType.KernelMessageRelay)).thenReturn(relaySelection)

    val stdin = system.actorOf(Props(classOf[Stdin], socketFactory, actorLoader))

    describe("#receive") {
      it("( KernelMessage ) should reply with a ZMQMessage via the socket") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        stdin ! MockKernelMessage
        socketProbe.expectMsg(MockZMQMessage)
      }

      it("( ZMQMessage ) should forward ZMQ Strings and KernelMessage to Relay") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        stdin ! MockZMQMessage

        // Should get the last four (assuming no buffer) strings in UTF-8
        val zmqStrings = MockZMQMessage.frames.map((byteString: ByteString) =>
          new String(byteString.toArray, Charset.forName("UTF-8"))
        ).takeRight(4)

        val kernelMessage: KernelMessage = MockZMQMessage

        relayProbe.expectMsg(MaxAkkaTestTimeout, (zmqStrings, kernelMessage))
      }
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.kernel.socket

import java.nio.charset.Charset

import akka.actor.{ActorSelection, ActorRef, ActorSystem, Props}
import akka.testkit.{ImplicitSender, TestKit, TestProbe}
import akka.util.ByteString
import org.apache.toree.communication.ZMQMessage
import org.apache.toree.kernel.protocol.v5._
import org.apache.toree.kernel.protocol.v5.kernel.{ActorLoader, Utilities}
import org.apache.toree.kernel.protocol.v5Test._
import Utilities._
import com.typesafe.config.ConfigFactory
import org.mockito.Matchers._
import org.mockito.Mockito._
import org.scalatest.mock.MockitoSugar
import org.scalatest.{FunSpecLike, Matchers}
import test.utils.MaxAkkaTestTimeout

object ShellSpec {
  val config ="""
    akka {
      loglevel = "WARNING"
    }"""
}

class ShellSpec extends TestKit(
  ActorSystem(
    "ShellActorSpec",
    ConfigFactory.parseString(ShellSpec.config),
    org.apache.toree.Main.getClass.getClassLoader
  ))
  with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {

  describe("Shell") {
    val socketFactory = mock[SocketFactory]
    val actorLoader = mock[ActorLoader]
    val socketProbe : TestProbe = TestProbe()
    when(socketFactory.Shell(any(classOf[ActorSystem]), any(classOf[ActorRef]))).thenReturn(socketProbe.ref)

    val relayProbe : TestProbe = TestProbe()
    val relaySelection : ActorSelection = system.actorSelection(relayProbe.ref.path)
    when(actorLoader.load(SystemActorType.KernelMessageRelay)).thenReturn(relaySelection)

    val shell = system.actorOf(Props(classOf[Shell], socketFactory, actorLoader))

    describe("#receive") {
      it("( KernelMessage ) should reply with a ZMQMessage via the socket") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        shell ! MockKernelMessage
        socketProbe.expectMsg(MockZMQMessage)
      }

      it("( ZMQMessage ) should forward ZMQ Strings and KernelMessage to Relay") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        shell ! MockZMQMessage

        // Should get the last four (assuming no buffer) strings in UTF-8
        val zmqStrings = MockZMQMessage.frames.map((byteString: ByteString) =>
          new String(byteString.toArray, Charset.forName("UTF-8"))
        ).takeRight(4)

        val kernelMessage: KernelMessage = MockZMQMessage

        relayProbe.expectMsg(MaxAkkaTestTimeout, (zmqStrings, kernelMessage))
      }
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.handler
import akka.actor.{ActorSelection, ActorSystem, Props}
import akka.testkit.{ImplicitSender, TestKit, TestProbe}
import org.apache.toree.Main
import org.apache.toree.kernel.protocol.v5.content.KernelInfoReply
import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
import org.apache.toree.kernel.protocol.v5._
import org.mockito.AdditionalMatchers.{not => mockNot}
import org.mockito.Matchers.{eq => mockEq}
import com.typesafe.config.ConfigFactory
import org.mockito.Mockito._
import org.scalatest.mock.MockitoSugar
import org.scalatest.{FunSpecLike, Matchers}
import play.api.libs.json.Json
import test.utils.MaxAkkaTestTimeout

object KernelInfoRequestHandlerSpec {
  val config = """
    akka {
      loglevel = "WARNING"
    }"""
}

class KernelInfoRequestHandlerSpec extends TestKit(
  ActorSystem("KernelInfoRequestHandlerSpec",
    ConfigFactory.parseString(KernelInfoRequestHandlerSpec.config),
    Main.getClass.getClassLoader)
) with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {
  val actorLoader: ActorLoader =  mock[ActorLoader]
  val actor = system.actorOf(Props(classOf[KernelInfoRequestHandler], actorLoader, LanguageInfo("test", "1.0.0", Some(".test"))))

  val relayProbe : TestProbe = TestProbe()
  val relaySelection : ActorSelection =
    system.actorSelection(relayProbe.ref.path)
  when(actorLoader.load(SystemActorType.KernelMessageRelay))
    .thenReturn(relaySelection)
  when(actorLoader.load(mockNot(mockEq(SystemActorType.KernelMessageRelay))))
    .thenReturn(system.actorSelection(""))

  val header = Header("","","","","")
  val kernelMessage = new KernelMessage(
    Seq[Array[Byte]](), "test message", header, header, Metadata(), "{}"
  )

  describe("Kernel Info Request Handler") {
    it("should return a KernelMessage containing kernel info response") {
      actor ! kernelMessage
      val reply = relayProbe.receiveOne(MaxAkkaTestTimeout).asInstanceOf[KernelMessage]
      val kernelInfo = Json.parse(reply.contentString).as[KernelInfoReply]
      kernelInfo.implementation should be ("spark")
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.handler

import akka.actor.{ActorSystem, Props, ActorRef, ActorSelection}
import akka.testkit.{ImplicitSender, TestKit, TestProbe}
import org.apache.toree.kernel.protocol.v5._
import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
import org.apache.toree.kernel.protocol.v5Test._
import org.mockito.Mockito._
import org.scalatest.mock.MockitoSugar
import org.scalatest.{Matchers, FunSpecLike}
import test.utils.MaxAkkaTestTimeout

class GenericSocketMessageHandlerSpec extends TestKit(
  ActorSystem(
    "GenericSocketMessageHandlerSystem",
    None,
    Some(org.apache.toree.Main.getClass.getClassLoader)
  ))
with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {
  describe("GenericSocketMessageHandler( ActorLoader, SocketType )") {
    //  Create a mock ActorLoader for the Relay we are going to test
    val actorLoader: ActorLoader = mock[ActorLoader]

    //  Create a probe for the ActorSelection that the ActorLoader will return
    val selectionProbe: TestProbe = TestProbe()
    val selection: ActorSelection = system.actorSelection(selectionProbe.ref.path.toString)
    when(actorLoader.load(SocketType.Control)).thenReturn(selection)

    //  The Relay we are going to be testing against
    val genericHandler: ActorRef = system.actorOf(
      Props(classOf[GenericSocketMessageHandler], actorLoader, SocketType.Control)
    )

    describe("#receive( KernelMessage )") {
      genericHandler ! MockKernelMessage

      it("should send the message to the selected actor"){
        selectionProbe.expectMsg(MaxAkkaTestTimeout, MockKernelMessage)
      }
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.kernel.socket

import java.nio.charset.Charset

import akka.actor.{ActorSelection, ActorSystem, ActorRef, Actor}
import akka.util.ByteString
import org.apache.toree.communication.ZMQMessage

//import org.apache.toree.kernel.protocol.v5.kernel.ZMQMessage
import org.apache.toree.kernel.protocol.v5.KernelMessage
import org.apache.toree.kernel.protocol.v5.kernel.Utilities._
import org.apache.toree.utils.MessageLogSupport


abstract class ZeromqKernelMessageSocket(
  actorSocketFunc: (ActorSystem, ActorRef) => ActorRef,
  actorForwardFunc: () => ActorSelection
) extends Actor with MessageLogSupport {
  val actorSocketRef = actorSocketFunc(context.system, self)
  val actorForwardRef = actorForwardFunc()

  override def receive: Receive = {
    case message: ZMQMessage =>
      val kernelMessage: KernelMessage = message
      logMessage(kernelMessage)

      // Grab the strings to use for signature verification
      val zmqStrings = message.frames.map((byteString: ByteString) =>
        new String(byteString.toArray, Charset.forName("UTF-8"))
      ).takeRight(4) // TODO: This assumes NO extra buffers, refactor?

      // Forward along our message (along with the strings used for
      // signatures)
      actorForwardRef ! ((zmqStrings, kernelMessage))

    case message: KernelMessage =>
      val zmqMessage: ZMQMessage = message
      logMessage(message)
      actorSocketRef ! zmqMessage
  }
} 

package org.apache.toree.kernel.protocol.v5.client

import akka.actor.{ActorRefFactory, ActorSelection}


  def load(actorEnum: Enumeration#Value): ActorSelection
}

case class SimpleActorLoader(actorRefFactory: ActorRefFactory)
  extends ActorLoader
{
  private val userActorDirectory: String = "/user/%s"

  override def load(actorEnum: Enumeration#Value): ActorSelection = {
    actorRefFactory.actorSelection(
      userActorDirectory.format(actorEnum.toString)
    )
  }
} 

package org.squbs.pattern.orchestration.japi

import java.util.concurrent.CompletableFuture

import akka.actor.{Actor, ActorRef, ActorSelection}
import akka.pattern.{AskableActorRef, AskableActorSelection}
import akka.util.Timeout
import org.squbs.pattern.orchestration.Orchestrator

import scala.concurrent.Future
import scala.util.{Failure, Success, Try}

abstract class AbstractOrchestrator extends Actor with Orchestrator {

  override def receive = super.receive

  def ask(actor: ActorRef, message: Any, timeout: Timeout): Ask = {
    val future = new AskableActorRef(actor).ask(message)(timeout)
    new Ask(future)
  }

  def ask(actor: ActorSelection, message: Any, timeout: Timeout): Ask = {
    val future = new AskableActorSelection(actor).ask(message)(timeout)
    new Ask(future)
  }

  class Ask private[japi](private val future: Future[Any]) {

    def thenComplete[T](cFuture: CompletableFuture[T]): Unit = {
      // Dragons here: DO NOT call nextMessageId from inside future.onComplete as that executes
      // outside the context of the actor. Instead, obtain the (val) id eagerly inside the actor and
      // give it to the function so it becomes pre-assigned.
      val nextId = nextMessageId
      import context.dispatcher
      future onComplete { self ! UniqueTryWrapper(nextId, _) }
      expectOnce {
        case UniqueTryWrapper(`nextId`, tt: Try[_]) =>
          tt match {
            case Success(t) => cFuture.complete(t.asInstanceOf[T])
            case Failure(e) => cFuture.completeExceptionally(e)
          }
      }
    }
  }
} 

import akka.actor.{ Actor, ActorSelection }
import akka.persistence.{ AtLeastOnceDelivery, PersistentActor }

sealed trait Cmd
case class SayHello(deliveryId: Long, s: String) extends Cmd
case class ReceiveHello(deliveryId: Long) extends Cmd

sealed trait Evt
case class HelloSaid(s: String) extends Evt
case class HelloReceived(deliveryId: Long) extends Evt


class SendActor(destination: ActorSelection)
    extends PersistentActor with AtLeastOnceDelivery {


  override def persistenceId: String = "persistence-id"


  override def receiveCommand: Receive = {
    case s: String =>
      persist(HelloSaid(s))(updateState)
    case ReceiveHello(deliveryId) =>
      persist(HelloReceived(deliveryId))(updateState)
  }


  override def receiveRecover: Receive = {
    case evt: Evt => updateState(evt)
  }

  def updateState(evt: Evt): Unit = evt match {
    case HelloSaid(s) =>
      deliver(destination)(deliveryId => SayHello(deliveryId, s))

    case HelloReceived(deliveryId) =>
      confirmDelivery(deliveryId)
  }
}

class ReceiveActor extends Actor {
  def receive = {
    case SayHello(deliveryId, s) =>
      // ... do something with s
      sender() ! ReceiveHello(deliveryId)
  }
} 

package com.rarebooks.library

import akka.actor.{ActorSelection, ActorSystem, Address, RootActorPath}
import akka.event.Logging

import scala.annotation.tailrec
import scala.concurrent.Await
import scala.concurrent.duration.{Duration, FiniteDuration, MILLISECONDS => Millis}
import scala.io.StdIn

object CustomerApp {

  
  protected def createCustomer(count: Int, odds: Int, tolerance: Int): Unit = {
    val selection: ActorSelection  =
        system.actorSelection(
        RootActorPath(rareBooksAddress) /
        "user" / "rare-books")

    selection.resolveOne(resolveTimeout).onComplete {
      case scala.util.Success(rareBooks) =>
        for (_ <- 1 to count)
          system.actorOf(Customer.props(rareBooks, odds, tolerance))
      case scala.util.Failure(ex) =>
        log.error(ex, ex.getMessage)
    }
  }
} 

package cmwell.ctrl.controllers

import akka.actor.ActorSelection
import cmwell.ctrl.controllers.CassandraController._
import cmwell.ctrl.config.Config
import cmwell.ctrl.utils.ProcUtil
import com.typesafe.scalalogging.LazyLogging
import k.grid.Grid

import scala.concurrent.{blocking, Future}
import scala.util.{Failure, Success}
import scala.concurrent.duration._
import scala.concurrent.ExecutionContext.Implicits.global


abstract class ComponentController(startScriptLocation: String, psIdentifier: String, dirIdentifier: Set[String]) {
  object ComponentControllerLogger extends LazyLogging {
    lazy val l = logger
  }

  protected val startScriptPattern: String = "start[0-9]*.sh"

  def getStartScriptLocation = startScriptLocation

  def getStartScripts(location: String): Set[String] = {
    ProcUtil.executeCommand(s"ls -1 $location/ | grep $startScriptPattern") match {
      case Success(str) =>
        str.trim.split("\n").toSet
      case Failure(err) => Set.empty[String]
    }
  }

  def getDataDirs(location: String, id: String): Set[String] = {
    ProcUtil.executeCommand(s"ls -1 $location | grep $id[0-9]*") match {
      case Success(str) =>
        str.trim.split("\n").toSet
      case Failure(err) => Set.empty[String]
    }
  }

  private def doStart: Unit = {
    getStartScripts(startScriptLocation).foreach { sScript =>
      val runScript = s"HAL=9000 $startScriptLocation/$sScript"
      ProcUtil.executeCommand(runScript)
    }
  }

  def start {
    blocking {
      Future {
        doStart
      }
    }
  }

  private def doStop(forceKill: Boolean = false, tries: Int = 5): Unit = {
    val cmd =
      s"ps aux | grep $psIdentifier | egrep -v 'grep|starter' | awk '{print $$2}' | xargs kill ${if (forceKill) "-9"
      else ""}"
    ComponentControllerLogger.l.info(s"executing $cmd")
    ProcUtil.executeCommand(cmd)
    val isDead =
      ProcUtil.executeCommand(s"ps aux | grep $psIdentifier | egrep -v 'grep|starter' | awk '{print $$2}'").get.isEmpty
    if (!isDead) {
      if (tries > 1) doStop(false, tries - 1) else doStop(true, tries - 1)
    }

  }
  def stop {
    Future {
      blocking {
        doStop()
      }
    }
  }

  def restart: Unit = {
    Future {
      blocking {
        doStop()
        doStart
      }
    }
  }

  def clearData {
    Future {
      blocking {
        dirIdentifier.foreach { id =>
          getDataDirs(s"${Config.cmwellHome}/data/", id).foreach { dir =>
            ProcUtil.executeCommand(s"rm -rf ${Config.cmwellHome}/data/$dir/")
          }
        }
      }
    }
  }
} 

package cmwell.ctrl.server

import akka.actor.{Actor, ActorLogging, ActorRef, ActorSelection}
import akka.pattern.pipe
import cmwell.ctrl.checkers._
import cmwell.ctrl.commands.{ControlCommand, StartElasticsearchMaster}
import cmwell.ctrl.config.{Config, Jvms}
import cmwell.ctrl.utils.ProcUtil
import k.grid.{Grid, GridJvm, GridJvm$}
import scala.concurrent.Future
import scala.sys.process._
import Config._

import scala.util.{Failure, Success, Try}
import scala.concurrent.ExecutionContext.Implicits.global


case class BashCommand(com: String)

object CommandActor {
  def select(host: String): ActorSelection = {
    Grid.selectActor(commandActorName, GridJvm(host, Some(Jvms.node)))
  }

  def all: Set[ActorSelection] = {
    Grid.availableMachines.map(host => Grid.selectActor(commandActorName, GridJvm(host, Some(Jvms.node))))
  }
}

class CommandActor extends Actor with ActorLogging {

  override def receive: Receive = {
    case BashCommand(com)   => sender ! ProcUtil.executeCommand(com)
    case CheckWeb           => WebChecker.check.pipeTo(sender())
    case CheckElasticsearch => ElasticsearchChecker.check.pipeTo(sender())
    case CheckCassandra     => CassandraChecker.check.pipeTo(sender())
    case cc: ControlCommand => cc.execute
  }
} 

package hydra.ingest.bootstrap

import akka.actor.{ActorSelection, ActorSystem}
import akka.pattern.ask
import akka.util.Timeout
import com.typesafe.config.Config
import hydra.common.util.ActorUtils
import hydra.ingest.services.IngestorRegistry
import hydra.ingest.services.IngestorRegistry.{FindByName, LookupResult}

import scala.concurrent.Future


class HydraIngestorRegistryClient(registryPath: String)(
    implicit val system: ActorSystem
) {

  lazy val registry: ActorSelection = system.actorSelection(registryPath)

  def lookupIngestor(
      name: String
  )(implicit timeout: Timeout): Future[LookupResult] = {
    (registry ? FindByName(name)).mapTo[LookupResult]
  }
}

object HydraIngestorRegistryClient {

  import hydra.common.config.ConfigSupport._

  def registryPath(config: Config) =
    config
      .getStringOpt("ingest.ingestor-registry.path")
      .getOrElse(
        s"/user/service/${ActorUtils.actorName(classOf[IngestorRegistry])}"
      )

  def apply(
      config: Config
  )(implicit system: ActorSystem): HydraIngestorRegistryClient = {
    new HydraIngestorRegistryClient(registryPath(config))(system)
  }
} 

package cmwell.rts

import akka.actor.{Actor, ActorRef, ActorSelection}
import cmwell.domain.Infoton
import k.grid.Grid
import com.typesafe.scalalogging.LazyLogging


case class AddSubscriber(subscriber: String, rule: Rule)
case class RemoveSubscriber(subscriber: String)

case class Publish(i: Vector[Infoton])
case class PublishOne(uuid: String)
//case class Publish(s : String)

class PublishAgent extends Actor with LazyLogging {
  // need a mapping rule (q) -> actor (label)
  var rules: Map[String, Rule] = Map.empty[String, Rule]
  var subMap: Map[String, ActorSelection] = Map.empty[String, ActorSelection]

  def publish_data(subscriber: String, i: Infoton) {
    // send
    val a = subMap(subscriber)
    logger.debug(s"Send data to $subscriber [$a]")
    a ! PublishOne(i.uuid)
  }

  def receive = {
    // add a rule to the internal map
    case AddSubscriber(subscriber: String, rule: Rule) =>
      val addr = sender().path.address
      val path = s"akka.tcp://${addr.system}@${addr.host.getOrElse("")}:${addr.port.getOrElse(0)}/user/$subscriber"
      rules += (subscriber -> rule)
      subMap += (subscriber -> context.actorSelection(path))
      logger.debug(s"AddRule rules [${rules}] sub map [${subMap}]")
    // remove the rule from the internal map
    case RemoveSubscriber(subscriber: String) =>
      rules -= (subscriber)
      subMap -= (subscriber)
      logger.debug(s"RemoveRule ${subscriber} rules [${rules}] sub map [${subMap}]")

    // this publish the infoton according the rule
    case Publish(infotonVec: Vector[Infoton]) => {
      logger.debug(s"in actor $infotonVec")
      // first lets calc
      infotonVec.foreach { i =>
        rules.foreach {
          case (subscriber, rule) =>
            rule match {
              case NoFilter =>
                publish_data(subscriber, i)
              case PathFilter(path) =>
                if (path.check(i.systemFields.path))
                  publish_data(subscriber, i)
              case MatchFilter(f) =>
                if (i.fields.isDefined && f.check(i.fields.get))
                  publish_data(subscriber, i)
              case PMFilter(p, m) =>
                if (p.check(i.systemFields.path) && i.fields.isDefined && m.check(i.fields.get))
                  publish_data(subscriber, i)
            }
        }
      }
    }
    case _ =>
      logger.debug("Error")
  }

}

object Publisher {
  val publishAgentActor: ActorRef = Grid.create(classOf[PublishAgent], "publisher")
  val p: Publisher = new Publisher(publishAgentActor)
  def init: Unit = {}
  def publish(i: Vector[Infoton]): Unit = p.publish(i)
}

class Publisher(val publishAgentActor: ActorRef) {

  def publish(i: Vector[Infoton]): Unit = {
    // no block call here
    publishAgentActor ! Publish(i)
  }

} 

package com.olegych.scastie.sbt

import com.olegych.scastie.api._
import com.olegych.scastie.util._
import akka.actor.{Actor, ActorContext, ActorLogging, ActorRef, ActorSelection, ActorSystem, Props}

import scala.concurrent.duration._

case object SbtActorReady

class SbtActor(system: ActorSystem,
               runTimeout: FiniteDuration,
               sbtReloadTimeout: FiniteDuration,
               isProduction: Boolean,
               readyRef: Option[ActorRef],
               override val reconnectInfo: Option[ReconnectInfo])
    extends Actor
    with ActorLogging
    with ActorReconnecting {

  def balancer(context: ActorContext, info: ReconnectInfo): ActorSelection = {
    import info._
    context.actorSelection(
      s"akka.tcp://Web@$serverHostname:$serverAkkaPort/user/DispatchActor"
    )
  }

  override def tryConnect(context: ActorContext): Unit = {
    if (isProduction) {
      reconnectInfo.foreach { info =>
        import info._
        balancer(context, info) ! SbtRunnerConnect(actorHostname, actorAkkaPort)
      }
    }
  }

  override def preStart(): Unit = {
    log.info("*** SbtRunner preStart ***")

    readyRef.foreach(_ ! SbtActorReady)
    super.preStart()
  }

  override def postStop(): Unit = {
    log.info("*** SbtRunner postStop ***")

    super.postStop()
  }

  private val formatActor =
    context.actorOf(Props(new FormatActor()), name = "FormatActor")

  private val sbtRunner =
    context.actorOf(
      Props(
        new SbtProcess(
          runTimeout,
          sbtReloadTimeout,
          isProduction,
          javaOptions = Seq("-Xms512m", "-Xmx1g")
        )
      ),
      name = "SbtRunner"
    )

  override def receive: Receive = reconnectBehavior orElse [Any, Unit] {
    case SbtPing => {
      sender ! SbtPong
    }

    case format: FormatRequest => {
      formatActor.forward(format)
    }

    case task: SbtTask => {
      sbtRunner.forward(task)
    }

    case SbtUp => {
      log.info("SbtUp")
      reconnectInfo.foreach { info =>
        log.info("SbtUp sent")
        balancer(context, info) ! SbtUp
      }
    }

    case replay: Replay => {
      log.info("Replay")
      reconnectInfo.foreach { info =>
        log.info("Replay sent")
        balancer(context, info) ! replay
      }
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.kernel

import akka.actor.{ActorSelection, ActorSystem, Props}
import akka.testkit.{TestKit, TestProbe}
import org.apache.toree.kernel.protocol.v5.MessageType
import org.scalatest.{FunSpecLike, Matchers}
import test.utils.TestProbeProxyActor
import test.utils.MaxAkkaTestTimeout

class SimpleActorLoaderSpec extends TestKit(
  ActorSystem(
    "SimpleActorLoaderSpecSystem",
    None,
    Some(org.apache.toree.Main.getClass.getClassLoader)
  )
)
  with FunSpecLike with Matchers
{
  describe("SimpleActorLoader") {
    //val system = ActorSystem("SimpleActorLoaderSystem")
    val testMessage: String = "Hello Message"

    describe("#load( MessageType )") {
      it("should load a MessageType Actor"){
        //  Create a new test probe to verify our selection works
        val messageTypeProbe: TestProbe = new TestProbe(system)

        //  Add an actor to the system to send a message to
        system.actorOf(
          Props(classOf[TestProbeProxyActor], messageTypeProbe),
          name = MessageType.Outgoing.ExecuteInput.toString
        )

        //  Create the ActorLoader with our test system
        val actorLoader: SimpleActorLoader = SimpleActorLoader(system)

        //  Get the actor and send it a message
        val loadedMessageActor: ActorSelection =
          actorLoader.load(MessageType.Outgoing.ExecuteInput)

        loadedMessageActor ! testMessage

        //  Assert the probe received the message
        messageTypeProbe.expectMsg(MaxAkkaTestTimeout, testMessage)
      }
    }

  }
} 

package org.apache.toree.kernel.protocol.v5.kernel.socket

import java.nio.charset.Charset

import akka.actor.{Props, ActorSelection, ActorRef, ActorSystem}
import akka.testkit.{TestProbe, ImplicitSender, TestKit}
import akka.util.ByteString
import org.apache.toree.communication.ZMQMessage
import org.apache.toree.kernel.protocol.v5.kernel.Utilities._
import org.apache.toree.kernel.protocol.v5Test._
import org.apache.toree.kernel.protocol.v5.{KernelMessage, SystemActorType}
import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
import com.typesafe.config.ConfigFactory
import org.scalatest.mock.MockitoSugar
import org.scalatest.{Matchers, FunSpecLike}
import org.mockito.Mockito._
import org.mockito.Matchers._
import test.utils.MaxAkkaTestTimeout

object StdinSpec {
  val config ="""
    akka {
      loglevel = "WARNING"
    }"""
}

class StdinSpec extends TestKit(ActorSystem(
  "StdinActorSpec",
  ConfigFactory.parseString(StdinSpec.config),
  org.apache.toree.Main.getClass.getClassLoader
)) with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {
  describe("Stdin") {
    val socketFactory = mock[SocketFactory]
    val actorLoader = mock[ActorLoader]
    val socketProbe : TestProbe = TestProbe()
    when(socketFactory.Stdin(any(classOf[ActorSystem]), any(classOf[ActorRef]))).thenReturn(socketProbe.ref)

    val relayProbe : TestProbe = TestProbe()
    val relaySelection : ActorSelection = system.actorSelection(relayProbe.ref.path)
    when(actorLoader.load(SystemActorType.KernelMessageRelay)).thenReturn(relaySelection)

    val stdin = system.actorOf(Props(classOf[Stdin], socketFactory, actorLoader))

    describe("#receive") {
      it("( KernelMessage ) should reply with a ZMQMessage via the socket") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        stdin ! MockKernelMessage
        socketProbe.expectMsg(MockZMQMessage)
      }

      it("( ZMQMessage ) should forward ZMQ Strings and KernelMessage to Relay") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        stdin ! MockZMQMessage

        // Should get the last four (assuming no buffer) strings in UTF-8
        val zmqStrings = MockZMQMessage.frames.map((byteString: ByteString) =>
          new String(byteString.toArray, Charset.forName("UTF-8"))
        ).takeRight(4)

        val kernelMessage: KernelMessage = MockZMQMessage

        relayProbe.expectMsg(MaxAkkaTestTimeout, (zmqStrings, kernelMessage))
      }
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.kernel.socket

import java.nio.charset.Charset

import akka.actor.{ActorSelection, ActorRef, ActorSystem, Props}
import akka.testkit.{ImplicitSender, TestKit, TestProbe}
import akka.util.ByteString
import org.apache.toree.communication.ZMQMessage
import org.apache.toree.kernel.protocol.v5._
import org.apache.toree.kernel.protocol.v5.kernel.{ActorLoader, Utilities}
import org.apache.toree.kernel.protocol.v5Test._
import Utilities._
import com.typesafe.config.ConfigFactory
import org.mockito.Matchers._
import org.mockito.Mockito._
import org.scalatest.mock.MockitoSugar
import org.scalatest.{FunSpecLike, Matchers}
import test.utils.MaxAkkaTestTimeout

object ShellSpec {
  val config ="""
    akka {
      loglevel = "WARNING"
    }"""
}

class ShellSpec extends TestKit(
  ActorSystem(
    "ShellActorSpec",
    ConfigFactory.parseString(ShellSpec.config),
    org.apache.toree.Main.getClass.getClassLoader
  ))
  with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {

  describe("Shell") {
    val socketFactory = mock[SocketFactory]
    val actorLoader = mock[ActorLoader]
    val socketProbe : TestProbe = TestProbe()
    when(socketFactory.Shell(any(classOf[ActorSystem]), any(classOf[ActorRef]))).thenReturn(socketProbe.ref)

    val relayProbe : TestProbe = TestProbe()
    val relaySelection : ActorSelection = system.actorSelection(relayProbe.ref.path)
    when(actorLoader.load(SystemActorType.KernelMessageRelay)).thenReturn(relaySelection)

    val shell = system.actorOf(Props(classOf[Shell], socketFactory, actorLoader))

    describe("#receive") {
      it("( KernelMessage ) should reply with a ZMQMessage via the socket") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        shell ! MockKernelMessage
        socketProbe.expectMsg(MockZMQMessage)
      }

      it("( ZMQMessage ) should forward ZMQ Strings and KernelMessage to Relay") {
        //  Use the implicit to convert the KernelMessage to ZMQMessage
        val MockZMQMessage : ZMQMessage = MockKernelMessage

        shell ! MockZMQMessage

        // Should get the last four (assuming no buffer) strings in UTF-8
        val zmqStrings = MockZMQMessage.frames.map((byteString: ByteString) =>
          new String(byteString.toArray, Charset.forName("UTF-8"))
        ).takeRight(4)

        val kernelMessage: KernelMessage = MockZMQMessage

        relayProbe.expectMsg(MaxAkkaTestTimeout, (zmqStrings, kernelMessage))
      }
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.handler
import akka.actor.{ActorSelection, ActorSystem, Props}
import akka.testkit.{ImplicitSender, TestKit, TestProbe}
import org.apache.toree.Main
import org.apache.toree.kernel.protocol.v5.content.KernelInfoReply
import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
import org.apache.toree.kernel.protocol.v5._
import org.mockito.AdditionalMatchers.{not => mockNot}
import org.mockito.Matchers.{eq => mockEq}
import com.typesafe.config.ConfigFactory
import org.mockito.Mockito._
import org.scalatest.mock.MockitoSugar
import org.scalatest.{FunSpecLike, Matchers}
import play.api.libs.json.Json
import test.utils.MaxAkkaTestTimeout

object KernelInfoRequestHandlerSpec {
  val config = """
    akka {
      loglevel = "WARNING"
    }"""
}

class KernelInfoRequestHandlerSpec extends TestKit(
  ActorSystem("KernelInfoRequestHandlerSpec",
    ConfigFactory.parseString(KernelInfoRequestHandlerSpec.config),
    Main.getClass.getClassLoader)
) with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {
  val actorLoader: ActorLoader =  mock[ActorLoader]
  val actor = system.actorOf(Props(classOf[KernelInfoRequestHandler], actorLoader, LanguageInfo("test", "1.0.0", Some(".test"))))

  val relayProbe : TestProbe = TestProbe()
  val relaySelection : ActorSelection =
    system.actorSelection(relayProbe.ref.path)
  when(actorLoader.load(SystemActorType.KernelMessageRelay))
    .thenReturn(relaySelection)
  when(actorLoader.load(mockNot(mockEq(SystemActorType.KernelMessageRelay))))
    .thenReturn(system.actorSelection(""))

  val header = Header("","","","","")
  val kernelMessage = new KernelMessage(
    Seq[Array[Byte]](), "test message", header, header, Metadata(), "{}"
  )

  describe("Kernel Info Request Handler") {
    it("should return a KernelMessage containing kernel info response") {
      actor ! kernelMessage
      val reply = relayProbe.receiveOne(MaxAkkaTestTimeout).asInstanceOf[KernelMessage]
      val kernelInfo = Json.parse(reply.contentString).as[KernelInfoReply]
      kernelInfo.implementation should be ("spark")
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.handler

import akka.actor.{ActorSystem, Props, ActorRef, ActorSelection}
import akka.testkit.{ImplicitSender, TestKit, TestProbe}
import org.apache.toree.kernel.protocol.v5._
import org.apache.toree.kernel.protocol.v5.kernel.ActorLoader
import org.apache.toree.kernel.protocol.v5Test._
import org.mockito.Mockito._
import org.scalatest.mock.MockitoSugar
import org.scalatest.{Matchers, FunSpecLike}
import test.utils.MaxAkkaTestTimeout

class GenericSocketMessageHandlerSpec extends TestKit(
  ActorSystem(
    "GenericSocketMessageHandlerSystem",
    None,
    Some(org.apache.toree.Main.getClass.getClassLoader)
  ))
with ImplicitSender with FunSpecLike with Matchers with MockitoSugar {
  describe("GenericSocketMessageHandler( ActorLoader, SocketType )") {
    //  Create a mock ActorLoader for the Relay we are going to test
    val actorLoader: ActorLoader = mock[ActorLoader]

    //  Create a probe for the ActorSelection that the ActorLoader will return
    val selectionProbe: TestProbe = TestProbe()
    val selection: ActorSelection = system.actorSelection(selectionProbe.ref.path.toString)
    when(actorLoader.load(SocketType.Control)).thenReturn(selection)

    //  The Relay we are going to be testing against
    val genericHandler: ActorRef = system.actorOf(
      Props(classOf[GenericSocketMessageHandler], actorLoader, SocketType.Control)
    )

    describe("#receive( KernelMessage )") {
      genericHandler ! MockKernelMessage

      it("should send the message to the selected actor"){
        selectionProbe.expectMsg(MaxAkkaTestTimeout, MockKernelMessage)
      }
    }
  }
} 

package org.apache.toree.kernel.protocol.v5.kernel.socket

import java.nio.charset.Charset

import akka.actor.{ActorSelection, ActorSystem, ActorRef, Actor}
import akka.util.ByteString
import org.apache.toree.communication.ZMQMessage

//import org.apache.toree.kernel.protocol.v5.kernel.ZMQMessage
import org.apache.toree.kernel.protocol.v5.KernelMessage
import org.apache.toree.kernel.protocol.v5.kernel.Utilities._
import org.apache.toree.utils.MessageLogSupport


abstract class ZeromqKernelMessageSocket(
  actorSocketFunc: (ActorSystem, ActorRef) => ActorRef,
  actorForwardFunc: () => ActorSelection
) extends Actor with MessageLogSupport {
  val actorSocketRef = actorSocketFunc(context.system, self)
  val actorForwardRef = actorForwardFunc()

  override def receive: Receive = {
    case message: ZMQMessage =>
      val kernelMessage: KernelMessage = message
      logMessage(kernelMessage)

      // Grab the strings to use for signature verification
      val zmqStrings = message.frames.map((byteString: ByteString) =>
        new String(byteString.toArray, Charset.forName("UTF-8"))
      ).takeRight(4) // TODO: This assumes NO extra buffers, refactor?

      // Forward along our message (along with the strings used for
      // signatures)
      actorForwardRef ! ((zmqStrings, kernelMessage))

    case message: KernelMessage =>
      val zmqMessage: ZMQMessage = message
      logMessage(message)
      actorSocketRef ! zmqMessage
  }
} 

package org.apache.toree.kernel.protocol.v5.client

import akka.actor.{ActorRefFactory, ActorSelection}


  def load(actorEnum: Enumeration#Value): ActorSelection
}

case class SimpleActorLoader(actorRefFactory: ActorRefFactory)
  extends ActorLoader
{
  private val userActorDirectory: String = "/user/%s"

  override def load(actorEnum: Enumeration#Value): ActorSelection = {
    actorRefFactory.actorSelection(
      userActorDirectory.format(actorEnum.toString)
    )
  }
} 

package akka.pattern

import language.implicitConversions
import scala.concurrent.{ Future, ExecutionContext }
import scala.util.{ Failure, Success }
import akka.actor.{ Status, ActorRef, Actor }
import akka.actor.ActorSelection
import java.util.concurrent.CompletionStage
import java.util.function.BiConsumer

trait PipeToSupport {

  final class PipeableFuture[T](val future: Future[T])(implicit executionContext: ExecutionContext) {
    def pipeTo(recipient: ActorRef)(implicit sender: ActorRef = Actor.noSender): Future[T] = {
      future andThen {
        case Success(r) ⇒ recipient ! r
        case Failure(f) ⇒ recipient ! Status.Failure(f)
      }
    }
    def pipeToSelection(recipient: ActorSelection)(implicit sender: ActorRef = Actor.noSender): Future[T] = {
      future andThen {
        case Success(r) ⇒ recipient ! r
        case Failure(f) ⇒ recipient ! Status.Failure(f)
      }
    }
    def to(recipient: ActorRef): PipeableFuture[T] = to(recipient, Actor.noSender)
    def to(recipient: ActorRef, sender: ActorRef): PipeableFuture[T] = {
      pipeTo(recipient)(sender)
      this
    }
    def to(recipient: ActorSelection): PipeableFuture[T] = to(recipient, Actor.noSender)
    def to(recipient: ActorSelection, sender: ActorRef): PipeableFuture[T] = {
      pipeToSelection(recipient)(sender)
      this
    }
  }

  final class PipeableCompletionStage[T](val future: CompletionStage[T])(implicit executionContext: ExecutionContext) {
    def pipeTo(recipient: ActorRef)(implicit sender: ActorRef = Actor.noSender): CompletionStage[T] = {
      future whenComplete new BiConsumer[T, Throwable] {
        override def accept(t: T, ex: Throwable) {
          if (t != null) recipient ! t
          if (ex != null) recipient ! Status.Failure(ex)
        }
      }
    }
    def pipeToSelection(recipient: ActorSelection)(implicit sender: ActorRef = Actor.noSender): CompletionStage[T] = {
      future whenComplete new BiConsumer[T, Throwable] {
        override def accept(t: T, ex: Throwable) {
          if (t != null) recipient ! t
          if (ex != null) recipient ! Status.Failure(ex)
        }
      }
    }
    def to(recipient: ActorRef): PipeableCompletionStage[T] = to(recipient, Actor.noSender)
    def to(recipient: ActorRef, sender: ActorRef): PipeableCompletionStage[T] = {
      pipeTo(recipient)(sender)
      this
    }
    def to(recipient: ActorSelection): PipeableCompletionStage[T] = to(recipient, Actor.noSender)
    def to(recipient: ActorSelection, sender: ActorRef): PipeableCompletionStage[T] = {
      pipeToSelection(recipient)(sender)
      this
    }
  }

  
  implicit def pipeCompletionStage[T](future: CompletionStage[T])(implicit executionContext: ExecutionContext): PipeableCompletionStage[T] = new PipeableCompletionStage(future)
} 

package hydra.kafka.endpoints

import akka.actor.ActorSelection
import akka.http.scaladsl.common.EntityStreamingSupport
import akka.kafka.Subscriptions
import akka.kafka.scaladsl.Consumer
import akka.pattern.ask
import akka.util.Timeout
import hydra.core.http.RouteSupport
import hydra.kafka.consumer.KafkaConsumerProxy.{GetLatestOffsets, LatestOffsetsResponse}
import org.apache.kafka.clients.consumer.ConsumerRecord
import org.apache.kafka.common.TopicPartition

import scala.collection.immutable.Map
import scala.concurrent.duration._
import scala.concurrent.{Await, ExecutionContext, Future}


class TopicsEndpoint(consumerProxy:ActorSelection)(implicit ec:ExecutionContext) extends RouteSupport {

  import hydra.kafka.util.KafkaUtils._

  implicit val jsonStreamingSupport = EntityStreamingSupport.json()

  override val route =
    path("transports" / "kafka" / "consumer" / "topics" / Segment) {
      topicName =>
        get {
          extractRequestContext { ctx =>
            parameters('format.?, 'group.?, 'n ? 10, 'start ? "earliest") {
              (format, groupId, n, startOffset) =>
                val settings = loadConsumerSettings[Any, Any](
                  format.getOrElse("avro"),
                  groupId.getOrElse("hydra"),
                  startOffset
                )
                val offsets = latestOffsets(topicName)
                val source = Consumer
                  .plainSource(settings, Subscriptions.topics(topicName))
                  .initialTimeout(5.seconds)
                  .zipWithIndex
                  .takeWhile(rec =>
                    rec._2 <= n && !shouldCancel(offsets, rec._1)
                  )
                  .map(rec => rec._1.value().toString)
                  .watchTermination()((_, termination) =>
                    termination.failed.foreach {
                      case cause => ctx.fail(cause)
                    }
                  )
                complete(source)

            }
          }
        }
    }

  def shouldCancel(
      fpartitions: Future[Map[TopicPartition, Long]],
      record: ConsumerRecord[Any, Any]
  ): Boolean = {
    if (fpartitions.isCompleted) {
      val partitions = Await.result(fpartitions, 1.millis)
      val tp = new TopicPartition(record.topic(), record.partition())
      partitions.get(tp) match {
        case Some(offset) => record.offset() >= offset
        case None         => false
      }
    } else {
      false
    }

  }

  private def latestOffsets(
      topic: String
  ): Future[Map[TopicPartition, Long]] = {
    implicit val timeout = Timeout(5 seconds)
    (consumerProxy ? GetLatestOffsets(topic))
      .mapTo[LatestOffsetsResponse]
      .map(_.offsets)
  }

} 

package hydra.kafka.producer

import akka.actor.ActorSelection
import hydra.core.monitor.HydraMetrics
import hydra.core.transport.TransportCallback
import hydra.kafka.transport.KafkaTransport.RecordProduceError
import org.apache.kafka.clients.producer.{Callback, RecordMetadata}


case class HydraKafkaCallback(
    deliveryId: Long,
    record: KafkaRecord[_, _],
    producer: ActorSelection,
    callback: TransportCallback
) extends Callback {

  override def onCompletion(metadata: RecordMetadata, e: Exception): Unit = {
    Option(e) match {
      case Some(err) => ackError(err)
      case None      => doAck(metadata)
    }
  }

  private def doAck(md: RecordMetadata) = {
    val kmd = KafkaRecordMetadata(md, deliveryId, record.ackStrategy)
    producer ! kmd
    callback.onCompletion(deliveryId, Some(kmd), None)
  }

  private def ackError(e: Exception) = {
    producer ! RecordProduceError(deliveryId, record, e)
    callback.onCompletion(deliveryId, None, Some(e))
  }
}