<commit_before>package lila.socket

import akka.actor.{ ActorSystem, CoordinatedShutdown }
import chess.{ Centis, Color }
import io.lettuce.core._
import io.lettuce.core.pubsub.{ StatefulRedisPubSubConnection => PubSub }
import java.util.concurrent.atomic.AtomicReference
import java.util.concurrent.ConcurrentHashMap
import play.api.libs.json._
import scala.concurrent.duration._
import scala.concurrent.{ Future, Promise }
import scala.util.chaining._
import Socket.Sri

import lila.common.{ Bus, Lilakka }
import lila.hub.actorApi.Announce
import lila.hub.actorApi.relation.{ Follow, UnFollow }
import lila.hub.actorApi.round.Mlat
import lila.hub.actorApi.security.CloseAccount
import lila.hub.actorApi.socket.remote.{ TellSriIn, TellSriOut, TellUserIn }
import lila.hub.actorApi.socket.{ ApiUserIsOnline, SendTo, SendToAsync, SendTos }

final class RemoteSocket(
    redisClient: RedisClient,
    shutdown: CoordinatedShutdown
)(implicit
    ec: scala.concurrent.ExecutionContext,
    system: ActorSystem
) {

  import RemoteSocket._, Protocol._

  private var stopping = false

  private type UserIds = Set[String]

  private val requests = new ConcurrentHashMap[Int, Promise[String]](32)

  def request[R](sendReq: Int => Unit, readRes: String => R): Fu[R] = {
    val id = lila.common.ThreadLocalRandom.nextInt()
    sendReq(id)
    val promise = Promise[String]()
    requests.put(id, promise)
    promise.future map readRes
  }

  val onlineUserIds: AtomicReference[Set[String]] = new AtomicReference(Set("lichess"))

  val baseHandler: Handler = {
    case In.ConnectUser(userId) =>
      onlineUserIds.getAndUpdate(_ + userId).unit
    case In.DisconnectUsers(userIds) =>
      onlineUserIds.getAndUpdate(_ -- userIds).unit
    case In.NotifiedBatch(userIds) =>
      Bus.publish(lila.hub.actorApi.notify.NotifiedBatch(userIds), "notify")
    case In.Lags(lags) =>
      lags foreach (UserLagCache.put _).tupled
      // this shouldn't be necessary... ensure that users are known to be online
      onlineUserIds.getAndUpdate((x: UserIds) => x ++ lags.keys).unit
    case In.TellSri(sri, userId, typ, msg) =>
      Bus.publish(TellSriIn(sri.value, userId, msg), s"remoteSocketIn:$typ")
    case In.TellUser(userId, typ, msg) =>
      Bus.publish(TellUserIn(userId, msg), s"remoteSocketIn:$typ")
    case In.ReqResponse(reqId, response) =>
      requests
        .computeIfPresent(
          reqId,
          (_: Int, promise: Promise[String]) => {
            promise success response
            null // remove from promises
          }
        )
        .unit
    case In.Ping(id) => send(Out.pong(id))
    case In.WsBoot =>
      logger.warn("Remote socket boot")
      onlineUserIds set Set("lichess")
  }

  Bus.subscribeFun(
    "socketUsers",
    "announce",
    "mlat",
    "sendToFlag",
    "remoteSocketOut",
    "accountClose",
    "shadowban",
    "impersonate",
    "relation",
    "onlineApiUsers"
  ) {
    case SendTos(userIds, payload) =>
      val connectedUsers = userIds intersect onlineUserIds.get
      if (connectedUsers.nonEmpty) send(Out.tellUsers(connectedUsers, payload))
    case SendTo(userId, payload) =>
      if (onlineUserIds.get.contains(userId)) send(Out.tellUser(userId, payload))
    case SendToAsync(userId, makePayload) =>
      if (onlineUserIds.get.contains(userId)) makePayload() foreach { payload =>
        send(Out.tellUser(userId, payload))
      }
    case Announce(_, _, json) =>
      send(Out.tellAll(Json.obj("t" -> "announce", "d" -> json)))
    case Mlat(millis) =>
      send(Out.mlat(millis))
    case Socket.SendToFlag(flag, payload) =>
      send(Out.tellFlag(flag, payload))
    case TellSriOut(sri, payload) =>
      send(Out.tellSri(Sri(sri), payload))
    case CloseAccount(userId) =>
      send(Out.disconnectUser(userId))
    case lila.hub.actorApi.mod.Shadowban(userId, v) =>
      send(Out.setTroll(userId, v))
    case lila.hub.actorApi.mod.Impersonate(userId, modId) =>
      send(Out.impersonate(userId, modId))
    case ApiUserIsOnline(userId, value) =>
      send(Out.apiUserOnline(userId, value))
      if (value) onlineUserIds.getAndUpdate(_ + userId).unit
    case Follow(u1, u2)   => send(Out.follow(u1, u2))
    case UnFollow(u1, u2) => send(Out.unfollow(u1, u2))
  }

  final class StoppableSender(val conn: PubSub[String, String], channel: Channel) extends Sender {
    def apply(msg: String): Unit               = if (!stopping) logAndSend(channel, msg)
    def sticky(_id: String, msg: String): Unit = apply(msg)
  }

  final class RoundRobinSender(val conn: PubSub[String, String], channel: Channel, parallelism: Int)
      extends Sender {
    def apply(msg: String): Unit = publish(msg.hashCode.abs % parallelism, msg)
    // use the ID to select the channel, not the entire message
    def sticky(id: String, msg: String): Unit = publish(id.hashCode.abs % parallelism, msg)

    private def publish(subChannel: Int, msg: String) =
      if (!stopping) logAndSend(s"$channel:$subChannel", msg)
  }

  def makeSender(channel: Channel, parallelism: Int = 1): Sender =
    if (parallelism > 1) new RoundRobinSender(redisClient.connectPubSub(), channel, parallelism)
    else new StoppableSender(redisClient.connectPubSub(), channel)

  private val send: Send = makeSender("site-out").apply _

  def subscribe(channel: Channel, reader: In.Reader)(handler: Handler): Funit =
    connectAndSubscribe(channel) { str =>
      RawMsg(str) match {
        case None => logger.error(s"Invalid $channel $str")
        case Some(msg) =>
          lila.mon.ws.in(channel, msg.path).increment()
          reader(msg) collect handler match {
            case Some(_) => // processed
            case None    => logger.warn(s"Unhandled $channel $str")
          }
      }
    }

  def subscribeRoundRobin(channel: Channel, reader: In.Reader, parallelism: Int)(
      handler: Handler
  ): Funit =
    // subscribe to main channel
    subscribe(channel, reader)(handler) >> {
      // and subscribe to subchannels
      (0 to parallelism)
        .map { index =>
          subscribe(s"$channel:$index", reader)(handler)
        }
        .sequenceFu
        .void
    }

  private def connectAndSubscribe(channel: Channel)(f: String => Unit): Funit = {
    val conn = redisClient.connectPubSub()
    conn.addListener(new pubsub.RedisPubSubAdapter[String, String] {
      override def message(_channel: String, message: String): Unit = f(message)<commit_msg>don't monitor round ws channels<commit_after>package lila.socket

import akka.actor.{ ActorSystem, CoordinatedShutdown }
import chess.{ Centis, Color }
import io.lettuce.core._
import io.lettuce.core.pubsub.{ StatefulRedisPubSubConnection => PubSub }
import java.util.concurrent.atomic.AtomicReference
import java.util.concurrent.ConcurrentHashMap
import play.api.libs.json._
import scala.concurrent.duration._
import scala.concurrent.{ Future, Promise }
import scala.util.chaining._
import Socket.Sri

import lila.common.{ Bus, Lilakka }
import lila.hub.actorApi.Announce
import lila.hub.actorApi.relation.{ Follow, UnFollow }
import lila.hub.actorApi.round.Mlat
import lila.hub.actorApi.security.CloseAccount
import lila.hub.actorApi.socket.remote.{ TellSriIn, TellSriOut, TellUserIn }
import lila.hub.actorApi.socket.{ ApiUserIsOnline, SendTo, SendToAsync, SendTos }

final class RemoteSocket(
    redisClient: RedisClient,
    shutdown: CoordinatedShutdown
)(implicit
    ec: scala.concurrent.ExecutionContext,
    system: ActorSystem
) {

  import RemoteSocket._, Protocol._

  private var stopping = false

  private type UserIds = Set[String]

  private val requests = new ConcurrentHashMap[Int, Promise[String]](32)

  def request[R](sendReq: Int => Unit, readRes: String => R): Fu[R] = {
    val id = lila.common.ThreadLocalRandom.nextInt()
    sendReq(id)
    val promise = Promise[String]()
    requests.put(id, promise)
    promise.future map readRes
  }

  val onlineUserIds: AtomicReference[Set[String]] = new AtomicReference(Set("lichess"))

  val baseHandler: Handler = {
    case In.ConnectUser(userId) =>
      onlineUserIds.getAndUpdate(_ + userId).unit
    case In.DisconnectUsers(userIds) =>
      onlineUserIds.getAndUpdate(_ -- userIds).unit
    case In.NotifiedBatch(userIds) =>
      Bus.publish(lila.hub.actorApi.notify.NotifiedBatch(userIds), "notify")
    case In.Lags(lags) =>
      lags foreach (UserLagCache.put _).tupled
      // this shouldn't be necessary... ensure that users are known to be online
      onlineUserIds.getAndUpdate((x: UserIds) => x ++ lags.keys).unit
    case In.TellSri(sri, userId, typ, msg) =>
      Bus.publish(TellSriIn(sri.value, userId, msg), s"remoteSocketIn:$typ")
    case In.TellUser(userId, typ, msg) =>
      Bus.publish(TellUserIn(userId, msg), s"remoteSocketIn:$typ")
    case In.ReqResponse(reqId, response) =>
      requests
        .computeIfPresent(
          reqId,
          (_: Int, promise: Promise[String]) => {
            promise success response
            null // remove from promises
          }
        )
        .unit
    case In.Ping(id) => send(Out.pong(id))
    case In.WsBoot =>
      logger.warn("Remote socket boot")
      onlineUserIds set Set("lichess")
  }

  Bus.subscribeFun(
    "socketUsers",
    "announce",
    "mlat",
    "sendToFlag",
    "remoteSocketOut",
    "accountClose",
    "shadowban",
    "impersonate",
    "relation",
    "onlineApiUsers"
  ) {
    case SendTos(userIds, payload) =>
      val connectedUsers = userIds intersect onlineUserIds.get
      if (connectedUsers.nonEmpty) send(Out.tellUsers(connectedUsers, payload))
    case SendTo(userId, payload) =>
      if (onlineUserIds.get.contains(userId)) send(Out.tellUser(userId, payload))
    case SendToAsync(userId, makePayload) =>
      if (onlineUserIds.get.contains(userId)) makePayload() foreach { payload =>
        send(Out.tellUser(userId, payload))
      }
    case Announce(_, _, json) =>
      send(Out.tellAll(Json.obj("t" -> "announce", "d" -> json)))
    case Mlat(millis) =>
      send(Out.mlat(millis))
    case Socket.SendToFlag(flag, payload) =>
      send(Out.tellFlag(flag, payload))
    case TellSriOut(sri, payload) =>
      send(Out.tellSri(Sri(sri), payload))
    case CloseAccount(userId) =>
      send(Out.disconnectUser(userId))
    case lila.hub.actorApi.mod.Shadowban(userId, v) =>
      send(Out.setTroll(userId, v))
    case lila.hub.actorApi.mod.Impersonate(userId, modId) =>
      send(Out.impersonate(userId, modId))
    case ApiUserIsOnline(userId, value) =>
      send(Out.apiUserOnline(userId, value))
      if (value) onlineUserIds.getAndUpdate(_ + userId).unit
    case Follow(u1, u2)   => send(Out.follow(u1, u2))
    case UnFollow(u1, u2) => send(Out.unfollow(u1, u2))
  }

  final class StoppableSender(val conn: PubSub[String, String], channel: Channel) extends Sender {
    def apply(msg: String): Unit               = if (!stopping) logAndSend(channel, msg)
    def sticky(_id: String, msg: String): Unit = apply(msg)
  }

  final class RoundRobinSender(val conn: PubSub[String, String], channel: Channel, parallelism: Int)
      extends Sender {
    def apply(msg: String): Unit = publish(msg.hashCode.abs % parallelism, msg)
    // use the ID to select the channel, not the entire message
    def sticky(id: String, msg: String): Unit = publish(id.hashCode.abs % parallelism, msg)

    private def publish(subChannel: Int, msg: String) =
      if (!stopping) conn.async.publish(s"$channel:$subChannel", msg).unit
  }

  def makeSender(channel: Channel, parallelism: Int = 1): Sender =
    if (parallelism > 1) new RoundRobinSender(redisClient.connectPubSub(), channel, parallelism)
    else new StoppableSender(redisClient.connectPubSub(), channel)

  private val send: Send = makeSender("site-out").apply _

  def subscribe(channel: Channel, reader: In.Reader, monitor: Boolean = true)(handler: Handler): Funit =
    connectAndSubscribe(channel) { str =>
      RawMsg(str) match {
        case None => logger.error(s"Invalid $channel $str")
        case Some(msg) =>
          if (monitor) lila.mon.ws.in(channel, msg.path).increment()
          reader(msg) collect handler match {
            case Some(_) => // processed
            case None    => logger.warn(s"Unhandled $channel $str")
          }
      }
    }

  def subscribeRoundRobin(channel: Channel, reader: In.Reader, parallelism: Int)(
      handler: Handler
  ): Funit =
    // subscribe to main channel
    subscribe(channel, reader, monitor = false)(handler) >> {
      // and subscribe to subchannels
      (0 to parallelism)
        .map { index =>
          subscribe(s"$channel:$index", reader)(handler)
        }
        .sequenceFu
        .void
    }

  private def connectAndSubscribe(channel: Channel)(f: String => Unit): Funit = {
    val conn = redisClient.connectPubSub()
    conn.addListener(new pubsub.RedisPubSubAdapter[String, String] {
      override def message(_channel: String, message: String): Unit = f(message)
{'fraction_non_alphanumeric': 0.1247104247104247, 'fraction_numerical': 0.002007722007722008, 'mean_word_length': 3.4428816466552314, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 83, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5096532', 'n_tokens_mistral': 4115, 'n_tokens_neox': 4057, 'n_words': 936}