zio.stream.ZStream Scala Examples

package io.univalence.sparkzio

import zio.{ DefaultRuntime, IO, Ref, Task, UIO, ZIO, ZManaged }
import zio.clock.Clock
import zio.stream.{ Stream, ZSink, ZStream }
import zio.test.DefaultRunnableSpec
import zio.test._
import zio.test.Assertion._

object StreamTest {

  def assertForAll[R, E, A](zstream: ZStream[R, E, A])(f: A => TestResult): ZIO[R, E, TestResult] =
    zstream.fold(assert(Unit, Assertion.anything))((as, a) => as && f(a))

  def isSorted[A: Ordering]: Assertion[Iterable[A]] =
    Assertion.assertion("sorted")()(x => {
      val y = x.toList
      y.sorted == y
    })
}

object IteratorTest
    extends DefaultRunnableSpec(
      suite("iterator")(
        testM("to iterator should be lazy")({
          case class Element(n: Int, time: Long)

          (for {
            clock      <- ZIO.environment[Clock]
            n          <- Ref.make(0)
            incCounter <- n.update(_ + 1).forever.fork

          } yield {
            def element: UIO[Element] = n.get.zipWith(clock.clock.nanoTime)(Element)

            val in = Stream.repeatEffect(element)

            val iterator = Iterator.unwrapManaged(Iterator.fromStream(in))

            val out: ZStream[Any, Nothing, Element] =
              ZStream.fromIterator(iterator).mapConcatM(e => element.map(List(e, _)))

            implicit val ordering: Ordering[Element] = Ordering.by(x => x.n -> x.time)

            out.take(2000).runCollect.map(e => assert(e, StreamTest.isSorted))
          }).flatten
        }),
        testM("<=>")({
          val in: List[Int] = (1 to 100).toList

          (for {
            _ <- UIO.unit.toManaged_
            stream1 = ZStream.fromIterator(UIO(in.toIterator))
            iterator <- Iterator.fromStream(stream1)
            stream2 = ZStream.fromIterator(UIO(iterator))
            out <- stream2.runCollect.toManaged_
          } yield {
            assert(in, equalTo(out))
          }).use(x => ZIO.effect(x))

        }),
        testM("on exit")(
          (for {
            isOpen <- Ref.make(false).toManaged_
            stream = ZStream.managed(ZManaged.make(isOpen.update(_ => true))(_ => isOpen.set(false)))
            iterator <- Iterator.fromStream(stream)
          } yield {
            assert(iterator.toList, equalTo(List(true)))
          }).use(x => IO.effect(x))
        )
      )
    ) 

package io.univalence.sparkzio

import zio.{ Fiber, Promise, Queue, Ref, Semaphore, UIO, ZIO }
import zio.stream.ZStream
import zio.stream.ZStream.Pull

object ZDynamicConcurrency {

  private def releaseCapacity(semaphore: Semaphore, toRelease: UIO[Int]): UIO[Fiber[Nothing, Nothing]] =
    (for {
      concurrency <- Ref.make(0)
    } yield {
      def update(next: Int): UIO[Unit] =
        (for {
          current <- concurrency.get
          _       <- if (next >= current) semaphore.releaseN(next - current) else semaphore.acquireN(current - next)
          _       <- concurrency.set(next)
        } yield {}).uninterruptible

      val nextConcurrencyBounded: UIO[Int] = toRelease.map({
        case x if x < 0 => 0
        case x          => x
      })

      (nextConcurrencyBounded >>= update).forever.ensuring(update(0)).fork
    }).flatten

  implicit class Ops[-R, +E, +A](stream: ZStream[R, E, A]) {
    def mapMDynamicPar[R1 <: R, E1 >: E, B](concurrencyLevel: UIO[Int],
                                            maxConcurrency: Int = 16)(f: A => ZIO[R1, E1, B]): ZStream[R1, E1, B] =
      ZStream[R1, E1, B] {
        for {
          out                    <- Queue.bounded[Pull[R1, E1, B]](maxConcurrency).toManaged(_.shutdown)
          permits                <- Semaphore.make(permits = 0).toManaged_
          updateConcurrencyFiber <- releaseCapacity(permits, concurrencyLevel).toManaged_
          interruptWorkers       <- Promise.make[Nothing, Unit].toManaged_
          _ <- stream.foreachManaged { a =>
            for {
              latch <- Promise.make[Nothing, Unit]
              p     <- Promise.make[E1, B]
              _     <- out.offer(Pull.fromPromise(p))
              _     <- (permits.withPermit(latch.succeed(()) *> f(a).to(p)) race interruptWorkers.await).fork
              _     <- latch.await
            } yield ()
          }.foldCauseM(
              c => (interruptWorkers.succeed(()) *> out.offer(Pull.halt(c))).unit.toManaged_,
              _ => out.offer(Pull.end).unit.toManaged_
            )
            .ensuringFirst(interruptWorkers.succeed(()) *> updateConcurrencyFiber.interrupt)
            .fork
        } yield out.take.flatten
      }
  }
} 

package zio.interop.test

import cats.data.NonEmptyList
import cats.{ Bifunctor, Monad, MonadError, SemigroupK }
import zio._
import zio.interop.catz.core._
import zio.stream.interop.catz.core._
import zio.stream.{ Stream, ZStream }
import zio.test.Assertion._
import zio.test.{ DefaultRunnableSpec, test, _ }

object CoreSummonSpecextends extends DefaultRunnableSpec {
  override def spec =
    suite("summons from catz.core work with only a cats-core dependency")(
      test("ZIO instances") {
        val monad      = implicitly[Monad[UIO]]
        val monadError = implicitly[MonadError[Task, Throwable]]
        val semigroupK = implicitly[SemigroupK[IO[NonEmptyList[Unit], ?]]]
        val bifunctor  = implicitly[Bifunctor[IO]]

        monad.map(ZIO.unit)(identity)
        monadError.map(ZIO.unit)(identity)
        semigroupK.combineK(ZIO.unit, ZIO.unit)
        bifunctor.leftMap(ZIO.fromOption(None))(identity)

        assert(())(anything)
      },
      test("ZManaged instances") {
        val monad      = implicitly[Monad[ZManaged[Any, Nothing, ?]]]
        val monadError = implicitly[MonadError[Managed[Throwable, ?], Throwable]]
        val semigroupK = implicitly[SemigroupK[Managed[Nothing, ?]]]
        val bifunctor  = implicitly[Bifunctor[Managed]]

        monad.map(ZManaged.unit)(identity)
        monadError.map(ZManaged.unit)(identity)
        semigroupK.combineK(ZManaged.unit, ZManaged.unit)
        bifunctor.leftMap(ZManaged.fail(()))(identity)

        assert(())(anything)
      },
      test("ZStream instances") {
        val monad      = implicitly[Monad[ZStream[Any, Nothing, ?]]]
        val monadError = implicitly[MonadError[Stream[Throwable, ?], Throwable]]
        val semigroupK = implicitly[SemigroupK[Stream[Nothing, ?]]]
        val bifunctor  = implicitly[Bifunctor[Stream]]

        monad.map(ZStream.unit)(identity)
        monadError.map(ZStream.unit)(identity)
        semigroupK.combineK(ZStream.unit, ZStream.unit)
        bifunctor.leftMap(ZStream.fail(()))(identity)

        assert(())(anything)
      }
    )

} 

package caliban.schema

import caliban.CalibanError.ExecutionError
import caliban.Value.NullValue
import caliban.execution.FieldInfo
import caliban.{ InputValue, ResponseValue }
import zio.stream.ZStream
import zio.query.ZQuery

sealed trait Step[-R]

object Step {
  case class ListStep[-R](steps: List[Step[R]])                         extends Step[R]
  case class FunctionStep[-R](step: Map[String, InputValue] => Step[R]) extends Step[R]
  case class ObjectStep[-R](name: String, fields: Map[String, Step[R]]) extends Step[R]
  case class QueryStep[-R](query: ZQuery[R, Throwable, Step[R]])        extends Step[R]
  case class StreamStep[-R](inner: ZStream[R, Throwable, Step[R]])      extends Step[R]

