Session Stream

Session.scala

import cats.{Monad, StackSafeMonad}
import cats.effect.{IO, LiftIO}
import io.iteratee.{Enumeratee, Enumerator, Iteratee}
import cats.implicits._

trait Session {
  def close(): Unit
}

case class SessionContext[T](exec: Session => IO[T]) {
  def runSession(newSession: () => Session): IO[T] = {
    val session = newSession()
    exec(session).map { t =>
      session.close()
      t
    }
  }
}

object SessionContext {
  implicit val instance: Monad[SessionContext] with LiftIO[SessionContext] =
    new StackSafeMonad[SessionContext] with LiftIO[SessionContext] {
      override def pure[A](a: A): SessionContext[A] =
        SessionContext(Function.const(IO.pure(a)))

      override def flatMap[A, B](fa: SessionContext[A])
                                (f: A => SessionContext[B]): SessionContext[B] =
        SessionContext(session => fa.exec(session).flatMap(a => f(a).exec(session)))

      override def liftIO[A](io: IO[A]): SessionContext[A] =
        SessionContext(Function.const(io))
    }

  def session: SessionContext[Session] = SessionContext(s => IO.pure(s))
}

def asStream[A](chunkSize: Int = 1)(els: Session => Iterable[A]): Enumerator[SessionContext, A] =
  Enumerator.liftM(SessionContext.session).flatMap(s => Enumerator.enumIterable(els(s), chunkSize))

class CounterSession(var counter: Int = 0) extends Session {
  def inc(): Int = {
    counter += 1
    counter
  }
  def close(): Unit = println(s"closing counter at $counter")
}

val count = asStream() {
  case cs: CounterSession => Iterator.continually(cs.inc()).toIterable
  case _                  => Iterable.empty
}


def sumUntil(n: Int) =
  Iteratee.sum[SessionContext, Int].through(Enumeratee.takeWhile(_ < n))

val action = count.into(sumUntil(50))

action.runSession{ () =>
  println("starting at 0")
  new CounterSession(0)
}.unsafeRunSync()

action.runSession{ () =>
  println("starting at 30")
  new CounterSession(30)
}.unsafeRunSync()