graphql-ws.scala

package powermonitor.graphql

import cats.effect.*
import cats.syntax.all.*
import cats.effect.std.Queue
import scala.concurrent.duration.*
import fs2.Stream
import io.circe.{Json, parser}
import io.circe.syntax.*
import org.http4s.{Header, Headers, HttpRoutes, Response}
import org.http4s.dsl.Http4sDsl
import org.http4s.server.websocket.WebSocketBuilder2
import org.http4s.websocket.WebSocketFrame
import org.typelevel.ci.*
import org.typelevel.log4cats.Logger
import grackle.Mapping

// Implements the graphql-transport-ws protocol (the "graphql-ws" library's
// modern subprotocol) over a WebSocket. This is what Postman and Apollo use for
// GraphQL subscriptions; plain SSE is not understood by those clients.
//
// Flow: client connects with subprotocol graphql-transport-ws and sends
// connection_init; the server replies connection_ack. The client then sends
// `subscribe` messages (each with an id); the server streams `next` frames and a
// final `complete`. A client `complete` cancels that subscription.
object GraphQLWebSocket:

  def routes(wsb: WebSocketBuilder2[IO], mapping: Mapping[IO], logger: Logger[IO]): HttpRoutes[IO] =
    val dsl = new Http4sDsl[IO] {}
    import dsl.*
    HttpRoutes.of[IO] { case GET -> Root / "graphql" => handler(wsb, mapping, logger) }

  private type Subs = Ref[IO, Map[String, Fiber[IO, Throwable, Unit]]]

  private def text(j: Json): WebSocketFrame = WebSocketFrame.Text(j.noSpaces)
  private val ackMsg = Json.obj("type" -> "connection_ack".asJson)
  private val pongMsg = Json.obj("type" -> "pong".asJson)
  private def nextMsg(id: String, payload: Json) =
    Json.obj("id" -> id.asJson, "type" -> "next".asJson, "payload" -> payload)
  private def completeMsg(id: String) =
    Json.obj("id" -> id.asJson, "type" -> "complete".asJson)
  private def errorMsg(id: String, msg: String) =
    Json.obj("id" -> id.asJson, "type" -> "error".asJson, "payload" -> Json.arr(Json.obj("message" -> msg.asJson)))

  private def handler(wsb: WebSocketBuilder2[IO], mapping: Mapping[IO], logger: Logger[IO]): IO[Response[IO]] =
    for
      out <- Queue.unbounded[IO, WebSocketFrame]
      subs <- IO.ref(Map.empty[String, Fiber[IO, Throwable, Unit]])
      // A periodic Ping keeps the connection alive past Ember's idle timeout
      // while a subscription waits between (possibly rare) events; the client
      // auto-replies Pong, which the receive pipe filters out.
      heartbeat = Stream.awakeEvery[IO](20.seconds).as(WebSocketFrame.Ping())
      resp <- wsb
        .withHeaders(Headers(Header.Raw(ci"Sec-WebSocket-Protocol", "graphql-transport-ws")))
        .build(
          Stream.fromQueueUnterminated(out).merge(heartbeat),
          _.foreach(frame => handle(mapping, logger, out, subs, frame)).onFinalize(cancelAll(subs))
        )
    yield resp

  private def cancelAll(subs: Subs): IO[Unit] =
    subs.getAndSet(Map.empty).flatMap(_.values.toList.traverse_(_.cancel))

  private def handle(mapping: Mapping[IO], logger: Logger[IO], out: Queue[IO, WebSocketFrame], subs: Subs, frame: WebSocketFrame): IO[Unit] =
    frame match
      case t: WebSocketFrame.Text =>
        parser.parse(t.str).toOption.flatMap(_.asObject) match
          case None => IO.unit
          case Some(obj) =>
            obj("type").flatMap(_.asString) match
              case Some("connection_init") => out.offer(text(ackMsg))
              case Some("ping") => out.offer(text(pongMsg))
              case Some("subscribe") => startSub(mapping, logger, out, subs, obj)
              case Some("complete") => obj("id").flatMap(_.asString).traverse_(stopSub(subs, _))
              case _ => IO.unit // pong, connection ack from peer, etc.
      case _ => IO.unit

  private def startSub(
      mapping: Mapping[IO],
      logger: Logger[IO],
      out: Queue[IO, WebSocketFrame],
      subs: Subs,
      obj: io.circe.JsonObject
  ): IO[Unit] =
    (obj("id").flatMap(_.asString), obj("payload").flatMap(_.asObject)) match
      case (Some(id), Some(payload)) =>
        payload("query").flatMap(_.asString) match
          case None => out.offer(text(errorMsg(id, "Missing query")))
          case Some(query) =>
            val op = payload("operationName").flatMap(_.asString)
            val vars = payload("variables")
            val run =
              mapping
                .compileAndRunSubscription(query, op, vars)
                .evalMap(json => out.offer(text(nextMsg(id, json))))
                .compile
                .drain
                .attempt
                .flatMap {
                  case Right(_) => out.offer(text(completeMsg(id)))
                  case Left(e) =>
                    logger.warn(e)(s"WS subscription $id failed") *>
                      out.offer(text(errorMsg(id, Option(e.getMessage).getOrElse(e.toString))))
                }
            run.start.flatMap(fib => subs.update(_.updated(id, fib)))
      case _ => IO.unit

  private def stopSub(subs: Subs, id: String): IO[Unit] =
    subs.modify(m => (m - id, m.get(id))).flatMap(_.traverse_(_.cancel))