  // PureStep is both a Step and a ReducedStep so it is defined outside this object
  // This is to avoid boxing/unboxing pure values during step reduction
  type PureStep = caliban.schema.PureStep
  val PureStep: caliban.schema.PureStep.type = caliban.schema.PureStep

  val NullStep: PureStep = PureStep(NullValue)

  
case class PureStep(value: ResponseValue) extends Step[Any] with ReducedStep[Any] 

package caliban.interop.monix

import caliban.introspection.adt.__Type
import caliban.schema.Step.{ QueryStep, StreamStep }
import caliban.schema.{ Schema, Step, Types }
import caliban.{ CalibanError, GraphQL, GraphQLInterpreter, GraphQLResponse, InputValue }
import cats.effect.ConcurrentEffect
import monix.eval.{ Task => MonixTask }
import monix.reactive.Observable
import zio._
import zio.interop.catz._
import zio.interop.reactivestreams._
import zio.stream.ZStream
import zio.query.ZQuery

object MonixInterop {

  def executeAsync[R, E](graphQL: GraphQLInterpreter[R, E])(
    query: String,
    operationName: Option[String] = None,
    variables: Map[String, InputValue] = Map(),
    extensions: Map[String, InputValue] = Map(),
    skipValidation: Boolean = false,
    enableIntrospection: Boolean = true
  )(implicit runtime: Runtime[R]): MonixTask[GraphQLResponse[E]] =
    MonixTask.async { cb =>
      val execution = graphQL.execute(
        query,
        operationName,
        variables,
        extensions,
        skipValidation = skipValidation,
        enableIntrospection = enableIntrospection
      )
      runtime.unsafeRunAsync(execution)(exit => cb(exit.toEither))
    }

  def checkAsync[R](graphQL: GraphQLInterpreter[R, Any])(query: String)(implicit runtime: Runtime[R]): MonixTask[Unit] =
    MonixTask.async(cb => runtime.unsafeRunAsync(graphQL.check(query))(exit => cb(exit.toEither)))

  def interpreterAsync[R](
    graphQL: GraphQL[R]
  )(implicit runtime: Runtime[R]): MonixTask[GraphQLInterpreter[R, CalibanError]] =
    MonixTask.async(cb => runtime.unsafeRunAsync(graphQL.interpreter)(exit => cb(exit.toEither)))

  def taskSchema[R, A](implicit ev: Schema[R, A], ev2: ConcurrentEffect[MonixTask]): Schema[R, MonixTask[A]] =
    new Schema[R, MonixTask[A]] {
      override def toType(isInput: Boolean, isSubscription: Boolean): __Type = ev.toType(isInput, isSubscription)
      override def optional: Boolean                                         = ev.optional
      override def resolve(value: MonixTask[A]): Step[R] =
        QueryStep(ZQuery.fromEffect(value.to[Task].map(ev.resolve)))
    }

  def observableSchema[R, A](
    queueSize: Int
  )(implicit ev: Schema[R, A], ev2: ConcurrentEffect[MonixTask]): Schema[R, Observable[A]] =
    new Schema[R, Observable[A]] {
      override def optional: Boolean = true
      override def toType(isInput: Boolean, isSubscription: Boolean): __Type = {
        val t = ev.toType(isInput, isSubscription)
        if (isSubscription) t else Types.makeList(if (ev.optional) t else Types.makeNonNull(t))
      }
      override def resolve(value: Observable[A]): Step[R] =
        StreamStep(
          ZStream
            .fromEffect(
              MonixTask
                .deferAction(implicit sc =>
                  MonixTask.eval(value.toReactivePublisher.toStream(queueSize).map(ev.resolve))
                )
                .to[Task]
            )
            .flatten
        )
    }
} 

package caliban.federation

import caliban.federation.FederationData.characters._
import zio.stream.ZStream
import zio.{ Has, Queue, Ref, UIO, URIO, ZLayer }

object CharacterService {

  type CharacterService = Has[Service]

  trait Service {
    def getCharactersByEpisode(season: Int, episode: Int): UIO[List[Character]]

    def getCharacters(origin: Option[Origin]): UIO[List[Character]]

    def findCharacter(name: String): UIO[Option[Character]]

    def deleteCharacter(name: String): UIO[Boolean]

    def deletedEvents: ZStream[Any, Nothing, String]
  }

  def getCharacters(origin: Option[Origin]): URIO[CharacterService, List[Character]] =
    URIO.accessM(_.get.getCharacters(origin))

  def findCharacter(name: String): URIO[CharacterService, Option[Character]] =
    URIO.accessM(_.get.findCharacter(name))

  def deleteCharacter(name: String): URIO[CharacterService, Boolean] =
    URIO.accessM(_.get.deleteCharacter(name))

  def deletedEvents: ZStream[CharacterService, Nothing, String] =
    ZStream.accessStream(_.get.deletedEvents)

  def getCharactersByEpisode(season: Int, episode: Int): URIO[CharacterService, List[Character]] =
    URIO.accessM(_.get.getCharactersByEpisode(season, episode))

  def make(initial: List[Character]): ZLayer[Any, Nothing, CharacterService] = ZLayer.fromEffect {
    for {
      characters  <- Ref.make(initial)
      subscribers <- Ref.make(List.empty[Queue[String]])
    } yield new Service {

      def getCharacters(origin: Option[Origin]): UIO[List[Character]] =
        characters.get.map(_.filter(c => origin.forall(c.origin == _)))

      def findCharacter(name: String): UIO[Option[Character]] = characters.get.map(_.find(c => c.name == name))

      def deleteCharacter(name: String): UIO[Boolean] =
        characters
          .modify(list =>
            if (list.exists(_.name == name)) (true, list.filterNot(_.name == name))
            else (false, list)
          )
          .tap(deleted =>
            UIO.when(deleted)(
              subscribers.get.flatMap(
                // add item to all subscribers
                UIO.foreach(_)(queue =>
                  queue
                    .offer(name)
                    .catchSomeCause {
                      case cause if cause.interrupted =>
                        subscribers.update(_.filterNot(_ == queue)).as(false)
                    } // if queue was shutdown, remove from subscribers
                )
              )
            )
          )

      def deletedEvents: ZStream[Any, Nothing, String] = ZStream.unwrap {
        for {
          queue <- Queue.unbounded[String]
          _     <- subscribers.update(queue :: _)
        } yield ZStream.fromQueue(queue).ensuring(queue.shutdown)
      }

      override def getCharactersByEpisode(season: Int, episode: Int): UIO[List[Character]] =
        characters.get.map(_.filter(c => c.starredIn.exists(e => e.episode == episode && e.season == season)))
    }
  }
} 

package caliban

import scala.language.postfixOps
import caliban.ExampleData._
import caliban.ExampleService.ExampleService
import caliban.GraphQL.graphQL
import caliban.schema.Annotations.{ GQLDeprecated, GQLDescription }
import caliban.schema.GenericSchema
import caliban.wrappers.ApolloTracing.apolloTracing
import caliban.wrappers.Wrappers.{ maxDepth, maxFields, printSlowQueries, timeout }
import zio.URIO
import zio.clock.Clock
import zio.console.Console
import zio.duration._
import zio.stream.ZStream

object ExampleApi extends GenericSchema[ExampleService] {

  case class Queries(
    @GQLDescription("Return all characters from a given origin")
    characters: CharactersArgs => URIO[ExampleService, List[Character]],
    @GQLDeprecated("Use `characters`")
    character: CharacterArgs => URIO[ExampleService, Option[Character]]
  )
  case class Mutations(deleteCharacter: CharacterArgs => URIO[ExampleService, Boolean])
  case class Subscriptions(characterDeleted: ZStream[ExampleService, Nothing, String])

  implicit val roleSchema           = gen[Role]
  implicit val characterSchema      = gen[Character]
  implicit val characterArgsSchema  = gen[CharacterArgs]
  implicit val charactersArgsSchema = gen[CharactersArgs]

