Scala 未来的 flatMap 实现（链接）

Question

Trying to wrap my head around how async tasks are chained together, for ex futures and flatMap

val a: Future[Int] = Future { 123 }
val b: Future[Int] = a.flatMap(a => Future { a + 321 })

How would i implement something similar where i build a new Future by waiting for a first and applying a f on the result. I tried looking a scala's source code but got stuck here: https://github.com/scala/scala/blob/2.13.x/src/library/scala/concurrent/Future.scala#L216

i would imagine some code like:

def myFlatMap(a: Future[Int])(f: a => Future[Int]): Future[Int] = {
   Future {
      // somehow wait for a to complete and the apply 'f'
      a.onComplete {
        case Success(x) => f(a)
      }
   }
}

but i guess above will return a Future[Unit] instead. I would just like to understand the "pattern" of how async tasks are chained together.

Answer 1

(Disclaimer: I'm the main maintainer of Scala Futures)

You wouldn't want to implement flatMap/recoverWith/transformWith yourself—it is a very complex feature to implement safely (stack-safe, memory-safe, and concurrency-safe).

You can see what I am talking about here:

• https://github.com/scala/scala/blob/2.13.x/src/library/scala/concurrent/impl/Promise.scala#L442
• https://github.com/scala/scala/blob/2.13.x/src/library/scala/concurrent/impl/Promise.scala#L307
• https://github.com/scala/scala/blob/2.13.x/src/library/scala/concurrent/impl/Promise.scala#L274

There is more reading on the topic here: https://viktorklang.com/blog/Futures-in-Scala-2.12-part-9.html

Answer 2

I found that looking at the source for scala 2.12 was easier for me to understand how Future / Promise was implemented, i learned that Future is in fact a Promise and by looking at the implementation i now understand how "chaining" async tasks can be implemented.

Im posting the code that i wrote to help me better understand what was going on under the hood.

import java.util.concurrent.{ScheduledThreadPoolExecutor, TimeUnit}
import scala.collection.mutable.ListBuffer

trait MyAsyncTask[T] {
  type Callback = T => Unit
  var result: Option[T]

  // register a new callback to be called when task is completed.
  def onComplete(callback: Callback): Unit
  // complete the given task, and invoke all registered callbacks.
  def complete(value: T): Unit

  // create a new task, that will wait for the current task to complete and apply
  // the provided map function, and complete the new task when completed.
  def flatMap[B](f: T => MyAsyncTask[B]): MyAsyncTask[B] = {
    val wrapper = new DefaultAsyncTask[B]()
    onComplete { x =>
      f(x).onComplete { y =>
        wrapper.complete(y)
      }
    }

    wrapper
  }

  def map[B](f: T => B): MyAsyncTask[B] = flatMap(x => CompletedAsyncTask(f(x)))
}

/**
  * task with fixed pre-calculated result.
  */
case class CompletedAsyncTask[T](value: T) extends MyAsyncTask[T] {
  override var result: Option[T] = Some(value)
  override def onComplete(callback: Callback): Unit = {
    // we already have the result just call the callback.
    callback(value)
  }

  override def complete(value: T): Unit = () // noop nothing to complete.
}

class DefaultAsyncTask[T] extends MyAsyncTask[T] {
  override var result: Option[T] = None
  var isCompleted = false
  var listeners = new ListBuffer[Callback]()

  /**
    * register callback, to be called when task is completed.
    */
  override def onComplete(callback: Callback): Unit = {
    if (isCompleted) {
      // already completed just invoke callback
      callback(result.get)
    } else {
      // add the listener
      listeners.addOne(callback)
    }
  }

  /**
    * trigger all registered callbacks to `onComplete`
    */
  override def complete(value: T): Unit = {
    result = Some(value)
    listeners.foreach { listener =>
      listener(value)
    }
  }
}

object MyAsyncTask {
  def apply[T](body: (T => Unit) => Unit): MyAsyncTask[T] = {
    val task = new DefaultAsyncTask[T]()
    // pass in `complete` as callback.
    body(task.complete)

    task
  }
}

object MyAsyncFlatMap {

  /**
    * helper to simulate async task.
    */
  def delayedCall(body: => Unit, delay: Int): Unit = {
    val scheduler = new ScheduledThreadPoolExecutor(1)
    val run = new Runnable {
      override def run(): Unit = {
        body
      }
    }

    scheduler.schedule(run, delay, TimeUnit.SECONDS)
  }

  def main(args: Array[String]): Unit = {
    val getAge = MyAsyncTask[Int] { cb =>
      delayedCall({ cb(66) }, 1)
    }

    val getName = MyAsyncTask[String] { cb =>
      delayedCall({ cb("John") }, 2)
    }

    // same as: getAge.flatMap(age => getName.map(name => (age, name)))
    val result: MyAsyncTask[(Int, String)] = for {
      age <- getAge
      name <- getName
    } yield (age, name)

    result.onComplete {
      case (age, name) => println(s"hello $name $age")
    }
  }
}