如何在单个 Akka-http 流中读取多个 Web 套接字？

Question

I am currently practicing Akka-http by attempting to establish multiple websockets connections. My code for creating the websockets client flow (snippet) looks like:

val webSocketFlow =
  Http().webSocketClientFlow(WebSocketRequest(url), settings = customSettings)

val (upgradeResponse, closed) =
  outgoing
    .viaMat(webSocketFlow)(Keep.right)
    .viaMat(decoder)(Keep.left)
    .toMat(sink)(Keep.both)
    .run()

This currently works great if I have one url. I am curious about how can I scale this to connect to multiple urls. So for example, if I have an indefinite list of websockets endpoints List("ws://localhost:8080/foo", "ws://localhost:8080/bar", "ws://localhost:8080/baz") .

I have considered adding a new flow for each URL, but what if I have a long list of websockets endpoints/urls. Then that becomes cumbersome and overtly-manual. I have also considered wrapping this into a function and calling for each URL in a given iterable. But that also felt-over-kill.

Is there a way to have a pool of connections all sourced into one Flow (or something like this)? Further readings are also welcome. As a "nice-to-have", is there also a way to tag the incoming messages to signal with url they are coming from?

Update: for clarification, I am only reading from websockets (only client side) and not sending any messages back.

Answer 1

This should work (code is being written in the textbox...):

def taggedWebsocketForUrl(url: String, tag: Int): Source[(Int, Message), Future[WebSocketUpgradeResponse]] =
  outgoing.viaMat(Http().webSocketClientFlow(WebSocketRequest(url), settings = customSettings))(Keep.right).map(tag -> _)

val websocketMergedSource: Source[(Int, Message), Seq[Future[WebSocketUpgradeResponse]]] = {
  // You could replace this with a mess of headOptions etc., but...
  if (websocketUrls.isEmpty) Source.empty[(Int, Message)].mapMaterializedValue(_ => Seq(Future.failed(new NoSuchElementException("no websocket URLs"))))
  else {
    val first: Source[(Int, Message), List[Future[WebSocketUpgradeResponse]]] =
      taggedWebsocketForUrl(websocketUrls.head, 0).mapMaterializedValue(List(_))
    if (websocketUrls.tail.isEmpty) first
    else {
      websocketUrls.tail.foldLeft(first -> 1) {
        (acc, url) =>
          val newSource = acc._1.mergeMat(taggedWebsocketForUrl(url, acc._2)) {
            (futs: List[Future[WebSocketUpgradeResponse]], fut: Future[WebSocketUpgradeResponse]) =>
              fut :: futs // Will reverse at the end...
          }
          newSource -> (acc._2 + 1)
      }._1.mapMaterializedValue(_.reverse)
    }
  }
}

With this, you'll have many upgrade responses (you could mapMaterializedValue(Future.sequence _) to combine them into a Future[Seq[WebsocketUpgradeResponse]] which will fail if any fail). The messages from the n th url in the list will be tagged with n .

Note that websocketUrls being a List guides to building up as a fold: if there are n urls, the messages from the first url will go through n -1 merge stages and the last url will go through only 1 merge stage, so you'd want to put the urls you expect to generate more traffic towards the end of the list.

An alternative, more efficient approach would be if using an IndexedSeq like Vector or Array to divide and conquer to build up a tree of merge s.

Using the Akka Streams GraphDSL would also give you a lot of control, but I would tend to use that only as a last resort.

Answer 2

What you need is some way to merge the various WebSocket flows so that you can process the incoming messages as if they came from a single Source.

Since you do not require to send any data but only receiving the implementation is straightforward.

Let's start creating a function that will create a WebSocket Source for a given uri:

def webSocketSource(uri: Uri): Source[Message, Future[WebSocketUpgradeResponse]] = {
  Source.empty.viaMat(Http().webSocketClientFlow(uri))(Keep.right)
}

Since you do not care about sending data, the function immediately close the out channel by providing an empty Source. The result is a Source containing the message read from the WebSocket.

At this point we can use this function to create a dedicated source for each uri:

val wsSources: List[Source[Message, NotUsed]] = uris.map { uri =>
  webSocketSource(uri).mapMaterializedValue { respFuture =>
    respFuture.map {
      case _: ValidUpgrade => log.debug(s"Websocket upgrade for [${uri}] successful")
      case err: InvalidUpgradeResponse => log.error(s"Websocket upgrade for [${uri}] failed: ${err.cause}")
    }

    NotUsed
  }
}

Here we need to somehow take care of the materialized values since it's not possible (or at least not easy) to combine them given we do not know how many they are. So here we go with the simplest approach of just logging.

Now that we have our sources ready we can proceed to merge them:

val mergedSource: Source[Message, NotUsed] = wsSources match {
  case s1 :: s2 :: rest => Source.combine(s1, s2, rest: _*)(Merge(_))
  case s1 :: Nil => s1
  case Nil => Source.empty[Message]
}

Here the idea is that in case we have 2 or more uris, we actually do a merge operation, otherwise if we have a single one we just use it without any modification. Lastly we also cover the case where we do not have any uri at all by providing an empty Source which will simply terminate the stream without errors.

At this point we can combine this source to the flows and sink you already have and run it:

val done: Future[Done] = mergedSource.via(decoder).toMat(sink)(Keep.right).run

Which gives us back a single future which will be completed when all connection are completed or failed as soon as one connection fails.