Redis pub/sub と Akka Stream

import java.time.ZonedDateTime
import org.slf4j.LoggerFactory

import scala.concurrent.duration._
import scala.concurrent.ExecutionContext.Implicits._

import akka.actor.ActorSystem
import akka.stream._
import akka.stream.scaladsl._

object RedisPubSubStream {

  val logger = LoggerFactory.getLogger(this.getClass())

  def main(args: Array[String]): Unit = {
    implicit val system = ActorSystem("pubsub")
    implicit val materializer = ActorMaterializer()

    val channelName = "channel1"

    val publisherSink = new PublisherSink(channelName)
    Source
      .tick(1.seconds, 5.seconds, Message(s"Hello, "))
      .map(m => {
        m.copy(data = m.data + s"this is ${ZonedDateTime.now()}")
      })
      .runWith(Sink.fromGraph(publisherSink))

    val subscriberSource = new SubscriberSource(channelName)
    Source
      .fromGraph(subscriberSource)
      .runWith(Sink.foreach(x => logger.debug(x.toString)))
      .onComplete(_ => system.terminate())
  }

}

import akka.stream._
import akka.stream.stage.{GraphStage, GraphStageLogic, InHandler}
import org.slf4j.LoggerFactory

class PublisherSink(
    val channelName: String
) extends GraphStage[SinkShape[Message]] {

  val logger = LoggerFactory.getLogger(this.getClass)
  val in: Inlet[Message] = Inlet("SubscribeSink.In")
  val channel = Channel(channelName)

  override def shape: SinkShape[Message] = SinkShape(in)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {
      override def preStart(): Unit = {
        pull(in)
      }

      setHandler(
        in,
        new InHandler {
          override def onPush(): Unit = {
            val message = grab(in)
            logger.debug(s"will publish message: $message")
            RedisClient.publish(message, channel)
            pull(in)
          }
        }
      )
    }

}

import java.util.concurrent.CyclicBarrier

import akka.stream._
import akka.stream.scaladsl._
import akka.stream.stage.{GraphStage, GraphStageLogic, OutHandler}
import org.slf4j.LoggerFactory

class SubscriberSource(
    val channelName: String
) extends GraphStage[SourceShape[Message]] {

  val logger = LoggerFactory.getLogger(this.getClass())
  val out: Outlet[Message] = Outlet("SubscriberSource.Out")
  val channel = Channel(channelName)

  override def shape: SourceShape[Message] = new SourceShape(out)

  override def createLogic(inheritedAttributes: Attributes): GraphStageLogic =
    new GraphStageLogic(shape) {

      private val barrier = new CyclicBarrier(2)
      private var message: Option[Message] = None
      private var threadName: Option[String] = None

      private val messageReceiver = new MessageReceiver {
        override def handleMessage(m: Message): Unit = {
          message = Some(m.copy())
          barrier.await()
        }
      }

      override def preStart(): Unit = {
        super.preStart()
        threadName = Some(RedisClient.subscribe(channel, messageReceiver))
      }

      override def postStop(): Unit = {
        super.postStop()
        RedisClient.unsubscribe(threadName.get)
      }

      setHandler(
        out,
        new OutHandler() {
          override def onPull(): Unit = {
            barrier.await()
            logger.debug(s"subscribed message: ${message.get}")
            push(out, message.get)
            message = None
            barrier.reset()
          }
        }
      )
    }

}

import org.slf4j.LoggerFactory
import redis.clients.jedis.{Jedis, JedisPool, JedisPoolConfig, JedisPubSub}

object RedisClient {

  private val logger = LoggerFactory.getLogger(this.getClass())
  private val host = "localhost"
  private val pool = new JedisPool(new JedisPoolConfig(), host)

  private val subscribers =
    scala.collection.mutable.Map[String, (Thread, Jedis, JedisPubSub)]()

  def publish(message: Message, channel: Channel): Unit = {
    withClient { jedis =>
      jedis.publish(channel.name, message.data)
    }
  }

  def subscribe(channel: Channel, messageReceiver: MessageReceiver): String = {
    val jedis = new Jedis(host)
    val pubsub = createPubSub(messageReceiver)
    val thread = new Thread(() => {
      logger.debug(s"start subscribe to ${channel.name}")
      jedis.subscribe(pubsub, channel.name)
    })
    subscribers.put(thread.getName(), (thread, jedis, pubsub))
    thread.start()
    thread.getName()
  }

  def unsubscribe(name: String): Unit = {
    val (thread, jedis, pubsub) = subscribers(name)
    pubsub.unsubscribe()
    jedis.quit()
  }

  private def withClient[T](f: Jedis => T): T = {
    val jedis = new Jedis(host)
    try {
      f(jedis)
    } finally {
      jedis.quit()
    }
  }

  def createPubSub(messageReceiver: MessageReceiver): JedisPubSub = {
    new JedisPubSub() {
      override def onSubscribe(channel: String,
                               subscribedChannels: Int): Unit =
        logger.debug("onSubscribe")
      override def onUnsubscribe(channel: String,
                                 subscribedChannels: Int): Unit =
        logger.debug("onUnsubscribe")
      override def onPSubscribe(pattern: String,
                                subscribedChannels: Int): Unit =
        logger.debug("onPSubscribe")
      override def onPUnsubscribe(pattern: String,
                                  subscribedChannels: Int): Unit =
        logger.debug("onPUnsubscribe")
      override def onPMessage(pattern: String,
                              channel: String,
                              message: String): Unit =
        logger.debug("onPMessage")
      override def onMessage(channel: String, message: String): Unit =
        messageReceiver.handleMessage(Message(message))
    }
  }

}

trait MessageReceiver {
  def handleMessage(m: Message): Unit
}