  val api: GraphQL[Console with Clock with ExampleService] =
    graphQL(
      RootResolver(
        Queries(
          args => ExampleService.getCharacters(args.origin),
          args => ExampleService.findCharacter(args.name)
        ),
        Mutations(args => ExampleService.deleteCharacter(args.name)),
        Subscriptions(ExampleService.deletedEvents)
      )
    ) @@
      maxFields(200) @@               // query analyzer that limit query fields
      maxDepth(30) @@                 // query analyzer that limit query depth
      timeout(3 seconds) @@           // wrapper that fails slow queries
      printSlowQueries(500 millis) @@ // wrapper that logs slow queries
      apolloTracing // wrapper for https://github.com/apollographql/apollo-tracing

} 

package caliban

import caliban.ExampleData._
import zio.stream.ZStream
import zio.{ Has, Queue, Ref, UIO, URIO, ZLayer }

object ExampleService {

  type ExampleService = Has[Service]

  trait Service {
    def getCharacters(origin: Option[Origin]): UIO[List[Character]]

    def findCharacter(name: String): UIO[Option[Character]]

    def deleteCharacter(name: String): UIO[Boolean]

    def deletedEvents: ZStream[Any, Nothing, String]
  }

  def getCharacters(origin: Option[Origin]): URIO[ExampleService, List[Character]] =
    URIO.accessM(_.get.getCharacters(origin))

  def findCharacter(name: String): URIO[ExampleService, Option[Character]] =
    URIO.accessM(_.get.findCharacter(name))

  def deleteCharacter(name: String): URIO[ExampleService, Boolean] =
    URIO.accessM(_.get.deleteCharacter(name))

  def deletedEvents: ZStream[ExampleService, Nothing, String] =
    ZStream.accessStream(_.get.deletedEvents)

  def make(initial: List[Character]): ZLayer[Any, Nothing, ExampleService] = ZLayer.fromEffect {
    for {
      characters  <- Ref.make(initial)
      subscribers <- Ref.make(List.empty[Queue[String]])
    } yield new Service {

      def getCharacters(origin: Option[Origin]): UIO[List[Character]] =
        characters.get.map(_.filter(c => origin.forall(c.origin == _)))

      def findCharacter(name: String): UIO[Option[Character]] = characters.get.map(_.find(c => c.name == name))

      def deleteCharacter(name: String): UIO[Boolean] =
        characters
          .modify(list =>
            if (list.exists(_.name == name)) (true, list.filterNot(_.name == name))
            else (false, list)
          )
          .tap(deleted =>
            UIO.when(deleted)(
              subscribers.get.flatMap(
                // add item to all subscribers
                UIO.foreach(_)(queue =>
                  queue
                    .offer(name)
                    .catchSomeCause {
                      case cause if cause.interrupted =>
                        subscribers.update(_.filterNot(_ == queue)).as(false)
                    } // if queue was shutdown, remove from subscribers
                )
              )
            )
          )

      def deletedEvents: ZStream[Any, Nothing, String] = ZStream.unwrap {
        for {
          queue <- Queue.unbounded[String]
          _     <- subscribers.update(queue :: _)
        } yield ZStream.fromQueue(queue).ensuring(queue.shutdown)
      }
    }
  }
} 

package neotypes.zio.stream

import zio.Task
import zio.stream.ZStream

trait ZioStreams {
  implicit final val zioStream: neotypes.Stream.Aux[ZioStream, Task] =
    new neotypes.Stream[ZioStream] {
      override final type F[T] = Task[T]

      override final def init[T](value: () => Task[Option[T]]): ZioStream[T] =
        ZStream.unfoldM(()) { _: Unit =>
          value().map { optional =>
            optional.map { v =>
              (v, ())
            }
          }
        }

      override final def onComplete[T](s: ZioStream[T])(f: => Task[Unit]): ZioStream[T] =
        ZStream.bracket(Task(s)) { _ =>
          f.orDie
        }.flatten

      override final def fToS[T](f: Task[ZioStream[T]]): ZioStream[T] =
        ZStream.fromEffect(f).flatten
    }
} 

package polynote.server

import cats.instances.list._
import cats.syntax.traverse._
import fs2.concurrent.Topic
import polynote.buildinfo.BuildInfo
import polynote.kernel.util.Publish
import polynote.kernel.{BaseEnv, StreamThrowableOps}
import polynote.kernel.environment.{Env, PublishMessage, Config}
import polynote.kernel.interpreter.Interpreter
import polynote.kernel.logging.Logging
import polynote.messages._
import polynote.server.auth.IdentityProvider.checkPermission
import polynote.server.auth.{IdentityProvider, Permission, UserIdentity}
import uzhttp.websocket.Frame
import zio.stream.ZStream
import zio.stream.{Stream, Take}
import zio.Queue
import zio.{Promise, RIO, Task, URIO, ZIO}

import scala.collection.immutable.SortedMap

object SocketSession {

  def apply(in: Stream[Throwable, Frame], broadcastAll: Topic[Task, Option[Message]]): URIO[SessionEnv with NotebookManager, Stream[Throwable, Frame]] =
    for {
      output          <- Queue.unbounded[Take[Nothing, Message]]
      publishMessage  <- Env.add[SessionEnv with NotebookManager](Publish(output): Publish[Task, Message])
      env             <- ZIO.environment[SessionEnv with NotebookManager with PublishMessage]
      closed          <- Promise.make[Throwable, Unit]
      _               <- broadcastAll.subscribe(32).unNone.interruptAndIgnoreWhen(closed).through(publishMessage.publish).compile.drain.forkDaemon
      close            = closeQueueIf(closed, output)
    } yield parallelStreams(
        toFrames(ZStream.fromEffect(handshake) ++ Stream.fromQueue(output).unTake),
        in.handleMessages(close)(handler andThen errorHandler) ++ closeStream(closed, output),
        keepaliveStream(closed)).provide(env).catchAllCause {
      cause =>
        ZStream.empty
    }

  private val handler: Message => RIO[SessionEnv with PublishMessage with NotebookManager, Option[Message]] = {
    case ListNotebooks(_) =>
      NotebookManager.list().map {
        notebooks => Some(ListNotebooks(notebooks.map(ShortString.apply)))
      }

    case CreateNotebook(path, maybeContent) =>
      NotebookManager.assertValidPath(path) *>
        checkPermission(Permission.CreateNotebook(path)) *> NotebookManager.create(path, maybeContent).as(None)

    case RenameNotebook(path, newPath) =>
      (NotebookManager.assertValidPath(path) &> NotebookManager.assertValidPath(newPath)) *>
        checkPermission(Permission.CreateNotebook(newPath)) *> checkPermission(Permission.DeleteNotebook(path)) *>
        NotebookManager.rename(path, newPath).as(None)

    case CopyNotebook(path, newPath) =>
      (NotebookManager.assertValidPath(path) &> NotebookManager.assertValidPath(newPath)) *>
        checkPermission(Permission.CreateNotebook(newPath)) *>
        NotebookManager.copy(path, newPath).as(None)

    case DeleteNotebook(path) =>
      NotebookManager.assertValidPath(path) *>
        checkPermission(Permission.DeleteNotebook(path)) *> NotebookManager.delete(path).as(None)

    case RunningKernels(_) => for {
      paths          <- NotebookManager.listRunning()
      statuses       <- ZIO.collectAllPar(paths.map(NotebookManager.status))
      kernelStatuses  = paths.zip(statuses).map { case (p, s) => ShortString(p) -> s }
    } yield Some(RunningKernels(kernelStatuses))

    case other =>
      ZIO.succeed(None)
  }

  val errorHandler: RIO[SessionEnv with PublishMessage with NotebookManager, Option[Message]] => RIO[SessionEnv with PublishMessage with NotebookManager, Option[Message]] =
    _.catchAll {
      err => Logging.error(err).as(Some(Error(0, err)))
    }

