318 lines
9.8 KiB
Scala
318 lines
9.8 KiB
Scala
/**
|
|
* Copyright (C) 2009-2014 Typesafe Inc. <http://www.typesafe.com>
|
|
*/
|
|
package akka.stream.impl.io
|
|
|
|
import java.net.InetSocketAddress
|
|
import akka.io.{ IO, Tcp }
|
|
import scala.concurrent.Promise
|
|
import scala.util.control.NoStackTrace
|
|
import akka.actor._
|
|
import akka.util.ByteString
|
|
import akka.io.Tcp._
|
|
import akka.stream.ActorFlowMaterializerSettings
|
|
import akka.stream.StreamTcpException
|
|
import org.reactivestreams.Processor
|
|
import akka.stream.impl._
|
|
import akka.actor.ActorLogging
|
|
|
|
/**
|
|
* INTERNAL API
|
|
*/
|
|
private[akka] object TcpStreamActor {
|
|
case object WriteAck extends Tcp.Event
|
|
|
|
def outboundProps(processorPromise: Promise[Processor[ByteString, ByteString]],
|
|
localAddressPromise: Promise[InetSocketAddress],
|
|
halfClose: Boolean,
|
|
connectCmd: Connect,
|
|
materializerSettings: ActorFlowMaterializerSettings): Props =
|
|
Props(new OutboundTcpStreamActor(processorPromise, localAddressPromise, halfClose, connectCmd,
|
|
materializerSettings)).withDispatcher(materializerSettings.dispatcher).withDeploy(Deploy.local)
|
|
|
|
def inboundProps(connection: ActorRef, halfClose: Boolean, settings: ActorFlowMaterializerSettings): Props =
|
|
Props(new InboundTcpStreamActor(connection, halfClose, settings)).withDispatcher(settings.dispatcher).withDeploy(Deploy.local)
|
|
|
|
}
|
|
|
|
/**
|
|
* INTERNAL API
|
|
*/
|
|
private[akka] abstract class TcpStreamActor(val settings: ActorFlowMaterializerSettings, halfClose: Boolean) extends Actor
|
|
with ActorLogging {
|
|
|
|
import TcpStreamActor._
|
|
|
|
val primaryInputs: Inputs = new BatchingInputBuffer(settings.initialInputBufferSize, writePump) {
|
|
override def inputOnError(e: Throwable): Unit = fail(e)
|
|
}
|
|
|
|
val primaryOutputs: Outputs = new SimpleOutputs(self, readPump)
|
|
|
|
def fullClose: Boolean = !halfClose
|
|
|
|
object tcpInputs extends DefaultInputTransferStates {
|
|
private var closed: Boolean = false
|
|
private var pendingElement: ByteString = null
|
|
private var connection: ActorRef = _
|
|
|
|
val subreceive = new SubReceive(Actor.emptyBehavior)
|
|
|
|
def setConnection(c: ActorRef): Unit = {
|
|
connection = c
|
|
// Prefetch
|
|
c ! ResumeReading
|
|
subreceive.become(handleRead)
|
|
readPump.pump()
|
|
}
|
|
|
|
def handleRead: Receive = {
|
|
case Received(data) ⇒
|
|
if (closed) connection ! ResumeReading
|
|
else {
|
|
pendingElement = data
|
|
readPump.pump()
|
|
}
|
|
case ConfirmedClosed ⇒
|
|
cancelWithoutTcpClose()
|
|
readPump.pump()
|
|
case PeerClosed ⇒
|
|
cancelWithoutTcpClose()
|
|
readPump.pump()
|
|
}
|
|
|
|
override def inputsAvailable: Boolean = pendingElement ne null
|
|
override def inputsDepleted: Boolean = closed && !inputsAvailable
|
|
override def isClosed: Boolean = closed
|
|
|
|
override def cancel(): Unit = {
|
|
if (!closed) {
|
|
closed = true
|
|
pendingElement = null
|
|
if (connection ne null) {
|
|
if (tcpOutputs.isClosed)
|
|
connection ! Abort
|
|
else
|
|
connection ! ResumeReading
|
|
}
|
|
}
|
|
}
|
|
|
|
def cancelWithoutTcpClose(): Unit = {
|
|
if (!closed) {
|
|
closed = true
|
|
pendingElement = null
|
|
}
|
|
}
|
|
|
|
override def dequeueInputElement(): Any = {
|
|
val elem = pendingElement
|
|
pendingElement = null
|
|
connection ! ResumeReading
|
|
elem
|
|
}
|
|
|
|
}
|
|
|
|
object tcpOutputs extends DefaultOutputTransferStates {
|
|
private var closed: Boolean = false
|
|
private var lastWriteAcked = true
|
|
private var connection: ActorRef = _
|
|
|
|
def isClosed: Boolean = closed
|
|
// Full-close mode needs to wait for the last write Ack before sending Close to avoid doing a connection reset
|
|
def isFlushed: Boolean = closed && (halfClose || lastWriteAcked)
|
|
|
|
private def initialized: Boolean = connection ne null
|
|
|
|
def setConnection(c: ActorRef): Unit = {
|
|
connection = c
|
|
writePump.pump()
|
|
subreceive.become(handleWrite)
|
|
}
|
|
|
|
val subreceive = new SubReceive(Actor.emptyBehavior)
|
|
|
|
def handleWrite: Receive = {
|
|
case WriteAck ⇒
|
|
lastWriteAcked = true
|
|
if (fullClose && closed) {
|
|
// Finish the closing after the last write has been flushed in full close mode.
|
|
connection ! Close
|
|
tryShutdown()
|
|
}
|
|
writePump.pump()
|
|
|
|
}
|
|
|
|
override def error(e: Throwable): Unit = {
|
|
if (!closed && initialized) connection ! Abort
|
|
closed = true
|
|
}
|
|
|
|
override def complete(): Unit = {
|
|
if (!closed && initialized) {
|
|
closed = true
|
|
if (tcpInputs.isClosed && (halfClose || lastWriteAcked)) {
|
|
// We can immediately close if
|
|
// - half close mode, and read size already finished
|
|
// - full close mode, and last write has been acked
|
|
//
|
|
// if in full close mode, and has a non-acknowledged write, we will do the closing in handleWrite
|
|
// when the Ack arrives
|
|
connection ! Close
|
|
tryShutdown()
|
|
} else
|
|
connection ! ConfirmedClose
|
|
}
|
|
}
|
|
|
|
override def cancel(): Unit = complete()
|
|
|
|
override def enqueueOutputElement(elem: Any): Unit = {
|
|
ReactiveStreamsCompliance.requireNonNullElement(elem)
|
|
connection ! Write(elem.asInstanceOf[ByteString], WriteAck)
|
|
lastWriteAcked = false
|
|
}
|
|
|
|
override def demandAvailable: Boolean = lastWriteAcked
|
|
}
|
|
|
|
object writePump extends Pump {
|
|
|
|
def running = TransferPhase(primaryInputs.NeedsInput && tcpOutputs.NeedsDemand) { () ⇒
|
|
var batch = ByteString.empty
|
|
while (primaryInputs.inputsAvailable) batch ++= primaryInputs.dequeueInputElement().asInstanceOf[ByteString]
|
|
tcpOutputs.enqueueOutputElement(batch)
|
|
}
|
|
|
|
override protected def pumpFinished(): Unit = {
|
|
if (fullClose) {
|
|
// In full close mode we shut down the read size immediately once the write side is finished
|
|
tcpInputs.cancel()
|
|
primaryOutputs.complete()
|
|
readPump.pump()
|
|
}
|
|
tcpOutputs.complete()
|
|
primaryInputs.cancel()
|
|
tryShutdown()
|
|
}
|
|
override protected def pumpFailed(e: Throwable): Unit = fail(e)
|
|
}
|
|
|
|
object readPump extends Pump {
|
|
|
|
def running = TransferPhase(tcpInputs.NeedsInput && primaryOutputs.NeedsDemand) { () ⇒
|
|
primaryOutputs.enqueueOutputElement(tcpInputs.dequeueInputElement())
|
|
}
|
|
|
|
override protected def pumpFinished(): Unit = {
|
|
tcpInputs.cancel()
|
|
primaryOutputs.complete()
|
|
tryShutdown()
|
|
}
|
|
override protected def pumpFailed(e: Throwable): Unit = fail(e)
|
|
}
|
|
|
|
final override def receive = new ExposedPublisherReceive(activeReceive, unhandled) {
|
|
override def receiveExposedPublisher(ep: ExposedPublisher): Unit = {
|
|
primaryOutputs.subreceive(ep)
|
|
context become activeReceive
|
|
}
|
|
}
|
|
|
|
def activeReceive =
|
|
primaryInputs.subreceive orElse
|
|
primaryOutputs.subreceive orElse
|
|
tcpInputs.subreceive orElse
|
|
tcpOutputs.subreceive orElse
|
|
commonCloseHandling
|
|
|
|
def commonCloseHandling: Receive = {
|
|
case Closed ⇒
|
|
tcpInputs.cancel()
|
|
tcpOutputs.complete()
|
|
writePump.pump()
|
|
readPump.pump()
|
|
case ErrorClosed(cause) ⇒ fail(new StreamTcpException(s"The connection closed with error $cause"))
|
|
case CommandFailed(cmd) ⇒ fail(new StreamTcpException(s"Tcp command [$cmd] failed"))
|
|
case Aborted ⇒ fail(new StreamTcpException("The connection has been aborted"))
|
|
}
|
|
|
|
readPump.nextPhase(readPump.running)
|
|
writePump.nextPhase(writePump.running)
|
|
|
|
def fail(e: Throwable): Unit = {
|
|
if (settings.debugLogging)
|
|
log.debug("fail due to: {}", e.getMessage)
|
|
tcpInputs.cancel()
|
|
tcpOutputs.error(e)
|
|
primaryInputs.cancel()
|
|
primaryOutputs.error(e)
|
|
tryShutdown()
|
|
}
|
|
|
|
def tryShutdown(): Unit =
|
|
if (primaryInputs.isClosed && tcpInputs.isClosed && tcpOutputs.isFlushed)
|
|
context.stop(self)
|
|
|
|
override def postStop(): Unit = {
|
|
// Close if it has not yet been done
|
|
tcpInputs.cancel()
|
|
tcpOutputs.complete()
|
|
primaryInputs.cancel()
|
|
primaryOutputs.complete()
|
|
super.postStop() // Remember, we have a Stash
|
|
}
|
|
}
|
|
|
|
/**
|
|
* INTERNAL API
|
|
*/
|
|
private[akka] class InboundTcpStreamActor(
|
|
val connection: ActorRef, _halfClose: Boolean, _settings: ActorFlowMaterializerSettings)
|
|
extends TcpStreamActor(_settings, _halfClose) {
|
|
|
|
connection ! Register(self, keepOpenOnPeerClosed = true, useResumeWriting = false)
|
|
tcpInputs.setConnection(connection)
|
|
tcpOutputs.setConnection(connection)
|
|
}
|
|
|
|
/**
|
|
* INTERNAL API
|
|
*/
|
|
private[akka] class OutboundTcpStreamActor(processorPromise: Promise[Processor[ByteString, ByteString]],
|
|
localAddressPromise: Promise[InetSocketAddress],
|
|
_halfClose: Boolean,
|
|
val connectCmd: Connect, _settings: ActorFlowMaterializerSettings)
|
|
extends TcpStreamActor(_settings, _halfClose) {
|
|
import TcpStreamActor._
|
|
import context.system
|
|
|
|
val initSteps = new SubReceive(waitingExposedProcessor)
|
|
|
|
override def activeReceive = initSteps orElse super.activeReceive
|
|
|
|
def waitingExposedProcessor: Receive = {
|
|
case StreamTcpManager.ExposedProcessor(processor) ⇒
|
|
IO(Tcp) ! connectCmd
|
|
initSteps.become(waitConnection(processor))
|
|
}
|
|
|
|
def waitConnection(exposedProcessor: Processor[ByteString, ByteString]): Receive = {
|
|
case Connected(remoteAddress, localAddress) ⇒
|
|
val connection = sender()
|
|
connection ! Register(self, keepOpenOnPeerClosed = true, useResumeWriting = false)
|
|
tcpOutputs.setConnection(connection)
|
|
tcpInputs.setConnection(connection)
|
|
localAddressPromise.success(localAddress)
|
|
processorPromise.success(exposedProcessor)
|
|
initSteps.become(Actor.emptyBehavior)
|
|
|
|
case f: CommandFailed ⇒
|
|
val ex = new StreamTcpException("Connection failed.")
|
|
localAddressPromise.failure(ex)
|
|
processorPromise.failure(ex)
|
|
fail(ex)
|
|
}
|
|
}
|