akka.io.WithUdpSend.scala Maven / Gradle / Ivy

/*
 * Copyright (C) 2009-2020 Lightbend Inc. 
 */

package akka.io

import java.net.InetSocketAddress
import java.nio.channels.{ DatagramChannel, SelectionKey }

import scala.util.control.NonFatal

import akka.actor.{ Actor, ActorLogging, ActorRef }
import akka.io.SelectionHandler._
import akka.io.Udp.{ CommandFailed, Send }
import akka.io.dns.DnsProtocol

/**
 * INTERNAL API
 */
private[io] trait WithUdpSend {
  me: Actor with ActorLogging =>

  private var pendingSend: Send = null
  private var pendingCommander: ActorRef = null
  // If send fails first, we allow a second go after selected writable, but no more. This flag signals that
  // pending send was already tried once.
  private var retriedSend = false
  private def hasWritePending = pendingSend ne null

  def channel: DatagramChannel
  def udp: UdpExt
  val settings = udp.settings

  import settings._

  def sendHandlers(registration: ChannelRegistration): Receive = {
    case send: Send if hasWritePending =>
      if (TraceLogging) log.debug("Dropping write because queue is full")
      sender() ! CommandFailed(send)

    case send: Send if send.payload.isEmpty =>
      if (send.wantsAck)
        sender() ! send.ack

    case send: Send =>
      pendingSend = send
      pendingCommander = sender()
      if (send.target.isUnresolved) {
        Dns.resolve(DnsProtocol.Resolve(send.target.getHostName), context.system, self) match {
          case Some(r) =>
            try {
              pendingSend = pendingSend.copy(target = new InetSocketAddress(r.address(), pendingSend.target.getPort))
              doSend(registration)
            } catch {
              case NonFatal(e) =>
                sender() ! CommandFailed(send)
                log.debug("Failure while sending UDP datagram to remote address [{}]: {}", send.target, e)
                retriedSend = false
                pendingSend = null
                pendingCommander = null
            }
          case None =>
            sender() ! CommandFailed(send)
            log.debug("Name resolution failed for remote address [{}]", send.target)
            retriedSend = false
            pendingSend = null
            pendingCommander = null
        }
      } else {
        doSend(registration)
      }

    case ChannelWritable => if (hasWritePending) doSend(registration)
  }

  private def doSend(registration: ChannelRegistration): Unit = {
    val buffer = udp.bufferPool.acquire()
    try {
      buffer.clear()
      pendingSend.payload.copyToBuffer(buffer)
      buffer.flip()
      val writtenBytes = channel.send(buffer, pendingSend.target)
      if (TraceLogging) log.debug("Wrote [{}] bytes to channel", writtenBytes)

      // Datagram channel either sends the whole message, or nothing
      if (writtenBytes == 0) {
        if (retriedSend) {
          pendingCommander ! CommandFailed(pendingSend)
          retriedSend = false
          pendingSend = null
          pendingCommander = null
        } else {
          registration.enableInterest(SelectionKey.OP_WRITE)
          retriedSend = true
        }
      } else {
        if (pendingSend.wantsAck) pendingCommander ! pendingSend.ack
        retriedSend = false
        pendingSend = null
        pendingCommander = null
      }
    } finally {
      udp.bufferPool.release(buffer)
    }
  }
}