ScalaCLI + ZIO + Selenium Example

ZIOSeleniumExample.scala

//> using scala "3.2.1"
//> using lib "dev.zio::zio:2.0.3"

import zio.*


case class AppConfig(root: String)

object AppConfig:
  val layer: ZLayer[Any, Nothing, AppConfig] =
    ZLayer.succeed(AppConfig("zio.dev"))


class WebDriver {
  def doHeavyJob(url: String): Unit = {
    println(s"Doing some heavy job using WebDriver on URL ${url}...")
    Thread.sleep(1000)
    println(s"Heavy job on ${url} completed.")
  }

  def close(): Unit = {
    println("Closing WebDriver...")
  }
}

class WebDriverPool(pool: ZPool[Throwable, WebDriver]):
  def get: RIO[Scope, WebDriver] = pool.get

object WebDriverPool:
  def get: RIO[WebDriverPool & Scope, WebDriver] =
    ZIO.serviceWithZIO[WebDriverPool](_.get)

  val layer: ZLayer[Any, Throwable, WebDriverPool] =
    ZLayer.scoped {
      def make: RIO[Scope, WebDriver] =
        ZIO.acquireRelease(
          ZIO.attempt {
            println("Making WebDriver...")
            WebDriver() 
          }
        )( driver =>
          ZIO.attempt(driver.close()).ignoreLogged
        )

      for pool <- ZPool.make(
          get = make,
          range = 0 to 4,
          timeToLive = 1.hour
        )
      yield WebDriverPool(pool)
    }


object SeleniumExample:

  enum DataSource:
    case Simple(some: Int, thing: Double)
    case Complex(url: String)

  type Result = Unit

  def loadSourceList(root: String): Task[Seq[DataSource]] =
    ZIO.foreach(1 to 1000) { i =>
      Random.nextDouble.map { d =>
        if (d < 0.01) DataSource.Complex(s"Complex source $i")
        else DataSource.Simple(i, d)
      }
    }

  def saveResults(results: Seq[Result]): Task[Unit] =
    ZIO.succeed(println(s"Saving ${results.length} results..."))
    

  def doSimpleJob(some: Int, thing: Double): UIO[Result] =
    ZIO.succeed { /* do something simple */ }

  def fetchFromComplexSPA(url: String): RIO[WebDriverPool, Result] =
    ZIO.scoped {
      for
        driver <- WebDriverPool.get
        result <- ZIO.attemptBlocking { driver.doHeavyJob(url) }
      yield result
    }


  val program: RIO[AppConfig & WebDriverPool, Unit] =
    for
      _ <- ZIO.succeed(println("Beginning program..."))

      config <- ZIO.service[AppConfig]

      urls <- loadSourceList(config.root)

      results <- ZIO.foreachPar(urls) {
        case DataSource.Simple(s, t) => doSimpleJob(s, t)
        case DataSource.Complex(url) => fetchFromComplexSPA(url)
      } @@ ZIOAspect.parallel(4)

      _ <- saveResults(results)

      _ <- ZIO.succeed(println("Done."))
    yield ()


object Main extends ZIOAppDefault:

  val run =
    SeleniumExample.program
      .provide(
        AppConfig.layer,
        WebDriverPool.layer
      )