Task manager

task-manager.scala

//> using scala 3
//> using dep "dev.zio::zio:2.1.6"

import zio.*
import zio.internal.FiberRuntime
import java.util.UUID

trait TaskManager {
    def add(task: Task[Unit]): UIO[UUID]
    def stop(id: UUID): UIO[Unit]
    def getAll: UIO[Set[UUID]]
}

case class InMemoryTaskManager(ref: Ref[Map[UUID, Fiber.Runtime[?, Unit]]]) extends TaskManager {

  override def add(task: Task[Unit]): UIO[UUID] = {
    val id = UUID.randomUUID()
    task
        .ensuring(ref.update(_ - id))
        .fork
        .flatMap(f => ref.update(_.updated(id, f)))
        .as(id)
  }
    

  override def getAll: UIO[Set[UUID]] = 
    ref.get.map(_.keySet)

  override def stop(id: UUID): UIO[Unit] = 
    ref.get.flatMap(map => map.get(id).fold(ZIO.logError("id not found"))(_.interrupt)) *>
     ref.update(_ - id)
}
object InMemoryTaskManager:
    def layer = ZLayer(
        Ref.make[Map[UUID, Fiber.Runtime[?, Unit]]](Map.empty).map(InMemoryTaskManager(_))
    )

// usage

object AnApp extends ZIOAppDefault {

    val app = 
        for {
            manager <- ZIO.service[TaskManager]
            _ <- manager.add(ZIO.log("Starting, will take 1 second") *> ZIO.sleep(2.second) *> ZIO.log("...done"))
            foreverTask <- manager.add(ZIO.log("Oh no I run forever!") *> ZIO.never)
            _ <- manager.getAll.map(tasks => assert(tasks.size == 2, "the manager now has 2 tasks"))
            _ <- ZIO.log("waiting 5 seconds")
            _ <- ZIO.sleep(5.seconds)
            _ <- ZIO.log("Tired of waiting, going to stop that long task")
            _ <- manager.stop(foreverTask)
            _ <- manager.getAll.map(tasks => assert(tasks.size == 0, "All tasks are interrupted or done"))
        } yield ()

    def run = app.provide(
        InMemoryTaskManager.layer
    )
        
}