  def handshake: RIO[SessionEnv, ServerHandshake] =
    for {
      factories <- Interpreter.Factories.access
      identity  <- UserIdentity.access
      config    <- Config.access
    } yield ServerHandshake(
      (SortedMap.empty[String, String] ++ factories.mapValues(_.head.languageName)).asInstanceOf[TinyMap[TinyString, TinyString]],
      serverVersion = BuildInfo.version,
      serverCommit = BuildInfo.commit,
      identity = identity.map(i => Identity(i.name, i.avatar.map(ShortString))),
      sparkTemplates = config.spark.flatMap(_.propertySets).getOrElse(Nil)
    )
} 

package sttp.client.httpclient.zio

import sttp.client._
import sttp.client.httpclient.WebSocketHandler
import sttp.client.impl.zio.RIOMonadAsyncError
import sttp.client.monad.MonadError
import sttp.client.testing.ConvertToFuture
import sttp.client.testing.HttpTest.wsEndpoint
import sttp.client.testing.websocket.HighLevelWebsocketTest
import sttp.client.ws.WebSocket
import zio.blocking.Blocking
import zio.stream.ZStream
import sttp.client.impl.zio._
import zio.clock.Clock
import zio.{Schedule, ZIO}

import scala.concurrent.duration._
import zio.duration.Duration

class HttpClientHighLevelZioWebsocketTest extends HighLevelWebsocketTest[BlockingTask, WebSocketHandler] {
  implicit val backend: SttpBackend[BlockingTask, ZStream[Blocking, Throwable, Byte], WebSocketHandler] =
    runtime.unsafeRun(HttpClientZioBackend())
  implicit val convertToFuture: ConvertToFuture[BlockingTask] = convertZioBlockingTaskToFuture
  implicit val monad: MonadError[BlockingTask] = new RIOMonadAsyncError

  def createHandler: Option[Int] => BlockingTask[WebSocketHandler[WebSocket[BlockingTask]]] = _ => ZioWebSocketHandler()

  it should "handle backpressure correctly" in {
    new ConvertToFutureDecorator(
      basicRequest
        .get(uri"$wsEndpoint/ws/echo")
        .openWebsocketF(createHandler(None))
        .flatMap { response =>
          val ws = response.result
          send(ws, 1000).flatMap(_ =>
            eventually(10.millis, 500) {
              ws.isOpen.map(_ shouldBe true)
            }
          )
        }
    ).toFuture()
  }

  override def eventually[T](interval: FiniteDuration, attempts: Int)(f: => BlockingTask[T]): BlockingTask[T] = {
    (ZIO.sleep(Duration.fromScala(interval)) *> f.retry(Schedule.recurs(attempts)))
      .provideSomeLayer[Blocking](Clock.live)
  }
} 

package zio.keeper.swim

import zio.ZIO
import zio.keeper.swim.PingPong.{ Ping, Pong }
import zio.keeper.{ ByteCodec, KeeperSpec, NodeAddress }
import zio.stream.ZStream
import zio.test.Assertion._
import zio.test._

object ProtocolSpec extends KeeperSpec {

  val protocolDefinition = Protocol[PingPong].make(
    {
      case Message.Direct(sender, _, Ping(i)) =>
        Message.direct(sender, Pong(i))
      case _ => Message.noResponse
    },
    ZStream.empty
  )

  val testNode = NodeAddress(Array(1, 2, 3, 4), 123)

  val spec = suite("protocol spec")(
    testM("request response") {
      for {
        protocol <- protocolDefinition
        response <- protocol.onMessage(Message.Direct(testNode, 1, Ping(123)))
      } yield assert(response)(equalTo(Message.Direct(testNode, 1, Pong(123))))
    },
    testM("binary request response") {
      for {
        protocol       <- protocolDefinition.map(_.binary)
        binaryMessage  <- ByteCodec.encode[PingPong](Ping(123))
        responseBinary <- protocol.onMessage(Message.Direct(testNode, 1, binaryMessage))
        response <- responseBinary match {
                     case Message.Direct(addr, conversationId, chunk) =>
                       ByteCodec.decode[PingPong](chunk).map(pp => Message.Direct(addr, conversationId, pp))
                     case _ => ZIO.succeed(Message.NoResponse)
                   }
      } yield assert(response)(equalTo(Message.Direct[PingPong](testNode, 1, Pong(123))))
    }
  ).provideLayer(ConversationId.live)

} 

package zio.keeper.swim

import izumi.reflect.Tags.Tag
import zio._
import zio.clock.Clock
import zio.logging.Logging
import zio.stream.{ Sink, ZStream }

object ProtocolRecorder {
  type ProtocolRecorder[A] = Has[ProtocolRecorder.Service[A]]

  trait Service[A] {
    def withBehavior(pf: PartialFunction[Message.Direct[A], Message[A]]): UIO[Service[A]]
    def collectN[B](n: Long)(pr: PartialFunction[Message[A], B]): UIO[List[B]]
    def send(msg: Message.Direct[A]): IO[zio.keeper.Error, Message[A]]
  }

  def apply[A: Tag](
    pf: PartialFunction[Message.Direct[A], Message[A]] = PartialFunction.empty
  ): ZIO[ProtocolRecorder[A], Nothing, Service[A]] =
    ZIO.accessM[ProtocolRecorder[A]](recorder => recorder.get.withBehavior(pf))

  def make[R, E, A: Tag](
    protocolFactory: ZIO[R, E, Protocol[A]]
  ): ZLayer[Clock with Logging with Nodes with R, E, ProtocolRecorder[A]] =
    ZLayer.fromEffect {
      for {
        behaviorRef  <- Ref.make[PartialFunction[Message.Direct[A], Message[A]]](PartialFunction.empty)
        protocol     <- protocolFactory
        messageQueue <- ZQueue.bounded[Message[A]](100)
        _            <- protocol.produceMessages.foreach(consumeMessages(messageQueue, _, behaviorRef, protocol)).fork
        stream       = ZStream.fromQueue(messageQueue)
      } yield new Service[A] {

        override def withBehavior(pf: PartialFunction[Message.Direct[A], Message[A]]): UIO[Service[A]] =
          behaviorRef.set(pf).as(this)

        override def collectN[B](n: Long)(pf: PartialFunction[Message[A], B]): UIO[List[B]] =
          stream.collect(pf).run(Sink.collectAllN[B](n))

        override def send(msg: Message.Direct[A]): IO[zio.keeper.Error, Message[A]] =
          protocol.onMessage(msg)
      }
    }

  private def consumeMessages[A](
    messageQueue: zio.Queue[Message[A]],
    message: Message[A],
    behaviorRef: Ref[PartialFunction[Message.Direct[A], Message[A]]],
    protocol: Protocol[A]
  ): ZIO[Clock with Logging with Nodes, zio.keeper.Error, Unit] =
    message match {
      case Message.WithTimeout(message, action, timeout) =>
        consumeMessages(messageQueue, message, behaviorRef, protocol).unit *>
          action.delay(timeout).flatMap(consumeMessages(messageQueue, _, behaviorRef, protocol)).unit
      case md: Message.Direct[A] =>
        messageQueue.offer(md) *>
          behaviorRef.get.flatMap { fn =>
            ZIO.whenCase(fn.lift(md)) {
              case Some(d: Message.Direct[A]) => protocol.onMessage(d)
            }
          }
      case msg =>
        messageQueue.offer(msg).unit
    }
} 

package zio.keeper.swim

import zio._
import zio.keeper.{ ByteCodec, KeeperSpec, NodeAddress }
import zio.keeper.swim.Messages.WithPiggyback
import zio.keeper.swim.PingPong.{ Ping, Pong }
import zio.logging.Logging
import zio.stream.ZStream
import zio.test.Assertion._
import zio.test._

object MessagesSpec extends KeeperSpec {

  val logger = Logging.console((_, line) => line)

  val messages = for {
    local     <- NodeAddress.local(1111).toManaged_
    transport <- TestTransport.make
    broadcast <- Broadcast.make(64000, 2).toManaged_
    messages  <- Messages.make(local, broadcast, transport)
  } yield (transport, messages)

