Simple and Safe Pickling of Closures in Scala 3

https://www.youtube.com/watch?v=9Fw0iQ0Rze8

Transcript

1. Simple and Safe Pickling of Closures in Scala 3 Jonas Spenger, Philipp Haller KTH Royal Institute of Technology, Stockholm, Sweden Scala Days 2025, 19 - 21 August, Lausanne, Switzerland 1
2. About The Speakers 2 Philipp Haller @philippkhaller (X) @phaller (GitHub) https://people.kth.se/~phaller/ Jonas Spenger @jspenger (GitHub) https://jspenger.github.io/ • PhD Student at KTH Royal Institute of Technology, Stockholm, Sweden • Durable computation (dataflow streaming, futures, actors) • Associate Professor at KTH Royal Institute of Technology, Stockholm, Sweden • Concurrency (actors, futures/async, joins) • Distributed programming • Type systems (capabilities, uniqueness)
3. Closures and Environments 3 • Note: threshold is not defined within the blue anonymous function! • Thus: threshold is a called a free variable (of the blue anonymous function) • Original definition of the term closure: An anonymous function “whose open bindings (free variables) have been closed by the lexical environment” (Peter J. Landin, 1964)" val numbers = List(6, 3, 9, 2, 4) def below(threshold: Int) = numbers.filter(num => num < threshold) assert(below(5) == List(3, 2, 4)) P. J. Landin: The Mechanical Evaluation of Expressions. Comput. J. 6(4): 308-320 (1964) https://doi.org/10.1093/comjnl/6.4.308 The blue lambda refers to threshold in its lexical context
4. object Global { val fld = 5 } val fun = { (num: Int) => num < Global.fld } assert(numbers.filter(fun), List(3, 2, 4)) What is Captured and What isn't? 4 object Global { val fld = 5 } val fun = { (num: Int) => num < Global.fld } assert(numbers.filter(fun), List(3, 2, 4)) val numbers = List(6, 3, 9, 2, 4) def mkFilter(threshold: Int) = (num: Int) => num < threshold val fun = mkFilter(5) assert(numbers.filter(fun), List(3, 2, 4)) class Filter(threshold: Int) extends Int => Bool: def apply(num: Int) = num < threshold val fun = Filter(5) assert(numbers.filter(fun), List(3, 2, 4)) Top-level object captured "captured" Not captured
5. class SparkExample { val distData = sc.parallelize(Array(1, 2, 3, 4, 5)) def transform(x: Int): Int = x+1 def test(): Unit = { val transformed = distData.map(elem => transform(elem)) transformed.collect().foreach(elem => println(elem)) } } Trouble in Paradise 5 Using Apache Spark™ Exception in thread "main" org.apache.spark.SparkException: Task not serializable ...
6. class SparkExample { val distData = sc.parallelize(Array(1, 2, 3, 4, 5)) def transform(x: Int): Int = x+1 def test(): Unit = { val transformed = distData.map(elem => transform(elem)) transformed.collect().foreach(elem => println(elem)) } } Trouble in Paradise 6 • distData is a distributed data set → each remote worker has a piece of the data • map sends its argument closure to each worker → argument closure must be serialized Uses transform method from enclosing scope
7. class SparkExample { val distData = sc.parallelize(Array(1, 2, 3, 4, 5)) def transform(x: Int): Int = x+1 def test(): Unit = { val transformed = distData.map(elem => this.transform(elem)) transformed.collect().foreach(elem => println(elem)) } } Trouble in Paradise 7 • distData is a distributed data set → each remote worker has a piece of the data • map sends its argument closure to each worker → argument closure must be serialized Actually: closure captures this! Closure is sent to remote workers → this must be serialized The type of this is SparkExample which is not serializable, hence... Exception in thread "main" org.apache.spark.SparkException: Task not serializable ...
8. Problematic Uses of Closures • Using closures in distributed settings is a safety risk ◦ Example: serializing closures can result in runtime errors (e.g., java.io.NotSerializableException on the JVM) • What about concurrency? 😇 8
9. Example: Concurrency 9 val customerData: mutable.Map[Int, CustomerInfo] = ... def averageAge(customers: List[Customer]): Future[Float] = Future { val infos = customers.flatMap { c => customerData.get(c.customerNo) match case Some(info) => List(info) case None => List() } val sumAges = infos.foldLeft(0)(_ + _.age).toFloat if (infos.nonEmpty) sumAges / infos.size else 0.0f } Possible data race! Mutable map
10. Problematic Uses of Closures • Using closures in distributed settings is a safety risk ◦ Example: serializing closures can result in runtime errors (e.g., java.io.NotSerializableException on the JVM) • Using closures in concurrent settings is a safety risk ◦ Example: running a closure on a concurrent thread could cause a data race if a captured variable refers to a shared mutable object • Anything else? ◦ Send a closure from a Scala.js frontend to a Scala/JVM backend? ◦ Requires portable pickling/serialization! 10
11. Short History of Spores • Goals of Spores: ◦ Safety checking of closures ◦ Increased flexibility, e.g., portable serialization • Main idea: “Spore” = a special kind of closure that ◦ consists of a closed (capture-free) function and an explicit environment ◦ where the environment is constrained using type classes • Spores were first proposed in SIP-21 in 2013 and presented at ECOOP'14 12 Miller, Haller, and Odersky. Spores: a type-based foundation for closures in the age of concurrency and distribution. ECOOP 2014 (Google Scholar: 56 citations) https://doi.org/10.1007/978-3-662-44202-9_13
12. Spores3 • A new take on Spores for Scala 3 • Uses a new approach for type-class-based pickling • A simpler and more robust implementation • Talk and paper (open access): 13 Philipp Haller: Enhancing closures in Scala 3 with Spores3. ACM SIGPLAN Scala Symposium, Berlin, Germany, June 2022. https://doi.org/10.1145/3550198.3550428 Talk at Strange Loop 2022: video, slides
13. This Talk • Plenty of innovations in the new version 0.2 • Spores are now even more flexible ◦ Partial application of Spores ◦ Support for context functions ◦ ... computation on pickled data made simple! 14
14. The Spore Factory // Spore without environment val spore: Spore[Int => Int] = Spore { (x: Int) => x + 1 } // Spore with environment val spore: Spore[Int => Int] = Spore(12) { (env: Int) => (x: Int) => x + env } ^^ ^^^ ^^^ env is replaced with 12 in body 15
15. // Spore with context environment val spore: Spore[Int => Int] = Spore.applyWithCtx(12) { (env: Int) ?=> (x: Int) => x + summon[Int]} The Spore Factory 16
16. Spore Auto Capture // Automatically captures environment variables // y, z captured* val y = 12 val z = “World” val spore: Spore[Int => String] = Spore.auto { (x: Int) => (x + y).toString() + z } ^ ^ 17 *if not top-level
17. Unwrapping / Applying a Spore 18 // unwraps the value in the spore val spore: Spore[Int => Int] = Spore { ... } spore.unwrap().apply(12) // or, using imported implicit conversion import spores.default.given spore(12) // Spore[Int => Int] is converted to Int => Int
18. val y = 12 val spore: Spore[Int => Int] = Spore { (x: Int) => x + y } The Spore Factory - Compile-Time Checks 19 More examples here and here
19. val y = 12 val spore: Spore[Int => Int] = Spore { (x: Int) => x + y } ^^^ Error: Invalid capture of variable `y`. The Spore Factory - Compile-Time Checks 20 More examples here and here
20. class Foo: val x = 12 def spore = Spore.apply { () => 42 * x } The Spore Factory - Compile-Time Checks 21 More examples here and here
21. class Foo: val x = 12 def spore = Spore.apply { () => 42 * x } ^^^ (i.e., this.x) Error: Invalid capture of `this` from class Foo. The Spore Factory - Compile-Time Checks 22 More examples here and here
22. // What about this example? class Outer: class Foo def spore = Spore.apply { () => class Bar extends Foo new Bar() } The Spore Factory - Compile-Time Checks 23
23. // What about this example? class Outer: class Foo def spore = Spore.apply { () => class Bar extends Foo ^^^ new Bar() } Error: Invalid capture of variable `Outer.Foo`. The Spore Factory - Compile-Time Checks 24
24. // What about this example? ... class Bar def foo(using Bar): Spore[Bar] = Spore.apply { summon } The Spore Factory - Compile-Time Checks 25
25. // What about this example? ... class Bar def foo(using Bar): Spore[Bar] = Spore.apply { summon } ^^^^^^ Error: Invalid capture of variable `x$1`. The Spore Factory - Compile-Time Checks 26 Captured context parameter!
26. • Problematic capturing is prevented • This ensures that it is always safe to pickle / unpickle a Spore • Pickling example: import upickle.default.* val spore: Spore[Int => Int] = ... val pickled: String = write(spore) // JSON String val unpickled = read[Spore[Int => Int]](pickled) unpickled: Spore[Int => Int] Safe Pickling of Spores 27 ‘write’ to JSON String ‘read’ from JSON String
27. Macro for Checking Captures • Macros.checkBodyExpr(bodyExpr) ◦ Collect all symbols which might be captured ▪ Ident, This, ClassDefs, New, etc. ▪ Ignore some special cases, e.g. ignore U in new T[U]() ◦ Remove symbols owned by the closure (parameters, etc.) ◦ Remove top-level symbols ◦ Remaining symbols are captured 28
28. • Closures are capture free ◦ Checked by macros • Captured environment variables ◦ By partial application Spores3 Approach 29
29. // Apply (partially) a Spore to a value val spore: Spore[Int => Int => String] spore.withEnv(12): Spore[Int => String] ^^ requires a given: Spore[ReadWriter[Int]] // Apply (partially) a Spore to a Spore val spore: Spore[Int => Int => String] = ... val spore12: Spore[Int] = Spore { 12 } spore.withEnv2(spore12): Spore[Int => String] ^^^^^^^ no given needed Partial Application 30
30. // Also works for context parameters val spore2: Spore[Int ?=> Int => String] spore2.withCtx(12): Spore[Int => String] ... spore2.withCtx2(spore12): Spore[Int => String] Partial Application 31
31. Spore Builders 32 // Spore builders work cross-platform: // - Scala JVM, Scala.js, Scala Native // - Also: Spore builders have names (unlike Lambdas) object MyBuilder extends SporeBuilder[Int => Int]({ x => x + 1 }) val spore: Spore[Int => Int] = MyBuilder.build()
32. Creating a `given Spore[ReadWriter[T]]` 33 // Works if the `summon`ed value is top-level given Spore[ReadWriter[Int]] = Spore { summon } // Works for all top-level ReadWriters inline given [T: ReadWriter]: Spore[ReadWriter[T]] = Spore { summon[ReadWriter[T]] }
33. That’s It! 34
34. Digression: Spores and Values 35 • Spores can express functions, function applications, and values • Powerful idea! ◦ Enables leveraging known concepts ▪ type-class derivation ▪ combinators, etc.
35. 36 apply 1 (x: Int) => (y: Int) => x + y apply 3 AST: ((x: Int) => (y: Int) => x + y).apply(1).apply(3) Program: Value Domain Spore / Pickled Domain Spore.apply{(x: Int) => (y: Int) => x + y}.applyWithEnv2(Env(1)).applyWithEnv2(Env(3)) Program: applyWithEnv2 Spore(1) Spore {(x: Int) => (y: Int) => x + y} applyWithEnv2 Spore(3) AST: Spore domain mirrors the Value domain • Computations remain in the Spore domain • Partially applied Spore is safe to pickle!
36. • Internal Spore structure mimics the AST structure ◦ Lambda(function) // The capture-free function ◦ Env(value, sporeRW) // Environments ◦ WithEnv(spore, sporeEnv) // Partial application Spore to env ◦ WithCtx(spore, sporeEnv) // Partial context-application Digression: Spores and Values 37
37. • Use Spores for serializing actor behaviors • Periodically serialize execution state • Recover state after crash • https://github.com/jspenger/durable-actor (WIP) ◦ API inspired by Akka Typed Actors Example: Durable Actors 38
38. • Two actors: ◦ Repeat: Actor A sends Increment() to B ▪ Actor B counts # received Increment() ◦ Actor A sends Retrieve(self) to B ▪ Actor B sends RetrieveReply(count) to A • Runs on JVM, Scala.js, Scala Native • Program can be stopped and recovered on another platform • Link Demo: “Counting” Actor (Cross-Platform) 39
39. // Actor behavior contains Spore case class ReceiveBehavior[T]( sp: Spore[ActorContext[T] ?=> T => ReceiveBehavior[T]] ) Durable Actor Behaviors 40
40. Behavior Factory // Use Spore.auto method by using inline inline def receive[T]( inline fun: ActorContext[T] ?=> T => ReceiveBehavior[T] ) = ReceiveBehavior( Spore.auto(fun) ) 41 User does not pass a Spore!
41. Behavior Factory // Alt. create actor behavior from Spore def receive2[T]( sp: Spore[ActorContext[T] ?=> T => ReceiveBehavior[T]] ) = ReceiveBehavior(sp) 42
42. Preliminary Microbenchmarks • Compare Spores3* to Java** • Pickling performance: ◦ Speedup between 1.3x and 5.4x ◦ Size reduction between 2.5x and 9.2x • More details in the full slide deck (online) 43 *Serialized Spore by upickle writeBinary **Serialized by java.io.ObjectInputStream / java.io.ObjectOutputStream
43. Summary • Spores3: making closures in Scala safer and more flexible ◦ Spore is safe to pickle / unpickle ◦ Captured environment is checked ◦ Version 0.2 is out now! https://github.com/phaller/spores3 44
44. Thank You! • Spores3: making closures in Scala safer and more flexible ◦ Spore is safe to pickle / unpickle ◦ Captured environment is checked ◦ Version 0.2 is out now! https://github.com/phaller/spores3 45 Philipp Haller @philippkhaller (Twitter) @phaller (GitHub) https://people.kth.se/~phaller/ Jonas Spenger @jspenger (GitHub) https://jspenger.github.io/ Simple and Safe Pickling of Closures in Scala 3 Jonas Spenger, Philipp Haller, Scala Days 2025, 19 - 21 August, Lausanne, Switzerland
45. Preliminary Microbenchmarks • Compare Spores3* to Java** • Function1: x => () => x + 1 • Function2: (x => () => x + 1).apply(7) • Function3: (list => () => list.sum).apply(List(1, …, 10)) • Function4: (pred => list => list.filter(pred)).apply(x => x % 2 == 0) 46 *Serialized Spore by upickle writeBinary **Serialized by java.io.ObjectInputStream / java.io.ObjectOutputStream
46. Operations per Second • Measure time to serialize and deserialize with JMH ◦ MOPS (million operations per second) • Function1: Spores3: 0.70 Java: 0.13 (MOPS) • Function2: Spores3: 0.16 Java: 0.12 • Function3: Spores3: 0.11 Java: 0.050 • Function4: Spores3: 0.22 Java: 0.10 • Speedup between 1.3x and 5.4x 47
47. • Serialize to Array[Byte] ◦ Measure size in Byte • Function1: Spores3: 67 Java: 622 (Byte) • Function2: Spores3: 223 Java: 562 • Function3: Spores3: 359 Java: 1166 • Function4: Spores3: 186 Java: 862 • Size reduction between 2.5x and 9.2x Serialization Size 48
48. // `events` and `behaviors` are serializable val events = mutable.Queue[ActorSend] val behaviors = mutable.Map[ActorRef[Any], ReceiveBehavior[Any]] … Actor Runtime 49
49. // `events` and `behaviors` are serializable val events = mutable.Queue[ActorSend] val behaviors = mutable.Map[ActorRef[Any], ReceiveBehavior[Any]] … // unwrap and apply behavior to message val nxtEvent = events.dequeue val behavior = behaviors(nxtEvent.aref) val fun = behavior.sp.unwrap() // unwrap the function val msg = nxt.msg.unwrap() fun(msg) // apply function to message ... Actor Runtime 50
50. Demo: “Fibonacci” Actor 51 • Calculate Fibonacci number with actors: ◦ Idea: recursively spawn actors to calculate fib(n-1), fib(n-2) • Runs on JVM (uses Spore.auto) • Syntactically indistinguishable from other Actor frameworks • Link
51. Example: Fibonacci… What is captured? 52 def fib: ActorBehavior[FibReply] = receive[FibRequest]: ... receive[FibResult]: res1 => // uses Spore.auto receive[FibResult]: res2 => // uses Spore.auto val res = res1.result + res2.result ctx.send(replyTo, FibResult(res)) stopped
52. Example: Fibonacci… What is captured? 53 def fib: ActorBehavior[FibReply] = receive[FibRequest]: ... receive[FibResult]: res1 => // outer receive also captures `replyTo` receive[FibResult]: res2 => val res = res1.result + res2.result captured: ^^^^ ctx.send(replyTo, FibResult(res)) captured:^^^^^^^ stopped
53. class ListRW[T] extends SporeClassBuilder[ ReadWriter[T] ?=> ReadWriter[List[T]] ]({ summon }) given [T](using Spore[ReadWriter[T]]): Spore[ReadWriter[List[T]] = ListRW().build().withEnv2(summon) `given Spore[ReadWriter[List[T]]]` (Cross-Platform) 54