基于Akka Streams的WebSocket：如何修改代码实现任意内部事件消息广播？

问题描述 #

我目前有一个基于Akka HTTP的WebSocket服务，它可以向所有连接的客户端广播定时生成的整数ticks，但我需要修改它，让它能够广播任意内部系统事件产生的消息——也就是需要有一个专用的Source，我可以随时向它发送消息，然后自动转发给所有已连接的客户端。现有代码如下：

import akka.http.scaladsl.server.RouteResult.route2HandlerFlow
import akka.http.scaladsl.server.Directives._
import akka.stream.scaladsl.{BroadcastHub, Flow, Keep, Sink, Source}
import akka.stream.{ActorMaterializer, ActorSystem, ThrottleMode}
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import scala.concurrent.duration._

implicit val system = ActorSystem("Server")
implicit val mat = ActorMaterializer()

// The source to broadcast (just ints for simplicity)
val dataSource = Source(1 to 1000).throttle(1, 1.second, 1, ThrottleMode.Shaping).map(_.toString)

// Go via BroadcastHub to allow multiple clients to connect
val runnableGraph: RunnableGraph[Source[String, NotUsed]] = dataSource.toMat(BroadcastHub.sink(bufferSize = 256))(Keep.right)
val producer: Source[String, NotUsed] = runnableGraph.run()

// Optional - add sink to avoid backpressuring the original flow when no clients are attached
producer.runWith(Sink.ignore)

val wsHandler: Flow[Message, Message, NotUsed] = Flow[Message]
  .mapConcat(_ => Nil) // Ignore any data sent from the client
  .merge(producer) // Stream the data we want to the client
  .map(l => TextMessage(l.toString))

val route = path("ws") {
  handleWebSocketMessages(wsHandler)
}

val port = 8080
println("Starting up route")
Http().bindAndHandle(route2HandlerFlow(route), "127.0.0.1", port)
println(s"Started HTTP server on port $port")

解决方案 #

要实现这个需求，核心是使用Akka Streams的Source.queue——它允许你从流的外部（比如你的内部事件处理逻辑）主动推送消息到流中，再结合BroadcastHub实现多客户端的广播。具体修改步骤如下：

1. 创建一个可外部推送的Source队列：用Source.queue替代原来的定时ticks数据源，这个队列可以接收你发送的任意消息。
2. 将队列连接到BroadcastHub：确保所有连接的客户端都能订阅到这个队列的消息流。
3. 提供消息发送入口：封装一个方法，用来向队列发送内部事件消息。

修改后的完整代码：

import akka.http.scaladsl.server.RouteResult.route2HandlerFlow
import akka.http.scaladsl.server.Directives._
import akka.stream.scaladsl.{BroadcastHub, Flow, Keep, Sink, Source}
import akka.stream.{ActorMaterializer, ActorSystem, OverflowStrategy}
import akka.http.scaladsl.model.ws.{Message, TextMessage}
import scala.concurrent.Future

implicit val system = ActorSystem("Server")
implicit val mat = ActorMaterializer()
implicit val ec = system.dispatcher

// 1. 创建一个Source队列，用于接收内部事件消息
// OverflowStrategy.dropNew表示当队列满时丢弃新消息，你可以根据需求选择其他策略（比如backpressure）
val messageQueue = Source.queue[String](bufferSize = 1000, OverflowStrategy.dropNew)

// 2. 将队列连接到BroadcastHub，生成可广播的消息源
val runnableGraph = messageQueue.toMat(BroadcastHub.sink(bufferSize = 256))(Keep.both)
val (queue, producer) = runnableGraph.run()

// 可选：当没有客户端连接时，避免队列背压
producer.runWith(Sink.ignore)

// 3. 封装发送消息的方法，供内部事件调用
def broadcastInternalMessage(message: String): Future[Unit] = {
  queue.offer(message).map {
    case akka.stream.QueueOfferResult.Enqueued => ()
    case akka.stream.QueueOfferResult.Dropped =>
      system.log.warning(s"Failed to broadcast message: queue is full, message dropped: $message")
    case akka.stream.QueueOfferResult.Failure(ex) =>
      system.log.error(ex, s"Failed to broadcast message due to error")
    case akka.stream.QueueOfferResult.QueueClosed =>
      system.log.error("Cannot broadcast message: queue is closed")
  }
}

// WebSocket处理逻辑保持不变，只是现在producer来自消息队列
val wsHandler: Flow[Message, Message, NotUsed] = Flow[Message]
  .mapConcat(_ => Nil) // 忽略客户端发送的消息
  .merge(producer) // 将广播消息流合并到WebSocket输出
  .map(msg => TextMessage(msg))

val route = path("ws") {
  handleWebSocketMessages(wsHandler)
}

val port = 8080
println("Starting up WebSocket broadcast server")
Http().bindAndHandle(route2HandlerFlow(route), "127.0.0.1", port)
println(s"Server started on port $port")

// 示例：模拟内部事件发送消息
broadcastInternalMessage("System started successfully!")
// 你可以在任意内部逻辑中调用broadcastInternalMessage发送消息

关键细节解释 #

• Source.queue的缓冲策略：OverflowStrategy有多种选项，比如dropNew（丢弃新消息）、backpressure（向发送方施加背压）、dropHead（丢弃最早的消息），你需要根据系统的可靠性需求选择合适的策略。
• BroadcastHub的作用：它会为每个新连接的客户端创建一个独立的订阅流，确保所有客户端都能收到相同的消息，同时自动处理客户端的连接和断开。
• 消息发送的异步处理：queue.offer返回一个Future，你可以根据返回的QueueOfferResult处理发送失败的情况，比如记录日志或者重试。

内容的提问来源于stack exchange，提问作者src091