  val spec = suite("messages")(
    testM("receiveMessage") {
      val testNodeAddress = NodeAddress(Array(1, 2, 3, 4), 1111)

      val protocol = Protocol[PingPong].make(
        {
          case Message.Direct(sender, _, Ping(i)) =>
            Message.direct(sender, Pong(i))
          case _ =>
            Message.noResponse
        },
        ZStream.empty
      )

      messages.use {
        case (testTransport, messages) =>
          for {
            dl    <- protocol
            _     <- messages.process(dl.binary)
            ping1 <- ByteCodec[PingPong].toChunk(PingPong.Ping(123))
            ping2 <- ByteCodec[PingPong].toChunk(PingPong.Ping(321))
            _     <- testTransport.incommingMessage(WithPiggyback(testNodeAddress, 1, ping1, List.empty))
            _     <- testTransport.incommingMessage(WithPiggyback(testNodeAddress, 2, ping2, List.empty))
            messages <- testTransport.outgoingMessages
                         .mapM {
                           case WithPiggyback(_, _, chunk, _) =>
                             ByteCodec.decode[PingPong](chunk)
                         }
                         .take(2)
                         .runCollect
          } yield assert(messages)(hasSameElements(List(PingPong.Pong(123), PingPong.Pong(321))))
      }
    },
    testM("should not exceed size of message") {
      val testNodeAddress = NodeAddress(Array(1, 2, 3, 4), 1111)

      val protocol = Protocol[PingPong].make(
        {
          case Message.Direct(sender, _, Ping(i)) =>
            ZIO.succeed(
              Message.Batch(
                Message.Direct(sender, 1, Pong(i)),
                Message.Broadcast(Pong(i)),
                List.fill(2000)(Message.Broadcast(Ping(1))): _*
              )
            )
          case _ =>
            Message.noResponse
        },
        ZStream.empty
      )

      messages.use {
        case (testTransport, messages) =>
          for {
            dl    <- protocol
            _     <- messages.process(dl.binary)
            ping1 <- ByteCodec[PingPong].toChunk(PingPong.Ping(123))
            _     <- testTransport.incommingMessage(WithPiggyback(testNodeAddress, 1, ping1, List.empty))
            m :: Nil <- testTransport.outgoingMessages
                         .take(1)
                         .runCollect
            bytes <- ByteCodec[WithPiggyback].toChunk(m)
          } yield assert(m.gossip.size)(equalTo(1453)) && assert(bytes.size)(equalTo(62580))
      }
    }
  ).provideCustomLayer(logger ++ ConversationId.live)
} 

package zio.keeper

import izumi.reflect.Tags.Tag
import zio._
import zio.clock.Clock
import zio.keeper.swim.{ Message, Nodes, Protocol }
import zio.logging.Logging
import zio.stream.{ Sink, ZStream }

object ProtocolRecorder {
  type ProtocolRecorder[A] = Has[ProtocolRecorder.Service[A]]

  trait Service[A] {
    def withBehavior(pf: PartialFunction[Message.Direct[A], Message[A]]): UIO[Service[A]]
    def collectN[B](n: Long)(pr: PartialFunction[Message[A], B]): UIO[List[B]]
    def send(msg: Message.Direct[A]): IO[zio.keeper.Error, Message[A]]
  }

  def apply[A: Tag](
    pf: PartialFunction[Message.Direct[A], Message[A]] = PartialFunction.empty
  ): ZIO[ProtocolRecorder[A], Nothing, Service[A]] =
    ZIO.accessM[ProtocolRecorder[A]](recorder => recorder.get.withBehavior(pf))

  def make[R, E, A: Tag](
    protocolFactory: ZIO[R, E, Protocol[A]]
  ): ZLayer[Clock with Logging with Nodes with R, E, ProtocolRecorder[A]] =
    ZLayer.fromEffect {
      for {
        behaviorRef  <- Ref.make[PartialFunction[Message.Direct[A], Message[A]]](PartialFunction.empty)
        protocol     <- protocolFactory
        messageQueue <- ZQueue.bounded[Message[A]](100)
        _            <- protocol.produceMessages.foreach(consumeMessages(messageQueue, _, behaviorRef, protocol)).fork
        stream       = ZStream.fromQueue(messageQueue)
      } yield new Service[A] {

        override def withBehavior(pf: PartialFunction[Message.Direct[A], Message[A]]): UIO[Service[A]] =
          behaviorRef.set(pf).as(this)

        override def collectN[B](n: Long)(pf: PartialFunction[Message[A], B]): UIO[List[B]] =
          stream.collect(pf).run(Sink.collectAllN[B](n))

        override def send(msg: Message.Direct[A]): IO[zio.keeper.Error, Message[A]] =
          protocol.onMessage(msg)
      }
    }

  private def consumeMessages[A](
    messageQueue: zio.Queue[Message[A]],
    message: Message[A],
    behaviorRef: Ref[PartialFunction[Message.Direct[A], Message[A]]],
    protocol: Protocol[A]
  ): ZIO[Clock with Logging with Nodes, zio.keeper.Error, Unit] =
    message match {
      case Message.WithTimeout(message, action, timeout) =>
        consumeMessages(messageQueue, message, behaviorRef, protocol).unit *>
          action.delay(timeout).flatMap(consumeMessages(messageQueue, _, behaviorRef, protocol)).fork.unit
      case md: Message.Direct[A] =>
        messageQueue.offer(md) *>
          behaviorRef.get.flatMap { fn =>
            ZIO.whenCase(fn.lift(md)) {
              case Some(d: Message.Direct[A]) => protocol.onMessage(d)
            }
          }
      case msg =>
        messageQueue.offer(msg).unit
    }
} 

package zio.keeper.hyparview

import zio.clock.Clock
import zio.keeper.transport.Transport
import zio.keeper.{ Error, NodeAddress, SendError, TransportError }
import zio.logging.Logging
import zio.logging.log
import zio.stream.{ Stream, Take, ZStream }
import zio._
import zio.keeper.hyparview.ActiveProtocol._
import zio.duration._

object HyParView {

  def live[R <: Transport with TRandom with Logging with Clock with HyParViewConfig](
    localAddr: NodeAddress,
    seedNodes: List[NodeAddress],
    shuffleSchedule: Schedule[R, ViewState, Any]
  ): ZLayer[R, Error, PeerService] = {
    type R1 = R with Views
    val layer = ZLayer.identity[R] ++ Views.fromConfig(localAddr)
    layer >>> ZLayer.fromManaged {
      for {
        env <- ZManaged.environment[R1]
        cfg <- getConfig.toManaged_
        _ <- log
              .info(s"Starting HyParView on $localAddr with configuration:\n${cfg.prettyPrint}")
              .toManaged(_ => log.info("Shut down HyParView"))
        scope <- ZManaged.scope
        connections <- Queue
                        .bounded[
                          (NodeAddress, Chunk[Byte] => IO[TransportError, Unit], Stream[Error, Chunk[Byte]], UIO[_])
                        ](
                          cfg.connectionBuffer
                        )
                        .toManaged(_.shutdown)
        plumTreeMessages <- Queue
                             .sliding[Take[Error, (NodeAddress, PlumTreeProtocol)]](cfg.userMessagesBuffer)
                             .toManaged(_.shutdown)
        peerEvents <- Queue.sliding[PeerEvent](128).toManaged(_.shutdown)
        sendInitial0 = (to: NodeAddress, msg: InitialProtocol.InitialMessage) =>
          sendInitial(to, msg, scope, connections).provide(env)
        _ <- receiveInitialProtocol[R1, Error](Transport.bind(localAddr), cfg.concurrentIncomingConnections)
              .merge(ZStream.fromQueue(connections))
              .merge(neighborProtocol.scheduleElements(Schedule.spaced(2.seconds)))
              .flatMapParSwitch(cfg.activeViewCapacity) {
                case (addr, send, receive, release) =>
                  ZStream
                    .fromEffect(peerEvents.offer(PeerEvent.NeighborUp(addr)))
                    .ensuring(peerEvents.offer(PeerEvent.NeighborDown(addr)))
                    .flatMap(_ => runActiveProtocol[R1, Error](addr, send, sendInitial0)(receive).ensuring(release))
              }
              .into(plumTreeMessages)
              .toManaged_
              .fork
        _ <- periodic.doShuffle
              .repeat(shuffleSchedule)
              .toManaged_
              .fork
        _ <- periodic.doReport
              .repeat(Schedule.spaced(2.seconds))
              .toManaged_
              .fork
        _ <- ZIO.foreach_(seedNodes)(sendInitial0(_, InitialProtocol.Join(localAddr))).toManaged_
      } yield new PeerService.Service {
        override val identity: UIO[NodeAddress] =
          ZIO.succeed(localAddr)

        override val getPeers: UIO[Set[NodeAddress]] =
          Views.activeView.commit.provide(env)

        override def send(to: NodeAddress, message: PlumTreeProtocol): IO[SendError, Unit] =
          Views.send(to, message).provide(env)

        override val receive: ZStream[Any, Error, (NodeAddress, PlumTreeProtocol)] =
          ZStream.fromQueue(plumTreeMessages).unTake

        override val events: ZStream[Any, Nothing, PeerEvent] =
          ZStream.fromQueue(peerEvents)
      }
    }
  }
} 

package zio.keeper.swim.protocols

import upickle.default._
import zio.{ ZIO, keeper }
import zio.keeper.discovery._
import zio.keeper.{ ByteCodec, NodeAddress }
import zio.keeper.swim.Nodes.{ NodeState, _ }
import zio.keeper.swim.{ ConversationId, Message, Nodes, Protocol }
import zio.logging._
import zio.stream.ZStream

sealed trait Initial

object Initial {

