jove.notebook.components.KernelWS.scala Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of jove-notebook_2.11 Show documentation
jove-notebook
The newest version!
package jove
package notebook.components

import java.util.concurrent.atomic.AtomicInteger
import java.util.concurrent.{ThreadFactory, Executors, LinkedBlockingQueue}

import jove.MessageSocket.Channel
import scalaz._
import scalaz.concurrent.Task
import scalaz.stream.Process

object ScalazStreamSocketProcess extends LazyLogging {
  def apply(socket: MessageSocket, id: String = ""): Channel => Process[Task, RawMessage] = {
    // Using our own thread pool here, to not preempt threads from the main one, which has a fixed size
    implicit val executorService =
      Executors.newFixedThreadPool(Channel.channels.size, new ThreadFactory {
        val defaultThreadFactory = Executors.defaultThreadFactory()
        val c = new AtomicInteger()
        def newThread(r: Runnable) = {
          val t = defaultThreadFactory.newThread(r)
          t setDaemon true
          t setName s"SocketStreamQueueProxy-${c addAndGet 1} $id"
          t
        }
      })

    def readFrom(channel: Channel) =
      Task {
        \/.fromTryCatchNonFatal {
          def helper(): RawMessage = {
            logger debug s"Polling on $channel ($id)"

            if (socket.poll(channel, 1000L)) {
              logger debug s"Receiving on $channel ($id)"
              val r = socket.receive(channel)
              logger debug s"Received $r on $channel ($id)"
              r match {
                case -\/(err) =>
                  logger debug s"Error while receiving on $channel ($id): $err"
                  helper()
                case \/-(m) =>
                  m
              }
            } else
              helper()
          }

          helper()
        }
      }

    def proc(channel: Channel) =
      Process.repeatEval(readFrom(channel)).collect {case \/-(m) => m}

    val reqProc = proc(Channel.Requests)
    val pubProc = proc(Channel.Publish)
    val ctlProc = proc(Channel.Control)
    val stdinProc = proc(Channel.Stdin)
    val hbProc = proc(Channel.Heartbeat)

    {
      case Channel.Requests =>
        reqProc
      case Channel.Publish =>
        pubProc
      case Channel.Control =>
        ctlProc
      case Channel.Stdin =>
        stdinProc
      case Channel.Heartbeat =>
        hbProc
    }
  }
}

object ScalazStreamMessageSendSocket extends LazyLogging {
  def apply(id: String = ""): (MessageSendSocket, Channel => Process[Task, RawMessage]) = {

    // Using our own thread pool here, to not preempt threads from the main one, which has a fixed size
    implicit val executorService =
      Executors.newFixedThreadPool(Channel.channels.size, new ThreadFactory {
        val defaultThreadFactory = Executors.defaultThreadFactory()
        val c = new AtomicInteger()
        def newThread(r: Runnable) = {
          val t = defaultThreadFactory.newThread(r)
          t setDaemon true
          t setName s"SocketQueueProxy-${c addAndGet 1} $id"
          t
        }
      })

    val reqQueue, pubQueue, ctlQueue, stdinQueue, hbQueue = new LinkedBlockingQueue[RawMessage]

    val socket =
      new MessageSendSocket {
        def send(channel: Channel, msg: RawMessage, silently: Boolean) = {
          logger debug s"Queuing $msg on $channel"
          channel match {
            case Channel.Requests =>
              reqQueue add msg
            case Channel.Publish =>
              pubQueue add msg
            case Channel.Control =>
              ctlQueue add msg
            case Channel.Stdin =>
              stdinQueue add msg
            case Channel.Heartbeat =>
              hbQueue add msg
          }
        }
      }

    def readFrom(q: LinkedBlockingQueue[RawMessage], channel: Channel) =
      Task {
        \/.fromTryCatchNonFatal{
          logger debug s"Taking on $channel"
          val r = q.take()
          logger debug s"Took $r on $channel"
          r
        }
      }

    def proc(queue: LinkedBlockingQueue[RawMessage], channel: Channel) =
      Process.repeatEval(readFrom(queue, channel)).collect {case \/-(m) => m}

    val reqProc = proc(reqQueue, Channel.Requests)
    val pubProc = proc(pubQueue, Channel.Publish)
    val ctlProc = proc(ctlQueue, Channel.Control)
    val stdinProc = proc(stdinQueue, Channel.Stdin)
    val hbProc = proc(hbQueue, Channel.Heartbeat)

    (socket, {
      case Channel.Requests =>
        reqProc
      case Channel.Publish =>
        pubProc
      case Channel.Control =>
        ctlProc
      case Channel.Stdin =>
        stdinProc
      case Channel.Heartbeat =>
        hbProc
    })
  }
}

class KernelWS extends LazyLogging {
  private val _remoteSocketsLock = new AnyRef
  private var _remoteSockets = Map.empty[String, Channel => Process[Task, RawMessage]]

  def remoteProcess(socket: MessageSocket, id: String, channel: Channel): Process[Task, RawMessage] = _remoteSocketsLock.synchronized {
    val f =
      _remoteSockets.getOrElse(id, {
        val f = ScalazStreamSocketProcess(socket, id)
        _remoteSockets += id -> f
        f
      })

    f(channel)
  }


  private val _webSocketsLock = new AnyRef
  private var _webSockets = Map.empty[String, (MessageSendSocket, Channel => Process[Task, RawMessage])]

  def process(id: String, channel: Channel): Process[Task, RawMessage] = _webSocketsLock.synchronized {
    val f =
      _webSockets.get(id) match {
        case Some((_, f)) =>
          f
        case None =>
          val r = ScalazStreamMessageSendSocket(id)
          _webSockets += id -> r
          r._2
      }

    f(channel)
  }

  def socket(id: String): MessageSendSocket = _webSocketsLock.synchronized {
    _webSockets.get(id) match {
      case Some((s, _)) =>
        s
      case None =>
        val r = ScalazStreamMessageSendSocket(id)
        _webSockets += id -> r
        r._1
    }
  }
}