MonadLogTask + TreeLogger

gistfile1.scala

import scalaz._, concurrent._

trait MonadLogTask[F[_], L] extends Monad[F] with Catchable[F] {

  def log[A](label: => L, task: Task[A]): F[A]

  def heading[A](label: => L, fa: F[A]): F[A]

  def liftTask[A](task: Task[A]): F[A]

}

object MonadLogTask {
  def apply[F[_], L](implicit F: MonadLogTask[F, L]): MonadLogTask[F, L] = F

  implicit def taskLogTask[L]: MonadLogTask[Task, L] =
    new MonadLogTask[Task, L] {
      def log[A](label: => L, task: Task[A]) = task
      def heading[A](label: => L, task: Task[A]) = task
      def liftTask[A](task: Task[A]) = task
      def point[A](a: => A) = Task.taskInstance.point(a)
      def bind[A, B](fa: Task[A])(f: A => Task[B]) = Task.taskInstance.bind(fa)(f)
      def fail[A](e: Throwable) = Task.taskInstance.fail[A](e)
      def attempt[A](fa: Task[A]) = Task.taskInstance.attempt(fa)
    }
}

/**
 * Externalizable representation of `Throwable`
 *
 * Stolen from doobie
 * @see https://github.com/tpolecat/doobie
 */
final case class Thrown(className: String,
                        message: String,
                        stack: IList[StackTraceElement],
                        cause: Option[Thrown])

object Thrown {

  def fromThrowable(t: Throwable): Thrown =
    Thrown(t.getClass.getName,
      t.getMessage,
      t.getStackTrace.foldRight(IList[StackTraceElement]())(_ :: _),
      Option(t.getCause).map(fromThrowable))

}

sealed trait LogNode[L] {
  import LogNode._

  /** Commits this node if pending, returning this otherwise */
  def commit(endNanos: Long, err: Option[Throwable]): LogNode[L] =
    this match {
      case Pending(l, startNanos) => entryT(l.value, endNanos - startNanos, err)
      case _ => this
    }

}
private case class Root[L](label: Need[L]) extends LogNode[L]
private case class Pending[L](label: Need[L], startNanos: Long) extends LogNode[L]
private case class Entry[L](label: Need[L], ellapsedNanos: Long, error: Option[Thrown]) extends LogNode[L]

object LogNode {
  import scalaz.syntax.show._

  def root[L](l: => L): LogNode[L] =
    Root(Need(l))

  def pending[L](l: => L, startNanos: Long): LogNode[L] =
    Pending(Need(l), startNanos)

  def entry[L](l: => L, ellapsedNanos: Long, thrown: Option[Thrown]): LogNode[L] =
    Entry(Need(l), ellapsedNanos, thrown)

  def entryT[L](l: => L, ellapsedNanos: Long, err: Option[Throwable]): LogNode[L] =
    entry(l, ellapsedNanos, err map Thrown.fromThrowable)

  def success[L](l: => L, ellapsedNanos: Long): LogNode[L] =
    entry(l, ellapsedNanos, None)

  def failure[L](l: => L, ellapsedNanos: Long, err: Throwable): LogNode[L] =
    entryT(l, ellapsedNanos, Some(err))

  implicit def logNodeShow[L: Show]: Show[LogNode[L]] =
    Show.shows {
      case Root(l)               => l.value.shows
      case Pending(l, _)         => s"Pending: ${l.value.shows}"
      case Entry(l, ns, None)    => f"[ok] ${l.value.shows} (${ns / 1000000.0}%.3f ms)"
      case Entry(l, ns, Some(t)) => f"[er] ${l.value.shows} !! ${t.message} !! (${ns / 1000000.0}%.3f ms)"
    }

}

type Logged[L, +A] = Task[(Log[L], Throwable \/ A)]

/**
 * Logging using a structured rose tree log
 *
 * Stolen from doobie
 * @see https://github.com/tpolecat/doobie
 */
sealed trait TreeLogger[L, +A] {
  import TreeLogger._

  protected val tl: TreeL[L]#f[A]

  def map[B](f: A => B): TreeLogger[L, B] =
    TreeLogger(tl map f)

  def flatMap[B](f: A => TreeLogger[L, B]): TreeLogger[L, B] =
    TreeLogger(tl flatMap (f(_).tl))