  final case class Join(nodeAddress: NodeAddress) extends Initial
  case object Accept                              extends Initial
  final case class Reject(msg: String)            extends Initial

  implicit val joinCodec: ByteCodec[Join] =
    ByteCodec.fromReadWriter(macroRW[Join])

  implicit val acceptCodec: ByteCodec[Accept.type] =
    ByteCodec.fromReadWriter(macroRW[Accept.type])

  implicit val rejectCodec: ByteCodec[Reject] =
    ByteCodec.fromReadWriter(macroRW[Reject])

  implicit val byteCodec: ByteCodec[Initial] =
    ByteCodec.tagged[Initial][
      Join,
      Accept.type,
      Reject
    ]

  type Env = ConversationId with Nodes with Logging with Discovery

  def protocol(local: NodeAddress): ZIO[Env, keeper.Error, Protocol[Initial]] =
    Protocol[Initial].make(
      {
        case Message.Direct(_, _, Join(addr)) if addr == local =>
          Message.noResponse
        case Message.Direct(_, _, join @ Join(addr)) =>
          nodeState(addr)
            .as(Message.NoResponse)
            .orElse(
              addNode(addr) *>
                changeNodeState(addr, NodeState.Healthy) *>
                Message.direct(addr, Accept).map(accept => Message.Batch[Initial](accept, Message.Broadcast(join)))
            )
        case Message.Direct(sender, _, Accept) =>
          addNode(sender) *>
            changeNodeState(sender, NodeState.Healthy) *>
            Message.noResponse
        case Message.Direct(sender, _, Reject(msg)) =>
          log.error("Rejected from cluster: " + msg) *>
            disconnect(sender) *>
            Message.noResponse
      },
      ZStream
        .fromIterator(
          discoverNodes
            .tap(otherNodes => log.info("Discovered other nodes: " + otherNodes))
            .map(_.iterator)
        )
        .mapM(
          node =>
            NodeAddress
              .fromSocketAddress(node)
              .flatMap(
                nodeAddress => Message.direct(nodeAddress, Join(local))
              )
        )
    )

} 

package zio.kafka.consumer.internal

import zio._
import zio.stream.ZStream


private[internal] class RequestBuffer(ref: Ref[RequestBuffer.State]) {
  import RequestBuffer.State._

  def offer(request: Runloop.Request): UIO[Any] =
    ref.modify {
      case Shutdown      => ZIO.interrupt      -> Shutdown
      case Empty(notify) => notify.succeed(()) -> Filled(List(request))
      case Filled(items) => UIO.unit           -> Filled(request :: items)
    }.flatten

  def takeAll: UIO[List[Runloop.Request]] =
    Promise
      .make[Nothing, Unit]
      .flatMap { p =>
        ref.modify {
          case Shutdown          => ZIO.interrupt -> Shutdown
          case s @ Empty(notify) => (notify.await *> takeAll, s)
          case Filled(items)     => (UIO.succeed(items.reverse), Empty(p))
        }
      }
      .flatten

  def stream: ZStream[Any, Nothing, List[Runloop.Request]] =
    ZStream.repeatEffectOption {
      takeAll.catchAllCause(cause => if (cause.interrupted) ZIO.fail(None) else ZIO.halt(cause))
    }

  def shutdown: UIO[Any] =
    ref.set(Shutdown)
}

private[internal] object RequestBuffer {
  sealed trait State
  object State {
    case object Shutdown                                 extends State
    case class Empty(notifyFill: Promise[Nothing, Unit]) extends State
    case class Filled(items: List[Runloop.Request])      extends State
  }

  def make: UIO[RequestBuffer] =
    Promise.make[Nothing, Unit].flatMap(p => Ref.make(State.Empty(p): State).map(new RequestBuffer(_)))
} 

package zio.kafka.consumer

import org.apache.kafka.common.TopicPartition
import zio.RIO
import zio.blocking.Blocking
import zio.clock.Clock
import zio.stream.ZStream
import zio.kafka.serde.Deserializer

class SubscribedConsumer(
  private val underlying: RIO[Blocking, Consumer.Service]
) {

  def partitionedStream[R, K, V](keyDeserializer: Deserializer[R, K], valueDeserializer: Deserializer[R, V]): ZStream[
    Clock with Blocking,
    Throwable,
    (TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])
  ] =
    ZStream.fromEffect(underlying).flatMap(_.partitionedStream(keyDeserializer, valueDeserializer))

  def plainStream[R, K, V](
    keyDeserializer: Deserializer[R, K],
    valueDeserializer: Deserializer[R, V]
  ): ZStream[R with Clock with Blocking, Throwable, CommittableRecord[K, V]] =
    partitionedStream(keyDeserializer, valueDeserializer).flatMapPar(n = Int.MaxValue)(_._2)
}

class SubscribedConsumerFromEnvironment(
  private val underlying: RIO[Blocking with Consumer, Consumer.Service]
) {

  def partitionedStream[R, K, V](keyDeserializer: Deserializer[R, K], valueDeserializer: Deserializer[R, V]): ZStream[
    Clock with Blocking with Consumer,
    Throwable,
    (TopicPartition, ZStream[R, Throwable, CommittableRecord[K, V]])
  ] =
    ZStream.fromEffect(underlying).flatMap(_.partitionedStream(keyDeserializer, valueDeserializer))

  def plainStream[R, K, V](
    keyDeserializer: Deserializer[R, K],
    valueDeserializer: Deserializer[R, V]
  ): ZStream[R with Clock with Blocking with Consumer, Throwable, CommittableRecord[K, V]] =
    partitionedStream(keyDeserializer, valueDeserializer).flatMapPar(n = Int.MaxValue)(_._2)
} 

package zio.test

import zio.stream.ZStream
import zio.test.Assertion._
import zio.{ UIO, ZIO }

object SampleSpec extends ZIOBaseSpec {

  def spec = suite("SampleSpec")(
    testM("monad left identity") {
      val sample = Sample.shrinkIntegral(0)(5)
      val result = equalSamples(sample.flatMap(Sample.noShrink), sample)
      assertM(result)(isTrue)
    },
    testM("monad right identity") {
      val n                           = 5
      def f(n: Int): Sample[Any, Int] = Sample.shrinkIntegral(0)(n)
      val result                      = equalSamples(Sample.noShrink(n).flatMap(f), f(n))
      assertM(result)(isTrue)
    },
    testM("monad associativity") {
      val sample                      = Sample.shrinkIntegral(0)(2)
      def f(n: Int): Sample[Any, Int] = Sample.shrinkIntegral(0)(n + 3)
      def g(n: Int): Sample[Any, Int] = Sample.shrinkIntegral(0)(n + 5)
      val result                      = equalSamples(sample.flatMap(f).flatMap(g), sample.flatMap(a => f(a).flatMap(g)))
      assertM(result)(isTrue)
    },
    testM("traverse fusion") {
      val sample              = Sample.shrinkIntegral(0)(5)
      def f(n: Int): UIO[Int] = ZIO.succeed(n + 2)
      def g(n: Int): UIO[Int] = ZIO.succeed(n * 3)
      val result = equalEffects(
        sample.foreach(a => f(a).flatMap(g)),
        sample.foreach(f).flatMap(_.foreach(g))
      )
      assertM(result)(isTrue)
    }
  )

