275 lines
11 KiB
Scala
275 lines
11 KiB
Scala
|
/**
|
||
|
|
* Copyright (C) 2014-2017 Lightbend Inc. <http://www.lightbend.com>
|
||
|
|
*/
|
||
|
|
package akka.stream.remote.scaladsl
|
||
|
|
|
||
|
|
import java.util.Queue
|
||
|
|
|
||
|
|
import akka.actor.{ ActorRef, Terminated }
|
||
|
|
import akka.event.Logging
|
||
|
|
import akka.stream._
|
||
|
|
import akka.stream.actor.{ MaxInFlightRequestStrategy, RequestStrategy, WatermarkRequestStrategy }
|
||
|
|
import akka.stream.impl.FixedSizeBuffer
|
||
|
|
import akka.stream.remote.StreamRefs
|
||
|
|
import akka.stream.remote.impl.StreamRefsMaster
|
||
|
|
import akka.stream.scaladsl.Source
|
||
|
|
import akka.stream.stage._
|
||
|
|
import akka.util.ByteString
|
||
|
|
|
||
|
|
import scala.concurrent.{ Future, Promise }
|
||
|
|
|
||
|
|
object SinkRef {
|
||
|
|
def source[T](): Source[T, Future[SinkRef[T]]] =
|
||
|
|
Source.fromGraph(new SinkRefTargetSource[T]()) // TODO settings?
|
||
|
|
|
||
|
|
def bulkTransferSource(port: Int = -1): Source[ByteString, SinkRef[ByteString]] = {
|
||
|
|
???
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* This stage can only handle a single "sender" (it does not merge values);
|
||
|
|
* The first that pushes is assumed the one we are to trust
|
||
|
|
*/
|
||
|
|
final class SinkRefTargetSource[T]() extends GraphStageWithMaterializedValue[SourceShape[T], Future[SinkRef[T]]] {
|
||
|
|
val out: Outlet[T] = Outlet[T](s"${Logging.simpleName(getClass)}.out")
|
||
|
|
override def shape = SourceShape.of(out)
|
||
|
|
|
||
|
|
override def createLogicAndMaterializedValue(inheritedAttributes: Attributes) = {
|
||
|
|
val promise = Promise[SinkRef[T]]()
|
||
|
|
|
||
|
|
val logic = new TimerGraphStageLogic(shape) with StageLogging with OutHandler {
|
||
|
|
private[this] lazy val streamRefsMaster = StreamRefsMaster(ActorMaterializerHelper.downcast(materializer).system)
|
||
|
|
private[this] lazy val settings = streamRefsMaster.settings
|
||
|
|
|
||
|
|
private[this] var self: GraphStageLogic.StageActor = _
|
||
|
|
private[this] lazy val selfActorName = streamRefsMaster.nextSinkRefTargetSourceName()
|
||
|
|
private[this] implicit def selfSender: ActorRef = self.ref
|
||
|
|
|
||
|
|
// demand management ---
|
||
|
|
private val highDemandWatermark = 16
|
||
|
|
|
||
|
|
private var expectingSeqNr: Long = 0L
|
||
|
|
private var localCumulativeDemand: Long = 0L // initialized in preStart with settings.initialDemand
|
||
|
|
|
||
|
|
private val receiveBuffer = FixedSizeBuffer[T](highDemandWatermark)
|
||
|
|
|
||
|
|
// TODO configurable?
|
||
|
|
// Request strategies talk in terms of Request(n), which we need to translate to cumulative demand
|
||
|
|
// TODO the MaxInFlightRequestStrategy is likely better for this use case, yet was a bit weird to use so this one for now
|
||
|
|
private val requestStrategy: RequestStrategy = WatermarkRequestStrategy(highWatermark = highDemandWatermark)
|
||
|
|
// end of demand management ---
|
||
|
|
|
||
|
|
private var remotePartner: ActorRef = _
|
||
|
|
|
||
|
|
override def preStart(): Unit = {
|
||
|
|
localCumulativeDemand = settings.initialDemand.toLong
|
||
|
|
|
||
|
|
self = getStageActor(initialReceive, name = selfActorName)
|
||
|
|
log.warning("Allocated receiver: {}", self.ref)
|
||
|
|
|
||
|
|
promise.success(new SinkRef(self.ref, settings.initialDemand))
|
||
|
|
}
|
||
|
|
|
||
|
|
override def onPull(): Unit = {
|
||
|
|
tryPush()
|
||
|
|
triggerCumulativeDemand()
|
||
|
|
}
|
||
|
|
|
||
|
|
def triggerCumulativeDemand(): Unit =
|
||
|
|
if (remotePartner ne null) {
|
||
|
|
val remainingRequested = java.lang.Long.min(highDemandWatermark, localCumulativeDemand - expectingSeqNr).toInt
|
||
|
|
val addDemand = requestStrategy.requestDemand(remainingRequested)
|
||
|
|
|
||
|
|
// only if demand has increased we shoot it right away
|
||
|
|
// otherwise it's the same demand level, so it'd be triggered via redelivery anyway
|
||
|
|
if (addDemand > 0) {
|
||
|
|
localCumulativeDemand += addDemand
|
||
|
|
val demand = StreamRefs.CumulativeDemand(localCumulativeDemand)
|
||
|
|
|
||
|
|
log.warning("[{}] Demanding until [{}] (+{})", selfActorName, localCumulativeDemand, addDemand)
|
||
|
|
remotePartner ! demand
|
||
|
|
scheduleDemandRedelivery()
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
val DemandRedeliveryTimerKey = "DemandRedeliveryTimerKey"
|
||
|
|
def scheduleDemandRedelivery() = scheduleOnce(DemandRedeliveryTimerKey, settings.demandRedeliveryInterval)
|
||
|
|
override protected def onTimer(timerKey: Any): Unit = timerKey match {
|
||
|
|
case DemandRedeliveryTimerKey ⇒
|
||
|
|
log.debug("[{}] Scheduled re-delivery of demand until [{}]", selfActorName, localCumulativeDemand)
|
||
|
|
remotePartner ! StreamRefs.CumulativeDemand(localCumulativeDemand)
|
||
|
|
scheduleDemandRedelivery()
|
||
|
|
}
|
||
|
|
|
||
|
|
lazy val initialReceive: ((ActorRef, Any)) ⇒ Unit = {
|
||
|
|
case (sender, msg @ StreamRefs.SequencedOnNext(seqNr, payload)) ⇒
|
||
|
|
observeAndValidateSender(sender, "Illegal sender in SequencedOnNext")
|
||
|
|
observeAndValidateSequenceNr(seqNr, "Illegal sequence nr in SequencedOnNext")
|
||
|
|
log.warning("Received seq {} from {}", msg, sender)
|
||
|
|
|
||
|
|
triggerCumulativeDemand()
|
||
|
|
tryPush(payload)
|
||
|
|
|
||
|
|
case (sender, StreamRefs.RemoteSinkCompleted(seqNr)) ⇒
|
||
|
|
observeAndValidateSender(sender, "Illegal sender in RemoteSinkCompleted")
|
||
|
|
observeAndValidateSequenceNr(seqNr, "Illegal sequence nr in RemoteSinkCompleted")
|
||
|
|
log.debug("The remote Sink has completed, completing this source as well...")
|
||
|
|
|
||
|
|
self.unwatch(sender)
|
||
|
|
completeStage()
|
||
|
|
|
||
|
|
case (sender, StreamRefs.RemoteSinkFailure(reason)) ⇒
|
||
|
|
observeAndValidateSender(sender, "Illegal sender in RemoteSinkFailure")
|
||
|
|
log.debug("The remote Sink has failed, failing (reason: {})", reason)
|
||
|
|
|
||
|
|
self.unwatch(sender)
|
||
|
|
failStage(StreamRefs.RemoteStreamRefActorTerminatedException(s"Remote Sink failed, reason: $reason"))
|
||
|
|
}
|
||
|
|
|
||
|
|
def tryPush(): Unit =
|
||
|
|
if (isAvailable(out) && receiveBuffer.nonEmpty) {
|
||
|
|
val elem = receiveBuffer.dequeue()
|
||
|
|
log.warning(s"PUSHING SIGNALED ${elem} (capacity: ${receiveBuffer.used}/${receiveBuffer.capacity})")
|
||
|
|
push(out, elem)
|
||
|
|
}
|
||
|
|
def tryPush(payload: Any): Unit =
|
||
|
|
if (isAvailable(out)) {
|
||
|
|
if (receiveBuffer.nonEmpty) {
|
||
|
|
val elem = receiveBuffer.dequeue()
|
||
|
|
push(out, elem)
|
||
|
|
receiveBuffer.enqueue(payload.asInstanceOf[T])
|
||
|
|
log.warning(s"PUSHING SIGNALED ${elem} BUFFERING payload" + payload + s"(capacity: ${receiveBuffer.used}/${receiveBuffer.capacity})")
|
||
|
|
} else {
|
||
|
|
push(out, payload.asInstanceOf[T])
|
||
|
|
log.warning(s"PUSHING DIRECTLY ${payload}")
|
||
|
|
}
|
||
|
|
} else {
|
||
|
|
receiveBuffer.enqueue(payload.asInstanceOf[T])
|
||
|
|
log.warning("PUSHING BUFFERING payload" + payload + s"(capacity: ${receiveBuffer.used}/${receiveBuffer.capacity})")
|
||
|
|
}
|
||
|
|
|
||
|
|
@throws[StreamRefs.InvalidPartnerActorException]
|
||
|
|
def observeAndValidateSender(sender: ActorRef, msg: String): Unit =
|
||
|
|
if (remotePartner == null) {
|
||
|
|
log.debug("Received first message from {}, assuming it to be the remote partner for this stage", sender)
|
||
|
|
remotePartner = sender
|
||
|
|
self.watch(sender)
|
||
|
|
} else if (sender != remotePartner) {
|
||
|
|
throw StreamRefs.InvalidPartnerActorException(sender, remotePartner, msg)
|
||
|
|
}
|
||
|
|
|
||
|
|
@throws[StreamRefs.InvalidSequenceNumberException]
|
||
|
|
def observeAndValidateSequenceNr(seqNr: Long, msg: String): Unit =
|
||
|
|
if (isInvalidSequenceNr(seqNr)) {
|
||
|
|
throw StreamRefs.InvalidSequenceNumberException(expectingSeqNr, seqNr, msg)
|
||
|
|
} else {
|
||
|
|
expectingSeqNr += 1
|
||
|
|
}
|
||
|
|
def isInvalidSequenceNr(seqNr: Long): Boolean =
|
||
|
|
seqNr != expectingSeqNr
|
||
|
|
|
||
|
|
setHandler(out, this)
|
||
|
|
}
|
||
|
|
(logic, promise.future) // FIXME we'd want to expose just the ref!
|
||
|
|
}
|
||
|
|
|
||
|
|
override def toString: String =
|
||
|
|
s"${Logging.simpleName(getClass)}()}"
|
||
|
|
}
|
||
|
|
|
||
|
|
/**
|
||
|
|
* The "handed out" side of a SinkRef. It powers a Source on the other side.
|
||
|
|
* TODO naming!??!?!!?!?!?!
|
||
|
|
*
|
||
|
|
* Do not create this instance directly, but use `SinkRef` factories, to run/setup its targetRef
|
||
|
|
*/
|
||
|
|
final class SinkRef[In] private[akka] ( // TODO is it more of a SourceRefSink?
|
||
|
|
private[akka] val targetRef: ActorRef,
|
||
|
|
private[akka] val initialDemand: Long
|
||
|
|
) extends GraphStage[SinkShape[In]] with Serializable { stage ⇒
|
||
|
|
import akka.stream.remote.StreamRefs._
|
||
|
|
|
||
|
|
val in = Inlet[In](s"${Logging.simpleName(getClass)}($targetRef).in")
|
||
|
|
override def shape: SinkShape[In] = SinkShape.of(in)
|
||
|
|
|
||
|
|
override def createLogic(inheritedAttributes: Attributes) = new GraphStageLogic(shape) with StageLogging with InHandler {
|
||
|
|
|
||
|
|
private[this] lazy val streamRefsMaster = StreamRefsMaster(ActorMaterializerHelper.downcast(materializer).system)
|
||
|
|
private[this] lazy val selfActorName = streamRefsMaster.nextSinkRefName()
|
||
|
|
|
||
|
|
// we assume that there is at least SOME buffer space
|
||
|
|
private[this] var remoteCumulativeDemandReceived = initialDemand
|
||
|
|
|
||
|
|
// FIXME this one will be sent over remoting so we have to be able to make that work
|
||
|
|
private[this] var remoteCumulativeDemandConsumed = 0L
|
||
|
|
private[this] var self: GraphStageLogic.StageActor = _
|
||
|
|
implicit def selfSender: ActorRef = self.ref
|
||
|
|
|
||
|
|
override def preStart(): Unit = {
|
||
|
|
self = getStageActor(initialReceive, selfActorName)
|
||
|
|
self.watch(targetRef)
|
||
|
|
|
||
|
|
log.warning("Created SinkRef, pointing to remote Sink receiver: {}, local worker: {}", targetRef, self)
|
||
|
|
|
||
|
|
pull(in)
|
||
|
|
}
|
||
|
|
|
||
|
|
lazy val initialReceive: ((ActorRef, Any)) ⇒ Unit = {
|
||
|
|
case (_, Terminated(`targetRef`)) ⇒
|
||
|
|
failStage(failRemoteTerminated())
|
||
|
|
|
||
|
|
case (sender, CumulativeDemand(d)) ⇒
|
||
|
|
validatePartnerRef(sender)
|
||
|
|
|
||
|
|
if (remoteCumulativeDemandReceived < d) {
|
||
|
|
remoteCumulativeDemandReceived = d
|
||
|
|
log.warning("Received cumulative demand [{}], consumable demand: [{}]", CumulativeDemand(d), remoteCumulativeDemandReceived - remoteCumulativeDemandConsumed)
|
||
|
|
}
|
||
|
|
tryPull()
|
||
|
|
}
|
||
|
|
|
||
|
|
override def onPush(): Unit = {
|
||
|
|
val elem = grabSequenced(in)
|
||
|
|
targetRef ! elem
|
||
|
|
log.warning("Sending sequenced: {} to {}", elem, targetRef)
|
||
|
|
tryPull()
|
||
|
|
}
|
||
|
|
|
||
|
|
private def tryPull() =
|
||
|
|
if (remoteCumulativeDemandConsumed < remoteCumulativeDemandReceived && !hasBeenPulled(in))
|
||
|
|
pull(in)
|
||
|
|
|
||
|
|
private def grabSequenced[T](in: Inlet[T]): SequencedOnNext[T] = {
|
||
|
|
val onNext = SequencedOnNext(remoteCumulativeDemandConsumed, grab(in))
|
||
|
|
remoteCumulativeDemandConsumed += 1
|
||
|
|
onNext
|
||
|
|
}
|
||
|
|
|
||
|
|
override def onUpstreamFailure(ex: Throwable): Unit = {
|
||
|
|
targetRef ! StreamRefs.RemoteSinkFailure(ex.getMessage) // TODO yes / no? At least the message I guess
|
||
|
|
self.unwatch(targetRef)
|
||
|
|
super.onUpstreamFailure(ex)
|
||
|
|
}
|
||
|
|
|
||
|
|
override def onUpstreamFinish(): Unit = {
|
||
|
|
targetRef ! StreamRefs.RemoteSinkCompleted(remoteCumulativeDemandConsumed)
|
||
|
|
self.unwatch(targetRef)
|
||
|
|
super.onUpstreamFinish()
|
||
|
|
}
|
||
|
|
|
||
|
|
setHandler(in, this)
|
||
|
|
}
|
||
|
|
|
||
|
|
private def validatePartnerRef(ref: ActorRef) = {
|
||
|
|
if (ref != targetRef) throw new RuntimeException("Got demand from weird actor! Not the one I trust hmmmm!!!")
|
||
|
|
}
|
||
|
|
|
||
|
|
private def failRemoteTerminated() = {
|
||
|
|
RemoteStreamRefActorTerminatedException(s"Remote target receiver of data ${targetRef} terminated. Local stream terminating, message loss (on remote side) may have happened.")
|
||
|
|
}
|
||
|
|
|
||
|
|
override def toString = s"${Logging.simpleName(getClass)}($targetRef)"
|
||
|
|
|
||
|
|
}
|