tech.ytsaurus.spark.launcher.rest.ByopDiscoveryService.scala Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of spark-yt-cluster_2.12 Show documentation

Spark over YTsaurus

The newest version!

package tech.ytsaurus.spark.launcher.rest

import io.circe._
import io.circe.generic.semiauto.deriveEncoder
import io.circe.parser._
import tech.ytsaurus.spyt.wrapper.Utils.flatten
import tech.ytsaurus.spyt.wrapper.cypress.{YsonWriter, YsonableProduct}
import sttp.client._
import tech.ytsaurus.spyt.HostAndPort

import scala.util.{Failure, Try}

class ByopDiscoveryService(masterEndpoint: String, byopPort: Int) {

  import ByopDiscoveryService._

  private val baseUri = s"http://$masterEndpoint/v1/submissions"

  def discoveryInfo: DiscoveryInfo = {
    DiscoveryInfo(proxiesList.get)
  }

  private def proxiesList: Try[Seq[String]] = {
    implicit val backend = HttpURLConnectionBackend()
    basicRequest
      .get(uri"$baseUri/master")
      .send()
      .body
      .fold[Try[Seq[String]]](
        error => Failure(HttpError(error)),
        body => parseWorkersList(body).toTry.map(workersToProxies(_, byopPort))
      )
  }
}

case class DiscoveryInfo(proxies: Seq[String] = Nil)

object DiscoveryInfo {
  implicit val encoder: Encoder[DiscoveryInfo] = deriveEncoder
  implicit val ysonWriter: YsonWriter[DiscoveryInfo] = YsonableProduct.ysonWriter[DiscoveryInfo]
}

case class WorkerInfo(host: String, alive: Boolean)

object ByopDiscoveryService {
  def appendPort(host: String, port: Int): String = {
    HostAndPort(host, port).toString
  }

  def parseWorkerInfo(worker: Json): Either[Error, WorkerInfo] = {
    val workerCursor = worker.hcursor
    for {
      host <- workerCursor.downField("host").as[String]
      alive <- workerCursor.downField("alive").as[Boolean]
    } yield WorkerInfo(host, alive)
  }

  def parseWorkersList(body: String): Either[Error, Seq[WorkerInfo]] = {
    for {
      cursor <- parse(body).map(_.hcursor)
      workers <- cursor.downField("workers").as[Array[Json]]
      workersHosts <- flatten(workers.map(parseWorkerInfo))
    } yield workersHosts
  }

  def workersToProxies(workers: Seq[WorkerInfo], byopPort: Int): Seq[String] = {
    workers.flatMap {
      case WorkerInfo(host, true) => Some(appendPort(host, byopPort))
      case _ => None
    }
  }
}