  def equalEffects[A, B](
    left: ZIO[Any, Nothing, Sample[Any, A]],
    right: ZIO[Any, Nothing, Sample[Any, B]]
  ): UIO[Boolean] =
    left.flatMap(a => right.flatMap(b => equalSamples(a, b)))

  def equalSamples[A, B](left: Sample[Any, A], right: Sample[Any, B]): UIO[Boolean] =
    if (left.value != right.value) UIO.succeed(false) else equalShrinks(left.shrink, right.shrink)

  def equalShrinks[A, B](
    left: ZStream[Any, Nothing, Sample[Any, A]],
    right: ZStream[Any, Nothing, Sample[Any, B]]
  ): UIO[Boolean] =
    left.zip(right).mapM { case (a, b) => equalSamples(a, b) }.fold(true)(_ && _)
} 

package zio.test.mock

import zio.Chunk
import zio.stream.{ ZSink, ZStream }
import zio.test.mock.module.{ StreamModule, StreamModuleMock }
import zio.test.{ suite, Assertion, TestAspect, ZIOBaseSpec }

object BasicStreamMockSpec extends ZIOBaseSpec with MockSpecUtils[StreamModule] {

  import Assertion._
  import Expectation._
  import TestAspect._

  val A = ZStream.fromIterable(List(1, 2, 3))

  def spec =
    suite("BasicStreamMockSpec")(
      suite("capabilities")(
        suite("sink")(
          testValue("success")(
            StreamModuleMock.Sink(equalTo(1), value(ZSink.collectAll.map(_.toList))),
            StreamModule.sink(1).flatMap(A.run(_)),
            equalTo(List(1, 2, 3))
          ),
          testError("failure")(
            StreamModuleMock.Sink(equalTo(1), failure("foo")),
            StreamModule.sink(1).flatMap(A.run(_)),
            equalTo("foo")
          )
        ),
        suite("stream")(
          testValue("success")(
            StreamModuleMock.Stream(equalTo(1), value(A)),
            StreamModule.stream(1).flatMap(_.runCollect),
            equalTo(Chunk(1, 2, 3))
          )
        )
      )
    ) @@ exceptJS
} 

package zio.test.mock

import zio.internal.Executor
import zio.stream.{ ZSink, ZStream }
import zio.test.TestPlatform
import zio.{ Has, Runtime, Tag, URIO, URLayer, ZIO }


  protected def withRuntime[R]: URIO[R, Runtime[R]] =
    ZIO.runtime[R].map { runtime =>
      if (!TestPlatform.isJS) runtime
      else
        runtime.withExecutor {
          val ec = runtime.platform.executor.asEC
          Executor.fromExecutionContext(Int.MaxValue)(ec)
        }
    }

  abstract class Effect[I: Tag, E: Tag, A: Tag]               extends Capability[R, I, E, A](self)
  abstract class Method[I: Tag, E <: Throwable: Tag, A: Tag]  extends Capability[R, I, E, A](self)
  abstract class Sink[I: Tag, E: Tag, A: Tag, L: Tag, B: Tag] extends Capability[R, I, E, ZSink[Any, E, A, L, B]](self)
  abstract class Stream[I: Tag, E: Tag, A: Tag]               extends Capability[R, I, Nothing, ZStream[Any, E, A]](self)

  object Poly {

    object Effect {
      abstract class Input[E: Tag, A: Tag]  extends Capability.Poly.Input[R, E, A](self)
      abstract class Error[I: Tag, A: Tag]  extends Capability.Poly.Error[R, I, A, Any](self)
      abstract class Output[I: Tag, E: Tag] extends Capability.Poly.Output[R, I, E, Any](self)
      abstract class InputError[A: Tag]     extends Capability.Poly.InputError[R, A, Any](self)
      abstract class InputOutput[E: Tag]    extends Capability.Poly.InputOutput[R, E, Any](self)
      abstract class ErrorOutput[I: Tag]    extends Capability.Poly.ErrorOutput[R, I, Any, Any](self)
      abstract class InputErrorOutput       extends Capability.Poly.InputErrorOutput[R, Any, Any](self)
    }

    object Method {
      abstract class Input[E <: Throwable: Tag, A: Tag]  extends Capability.Poly.Input[R, E, A](self)
      abstract class Error[I: Tag, A: Tag]               extends Capability.Poly.Error[R, I, A, Throwable](self)
      abstract class Output[I: Tag, E <: Throwable: Tag] extends Capability.Poly.Output[R, I, E, Any](self)
      abstract class InputError[A: Tag]                  extends Capability.Poly.InputError[R, A, Throwable](self)
      abstract class InputOutput[E <: Throwable: Tag]    extends Capability.Poly.InputOutput[R, E, Any](self)
      abstract class ErrorOutput[I: Tag]                 extends Capability.Poly.ErrorOutput[R, I, Throwable, Any](self)
      abstract class InputErrorOutput                    extends Capability.Poly.InputErrorOutput[R, Throwable, Any](self)
    }
  }
}

object Mock {

  private[mock] case class Composed[R <: Has[_]: Tag](compose: URLayer[Has[Proxy], R]) extends Mock[R]
} 

package zio.examples.macros

import zio.console.Console
import zio.macros.accessible
import zio.random.Random
import zio.stream.{ZSink, ZStream}
import zio.{Chunk, Has, IO, RIO, UIO, URIO, ZIO, ZLayer, random}

@accessible
object AccessibleMacroExample {

  type AccessibleMacroExample = Has[AccessibleMacroExample.Service]

  trait Foo { val value: String }
  case class Bar(value: String) extends Foo
  case class Wrapped[T](value: T)

  trait Service {

    val foo: UIO[Unit]
    def bar(n: Int): UIO[Unit]
    def baz(x: Int, y: Int): IO[String, Int]
    def poly[A](a: A): IO[Long, A]
    def poly2[A <: Foo](a: Wrapped[A]): IO[String, List[A]]
    def dependent(n: Int): ZIO[Random, Long, Int]
    val value: String
    def function(n: Int): String
    def stream(n: Int): ZStream[Any, String, Int]
    def sink(n: Int): ZSink[Any, Nothing, Int, Nothing, Chunk[Int]]
  }

  val live: ZLayer[Console, Nothing, Has[Service]] =
    ZLayer.fromService(console =>
      new Service {
        val foo: UIO[Unit]                                      = UIO.unit
        def bar(n: Int): UIO[Unit]                              = console.putStrLn(s"bar $n")
        def baz(x: Int, y: Int): IO[String, Int]                = UIO.succeed(x + y)
        def poly[A](a: A): IO[Long, A]                          = UIO.succeed(a)
        def poly2[A <: Foo](a: Wrapped[A]): IO[String, List[A]] = UIO.succeed(List(a.value))
        def dependent(n: Int): ZIO[Random, Long, Int]           = random.nextIntBounded(n)
        val value: String                                       = "foo"
        def function(n: Int): String                            = s"foo $n"
        def stream(n: Int): ZStream[Any, String, Int]           = ZStream.fromIterable(List(1, 2, 3))
        def sink(n: Int): ZSink[Any, Nothing, Int, Nothing, Chunk[Int]]  = ZSink.collectAll
      }
    )