  /** Add a heading to this logger */
  def heading(label: => L): TreeLogger[L, A] =
    TreeLogger(for {
      p <- nanoTime.map(pending(label, _)).liftM[TreeLST[L]#g]
      _ <- TreeLSM[L].modify(_.insertDownLast(Tree(p)))
      r <- tl.run
      _ <- commit(p, r)
    } yield r)

  def ++>(label: => L): TreeLogger[L, A] =
    heading(label)

  def mapT(f: Task ~> Task): TreeLogger[L, A] =
    TreeLogger(StateT.StateMonadTrans[LZ[L]].hoist(f).apply(tl.run))

  def attempt: TreeLogger[L, Throwable \/ A] =
    TreeLogger(tl.run.map(_.right[Throwable]))

  def run(rootLabel: => L): Logged[L, A] =
    tl.run.leftMap(_.tree).run(Tree(root(rootLabel)).loc)

  def runZero(implicit L: Monoid[L]): Logged[L, A] =
    run(L.zero)

  def runLog(rootLabel: => L, write: String => Task[Unit])(implicit L: Show[L]): Task[A] =
    run(rootLabel) >>= (l => write(l._1.drawTree) *> l._2.fold(Task.fail, Task.now))

  def unLog(implicit L: Monoid[L]): Task[A] =
    runZero >>= (_._2.fold(Task.fail, Task.now))

}

object TreeLogger {
  private def apply[L, A](_tl: TreeL[L]#f[A]): TreeLogger[L, A] =
    new TreeLogger[L, A] {
      protected val tl = _tl
    }

  private def apply[L, A](tls: TreeLS[L]#f[Throwable \/ A]): TreeLogger[L, A] =
    new TreeLogger[L, A] {
      protected val tl = EitherT[TreeLS[L]#f, Throwable, A](tls)
    }

  def log[L, A](label: => L, t: Task[A]): TreeLogger[L, A] =
    TreeLogger(for {
      p <- nanoTime.map(pending(label, _)).liftM[TreeLST[L]#g]
      _ <- TreeLSM[L].modify(_.insertDownLast(Tree(p)))
      r <- t.attempt.liftM[TreeLST[L]#g]
      _ <- commit(p, r)
    } yield r)

  def silent[L, A](t: Task[A]): TreeLogger[L, A] =
    TreeLogger(t.attempt.liftM[TreeLST[L]#g])

  def fail[L](err: Throwable): TreeLogger[L, Nothing] =
    TreeLogger(TreeLSM[L].point(err.left[Nothing]))

  def toTask[L: Show](rootLabel: => L, write: String => Task[Unit]): ({type f[+a]=TreeLogger[L, a]})#f ~> Task =
    new (({type f[+a]=TreeLogger[L, a]})#f ~> Task) {
      def apply[A](tl: TreeLogger[L, A]): Task[A] =
        tl.runLog(rootLabel, write)
    }

  implicit def TreeLoggerMonadLogTask[L]: MonadLogTask[({type f[+a] = TreeLogger[L, a]})#f, L] =
    new MonadLogTask[({type f[+a] = TreeLogger[L, a]})#f, L] {
      def log[A](label: => L, task: Task[A]): TreeLogger[L, A] =
        TreeLogger.log(label, task)

      def heading[A](label: => L, fa: TreeLogger[L, A]): TreeLogger[L, A] =
        fa.heading(label)

      def liftTask[A](task: Task[A]): TreeLogger[L, A] =
        TreeLogger.silent(task)

      def bind[A, B](fa: TreeLogger[L, A])(f: A => TreeLogger[L, B]): TreeLogger[L, B] =
        fa flatMap f

      def point[A](a: => A): TreeLogger[L, A] =
        TreeLogger(TreeLM[L].point(a))

      def fail[A](err: Throwable) =
        TreeLogger.fail[L](err)

      def attempt[A](fa: TreeLogger[L, A]): TreeLogger[L, Throwable \/ A] =
        fa.attempt
    }

  ////

  private type ErrT[F[+_], +A] = EitherT[F, Throwable, A]
  private type LZ[L] = TreeLoc[LogNode[L]]
  private trait TreeLST[L] { type g[f[+_], +a] = StateT[f, LZ[L], a] }
  private trait TreeLS[L] { type f[+a] = TreeLST[L]#g[Task, a] }
  private trait TreeL[L] { type f[+a] = ErrT[TreeLS[L]#f, a] }

  implicit private def TreeLSM[L] = StateT.stateTMonadState[LZ[L], Task]
  implicit private def TreeLM[L] = EitherT.eitherTMonad[TreeLS[L]#f, Throwable]

  private def commit[L, A](n: LogNode[L], result: Throwable \/ A): TreeLS[L]#f[Unit] =
    for {
      e <- nanoTime.map(n.commit(_, result.swap.toOption)).liftM[TreeLST[L]#g]
      _ <- TreeLSM[L].modify(_.setLabel(e).parent.get)
    } yield ()

  private def nanoTime: Task[Long] =
    Task.delay(System.nanoTime)

}