  // can use accessors even in the same compilation unit
  val program: URIO[AccessibleMacroExample with Random, (Int, String, Long, List[Foo], Int, String, String, ZStream[Any, String, Int], ZSink[Any, Nothing, Int, Nothing, Chunk[Int]])] =
    for {
      _  <- AccessibleMacroExample.foo
      _  <- AccessibleMacroExample.bar(1)
      v1 <- AccessibleMacroExample.baz(2, 3).orDieWith(_ => new Exception)
      v2 <- AccessibleMacroExample.poly("foo").orDieWith(_ => new Exception)
      v3 <- AccessibleMacroExample.poly(4L).orDieWith(_ => new Exception)
      v4 <- AccessibleMacroExample.poly2(Wrapped(Bar("bar"))).orDieWith(_ => new Exception)
      v5 <- AccessibleMacroExample.dependent(5).orDieWith(_ => new Exception)
      v6 <- AccessibleMacroExample.value.orDie
      v7 <- AccessibleMacroExample.function(6).orDie
      v8 <- AccessibleMacroExample.stream(7)
      v9 <- AccessibleMacroExample.sink(8)
    } yield (v1, v2, v3, v4, v5, v6, v7, v8, v9)

  // sanity check
  val _foo                            : URIO[AccessibleMacroExample, Unit]                                         = AccessibleMacroExample.foo
  def _bar(n: Int)                    : URIO[AccessibleMacroExample, Unit]                                         = AccessibleMacroExample.bar(n)
  def _baz(x: Int, y: Int)            : ZIO[AccessibleMacroExample, String, Int]                                   = AccessibleMacroExample.baz(x, y)
  def _poly[A](a: A)                  : ZIO[AccessibleMacroExample, Long, A]                                       = AccessibleMacroExample.poly(a)
  def _poly2[A <: Foo](a: Wrapped[A]) : ZIO[AccessibleMacroExample, String, List[A]]                               = AccessibleMacroExample.poly2(a)
  def _dependent(n: Int)              : ZIO[AccessibleMacroExample with Random, Long, Int]                         = AccessibleMacroExample.dependent(n)
  def _value                          : RIO[AccessibleMacroExample, String]                                        = AccessibleMacroExample.value
  def _function(n: Int)               : RIO[AccessibleMacroExample, String]                                        = AccessibleMacroExample.function(n)
  def _stream(n: Int)                 : ZIO[AccessibleMacroExample, Nothing, ZStream[Any, String, Int]]            = AccessibleMacroExample.stream(n)
  def _sink(n: Int)                   : ZIO[AccessibleMacroExample, Nothing, ZSink[Any, Nothing, Int, Nothing, Chunk[Int]]] = AccessibleMacroExample.sink(n)

  // macro autogenerates accessors for `foo`, `bar`, `baz`, `poly`, `poly2`, `value` and `function` below
} 

package zio.rocksdb

import org.{ rocksdb => jrocks }
import zio.stream.ZStream
import zio.{ Has, RIO, Tagged }

abstract class Operations[R <: Has[S], S <: service.RocksDB](implicit tagged: Tagged[S]) {
  private val db: RIO[R, S]                                                       = RIO.access[R](_.get)
  def delete(key: Array[Byte]): RIO[R, Unit]                                      = db >>= (_.delete(key))
  def delete(cfHandle: jrocks.ColumnFamilyHandle, key: Array[Byte]): RIO[R, Unit] = db >>= (_.delete(cfHandle, key))
  def get(key: Array[Byte]): RIO[R, Option[Array[Byte]]]                          = db >>= (_.get(key))
  def get(cfHandle: jrocks.ColumnFamilyHandle, key: Array[Byte]): RIO[R, Option[Array[Byte]]] =
    db >>= (_.get(cfHandle, key))
  def multiGetAsList(keys: List[Array[Byte]]): RIO[R, List[Option[Array[Byte]]]] = db >>= (_.multiGetAsList(keys))
  def multiGetAsList(
    handles: List[jrocks.ColumnFamilyHandle],
    keys: List[Array[Byte]]
  ): RIO[R, List[Option[Array[Byte]]]] =
    db >>= (_.multiGetAsList(handles, keys))
  def newIterator: ZStream[R, Throwable, (Array[Byte], Array[Byte])] = ZStream.unwrap(db map (_.newIterator))
  def newIterator(cfHandle: jrocks.ColumnFamilyHandle): ZStream[R, Throwable, (Array[Byte], Array[Byte])] =
    ZStream.unwrap(db map (_.newIterator(cfHandle)))
  def newIterators(
    cfHandles: List[jrocks.ColumnFamilyHandle]
  ): ZStream[R, Throwable, (jrocks.ColumnFamilyHandle, ZStream[R, Throwable, (Array[Byte], Array[Byte])])] =
    ZStream.unwrap(db map (_.newIterators(cfHandles)))
  def put(key: Array[Byte], value: Array[Byte]): RIO[R, Unit] = db >>= (_.put(key, value))
  def put(cfHandle: jrocks.ColumnFamilyHandle, key: Array[Byte], value: Array[Byte]): RIO[R, Unit] =
    db >>= (_.put(cfHandle, key, value))
} 

package zio.metrics

import zio.{ Fiber, Queue, RIO, Task, UIO, URIO, ZManaged, ZQueue }
import zio.clock.Clock
import zio.stream.ZStream
import zio.duration.Duration.Finite
import zio.metrics.encoders._
import java.util.concurrent.ThreadLocalRandom

class Client(val bufferSize: Long, val timeout: Long, val queueCapacity: Int, host: Option[String], port: Option[Int]) {

  type UDPQueue = ZQueue[Nothing, Any, Encoder, Throwable, Nothing, Metric]

  val queue: UIO[Queue[Metric]] = ZQueue.bounded[Metric](queueCapacity)
  private val duration: Finite  = Finite(timeout)

  val udpClient: ZManaged[Any, Throwable, UDPClient] = (host, port) match {
    case (None, None)       => UDPClient()
    case (Some(h), Some(p)) => UDPClient(h, p)
    case (Some(h), None)    => UDPClient(h, 8125)
    case (None, Some(p))    => UDPClient("localhost", p)
  }

  val sample: List[Metric] => Task[List[Metric]] = metrics =>
    Task(
      metrics.filter(
        m =>
          m match {
            case sm: SampledMetric =>
              if (sm.sampleRate >= 1.0 || ThreadLocalRandom.current.nextDouble <= sm.sampleRate) true else false
            case _ => true
          }
      )
    )

  val udp: List[Metric] => RIO[Encoder, List[Int]] = metrics =>
    for {
      sde  <- RIO.environment[Encoder]
      flt  <- sample(metrics)
      msgs <- RIO.foreach(flt)(sde.get.encode(_))
      ints <- RIO.foreach(msgs.flatten)(s => udpClient.use(_.send(s)))
    } yield ints

  def listen(implicit queue: UDPQueue): URIO[Client.ClientEnv, Fiber[Throwable, Unit]] =
    listen[List, Int](udp)

  def listen[F[_], A](
    f: List[Metric] => RIO[Encoder, F[A]]
  )(implicit queue: UDPQueue): URIO[Client.ClientEnv, Fiber[Throwable, Unit]] =
    ZStream
      .fromQueue[Encoder, Throwable, Metric](queue)
      .groupedWithin(bufferSize, duration)
      .mapM(l => f(l))
      .runDrain
      .fork

  val send: Queue[Metric] => Metric => Task[Unit] = q =>
    metric =>
      for {
        _ <- q.offer(metric)
      } yield ()

  val sendAsync: Queue[Metric] => Metric => Task[Unit] = q =>
    metric =>
      for {
        _ <- q.offer(metric).fork
      } yield ()
}

object Client {

  type ClientEnv = Encoder with Clock //with Console

  def apply(): Client = apply(5, 5000, 100, None, None)

  def apply(bufferSize: Long, timeout: Long): Client =
    apply(bufferSize, timeout, 100, None, None)

  def apply(bufferSize: Long, timeout: Long, queueCapacity: Int): Client =
    apply(bufferSize, timeout, queueCapacity, None, None)

  def apply(bufferSize: Long, timeout: Long, queueCapacity: Int, host: Option[String], port: Option[Int]): Client =
    new Client(bufferSize, timeout, queueCapacity, host, port)

}