perf(avro): reuse write plumbing in writeDatum — attribution + fix for #119

.github/bench/bench_tools.py

@@ -42,6 +42,7 @@
 MODULE_BENCHES = {
     "avro/": [
         "AvroBytesBench",
+        "AvroEncodeRouteBench",
         "AvroJsonBridgeBench",
         "AvroVulcanBench",
         "OrderAvroBench",

avro/src/main/scala/dev/constructive/eo/avro/AvroBinaryCursor.scala

@@ -7,7 +7,14 @@ import java.io.{ByteArrayOutputStream, InputStream}
 import java.util.{Arrays, HashMap, List as JList}
 import org.apache.avro.Schema
 import org.apache.avro.generic.{GenericDatumReader, GenericDatumWriter, IndexedRecord}
-import org.apache.avro.io.{BinaryData, BinaryDecoder, Decoder, DecoderFactory, EncoderFactory}
+import org.apache.avro.io.{
+  BinaryData,
+  BinaryDecoder,
+  BinaryEncoder,
+  Decoder,
+  DecoderFactory,
+  EncoderFactory
+}
 
 /** Internal byte-offset locator behind [[AvroPrism]]'s byte-carried optic (`to`/`from`) and its
   * slice/graft surface.
@@ -357,17 +364,64 @@ private[avro] object AvroBinaryCursor:
     */
   private[avro] val leaves = new DatumReaders[Any]
 
+  /** Per-thread reusable binary-write plumbing — [[writeDatum]]'s engine, the write mirror of the
+    * [[DatumReaders]] reader cache + [[binaryDecoderCache]] decoder on the read side (issue #119).
+    *
+    * '''Why cache at all:''' every whole-record and leaf write funnels through [[writeDatum]]. Its
+    * first form allocated a fresh `ByteArrayOutputStream` (32-byte start — a multi-KB payload
+    * reallocates and copies its buffer a dozen times on the way up), a fresh `GenericDatumWriter`,
+    * and a fresh `BufferedBinaryEncoder` (2 KB internal buffer) on EVERY call. Measured on a
+    * 15-field record with an all-`Option` nested record and an 18-branch union (245 B output):
+    * 2,928 B/op of write-side plumbing against a 392 B/op reused-plumbing floor and 397 B/op for a
+    * hand-written direct-to-encoder writer — the plumbing, not `GenericDatumWriter` dispatch, is
+    * the bulk of the allocation gap reported in issue #119 over hand-written producers.
+    *
+    * '''Rebind, not rebuild:''' `EncoderFactory.binaryEncoder(out, reuse)` reconfigures a
+    * `BufferedBinaryEncoder` in place — position reset to 0, same 2 KB buffer (only replaced when
+    * the factory buffer size changes) — so steady-state writes allocate exactly the returned
+    * `toByteArray` result. The writer cache is keyed per `Schema` like the read cache is keyed per
+    * schema pair; `GenericDatumWriter` is mutable and NOT thread-safe, hence the `ThreadLocal` —
+    * the same reason [[DatumReaders.cache]] exists.
+    *
+    * '''Lifetime caveat (same as [[DatumReaders]]):''' the writer map has no eviction — it grows by
+    * one entry per distinct schema written on the thread. Producers write one schema repeatedly, so
+    * this is the intended shape; dynamically-built schemas on a huge pool would instead want fresh
+    * writers per call, which is exactly what the pre-#119 form did.
+    */
+  final private class DatumWriters:
+
+    private val out = new ByteArrayOutputStream(4096)
+    private var encoder: BinaryEncoder = EncoderFactory.get().binaryEncoder(out, null)
+    private val cache = new HashMap[Schema, GenericDatumWriter[Any]]()
+
+    /** Encode `datum` (of `schema` shape) to a fresh, detached `Array[Byte]`. Safe against a
+      * mid-write throw by ordering, not by state: `out` is reset at the START of every encode and
+      * the previous result was already copied out by `toByteArray`, so dirty bytes from an aborted
+      * write are never observable.
+      */
+    def write(datum: Any, schema: Schema): Array[Byte] =
+      out.reset()
+      encoder = EncoderFactory.get().binaryEncoder(out, encoder)
+      cache
+        .computeIfAbsent(schema, s => new GenericDatumWriter[Any](s))
+        .write(datum, encoder)
+      encoder.flush()
+      out.toByteArray
+
+  end DatumWriters
+
+  /** One [[DatumWriters]] plumbing set per writing thread — see its scaladoc for why reuse is
+    * thread-local.
+    */
+  private val writeCache: ThreadLocal[DatumWriters] =
+    ThreadLocal.withInitial(() => new DatumWriters)
+
   /** THE module's binary write — [[DatumReaders.read]]'s mirror: encode an `Any`-shaped `datum`
-    * under `schema` to its binary wire form. Fresh writer/encoder per call: `GenericDatumWriter`
-    * carries no resolution state worth caching.
+    * under `schema` to its binary wire form. Delegates to the per-thread [[DatumWriters]] plumbing
+    * (issue #119); the returned array is always freshly copied.
     */
   private[avro] def writeDatum(datum: Any, schema: Schema): Array[Byte] =
-    val out = new ByteArrayOutputStream()
-    val writer = new GenericDatumWriter[Any](schema)
-    val encoder = EncoderFactory.get().binaryEncoder(out, null)
-    writer.write(datum, encoder)
-    encoder.flush()
-    out.toByteArray
+    writeCache.get().write(datum, schema)
 
   /** Read a `ByteBuffer`'s remaining bytes without disturbing its position — how a `bytes` field
     * arrives in the generic runtime model. Shared by the bridges' structural walks (`AvroJson` /

avro/src/test/scala/dev/constructive/eo/avro/AvroWriteCorrectnessSpec.scala

@@ -479,4 +479,56 @@ class AvroWriteCorrectnessSpec extends Specification with ScalaCheck:
     codecPrism[FullName].field(_.first).getOption(bytes) === Some("Doe")
   }
 
+  // ---- reused write plumbing (issue #119) -----------------------------
+
+  // covers: writeDatum's per-thread reused ByteArrayOutputStream / BinaryEncoder /
+  //   GenericDatumWriter is safe on the three axes reuse introduces:
+  //   (a) RETENTION — a result array handed to the caller is never disturbed by later writes on
+  //       the same thread (every result is copied out by toByteArray); note the byte-face optic
+  //       contracts already promise freshly-returned arrays.
+  //   (b) ABORTED WRITE — a datum/schema mismatch throws mid-encode (surfaced as Left); the
+  //       dirty buffer it leaves behind must never leak into the next write's bytes.
+  //   (c) THREAD ISOLATION — the plumbing is ThreadLocal, so concurrent writers interleave
+  //       freely and every thread's bytes equal the single-writer golden bytes.
+  "reused write plumbing: retention, aborted write, thread isolation (issue #119)" >> {
+    val pc = summon[AvroCodec[Person]]
+    val p = Person("Alice", 42)
+    val q = Person("Bob", 404)
+
+    def encode(x: Person): Array[Byte] =
+      AvroCodec.encodeValue(x)(using pc).getOrElse(throw new RuntimeException("encode failed"))
+
+    val goldenP = encode(p)
+    val goldenQ = encode(q)
+    val snapshotP = goldenP.clone()
+    val snapshotQ = goldenQ.clone()
+
+    // 200 interleaved writes on THIS thread must not disturb the two retained results.
+    (1 to 200).foreach { n => encode(if n % 2 == 0 then p else q); () }
+    val churnOk = Arrays.equals(goldenP, snapshotP) && Arrays.equals(goldenQ, snapshotQ)
+
+    val aborted = AvroCodec.encodeRecord("definitely not a record", pc.schema)
+    val cleanAfterAbort = Arrays.equals(goldenP, encode(p))
+
+    val results = new java.util.concurrent.ConcurrentLinkedQueue[Boolean]()
+    val writers = (1 to 8).toList.map { _ =>
+      val t = new Thread(() =>
+        (1 to 50).foreach { n =>
+          val x = if n % 2 == 0 then p else q
+          val golden = if n % 2 == 0 then goldenP else goldenQ
+          results.add(Arrays.equals(golden, encode(x)))
+        }
+      )
+      t.start()
+      t
+    }
+    writers.foreach(_.join())
+
+    (aborted.isLeft === true)
+      .and(churnOk === true)
+      .and(cleanAfterAbort === true)
+      .and(results.size() === 400)
+      .and(results.contains(false) === false)
+  }
+
 end AvroWriteCorrectnessSpec

benchmarks/src/main/scala/dev/constructive/eo/bench/AvroEncodeRouteBench.scala

@@ -0,0 +1,222 @@
+package dev.constructive.eo
+package bench
+
+import scala.compiletime.uninitialized
+
+import java.io.ByteArrayOutputStream
+import java.util.concurrent.TimeUnit
+
+import avro.AvroCodec
+import hearth.kindlings.avroderivation.{AvroConfig, AvroDecoder, AvroEncoder, AvroSchemaFor}
+import org.apache.avro.Schema
+import org.apache.avro.generic.GenericDatumWriter
+import org.apache.avro.io.{BinaryEncoder, EncoderFactory}
+import org.openjdk.jmh.annotations.*
+import scala.jdk.CollectionConverters.*
+
+/** Fixtures + routes for [[AvroEncodeRouteBench]]. Kept top-level: kindlings' derivation, like
+  * hearth's constructor synthesis, must not see an outer accessor.
+  */
+object EncodeRouteImpls:
+
+  given AvroConfig = AvroConfig()
+
+  /** All-`Option` nested record — the issue's 55-field `metrics` shape, scaled to 6. */
+  final case class Metrics(
+      a1: Option[Double],
+      a2: Option[Double],
+      b1: Option[Long],
+      s1: Option[String],
+      i1: Option[Int],
+      b2: Option[Boolean],
+  )
+
+  object Metrics:
+    given AvroEncoder[Metrics] = AvroEncoder.derived
+    given AvroDecoder[Metrics] = AvroDecoder.derived
+    given AvroSchemaFor[Metrics] = AvroSchemaFor.derived
+
+  /** Multi-branch union — the issue's 18-branch slot, scaled to 6. */
+  enum Event:
+    case Ev0(v: Long)
+    case Ev1(v: Long)
+    case Ev2(v: Long)
+    case Ev3(v: Long)
+    case Ev4(v: Long)
+    case Ev5(v: Long)
+
+  object Event:
+    given AvroEncoder[Event] = AvroEncoder.derived
+    given AvroDecoder[Event] = AvroDecoder.derived
+    given AvroSchemaFor[Event] = AvroSchemaFor.derived
+
+  /** 10-field top level: strings, numerics, a boolean, one optional nested record, one union. */
+  final case class Payload(
+      id: String,
+      tenant: String,
+      source: String,
+      ts: Long,
+      seq: Long,
+      amount: Double,
+      flag: Boolean,
+      kind: Int,
+      metrics: Option[Metrics],
+      event: Event,
+  )
+
+  object Payload:
+    given AvroEncoder[Payload] = AvroEncoder.derived
+    given AvroDecoder[Payload] = AvroDecoder.derived
+    given AvroSchemaFor[Payload] = AvroSchemaFor.derived
+
+  val payload: Payload =
+    Payload(
+      "id-1234567890",
+      "tenant-xyz",
+      "src.system.a",
+      1_700_000_000L,
+      987_654L,
+      123_456.789,
+      flag = true,
+      kind = 3,
+      metrics = Some(Metrics(Some(1.5), Some(2.5), None, Some("alpha"), None, Some(true))),
+      event = Event.Ev3(99L),
+    )
+
+  val codec: AvroCodec[Payload] = summon[AvroCodec[Payload]]
+  val schema: Schema = codec.schema
+  val metricsUnion: Schema = schema.getField("metrics").schema()
+
+  /** Which arm of a `union<null, X>` is the null arm — read off the DERIVED schema once at object
+    * init, so the hand-written route below never guesses the spelling and never re-derives it on
+    * the hot path.
+    */
+  def nullArm(s: Schema): Int =
+    s.getTypes.asScala.indexWhere(_.getType == Schema.Type.NULL)
+
+  private def nonNullArm(s: Schema): Schema =
+    s.getTypes.asScala.find(_.getType != Schema.Type.NULL).get
+
+  val metricsSomeIdx: Int = 1 - nullArm(metricsUnion)
+  private val metricsSchema: Schema = nonNullArm(metricsUnion)
+
+  val metricsNullIdx: Array[Int] =
+    List("a1", "a2", "b1", "s1", "i1", "b2")
+      .map(n => nullArm(metricsSchema.getField(n).schema()))
+      .toArray
+
+/** Whole-record ENCODE attribution — the route decomposition issue #119 asked for.
+  *
+  * The report measured `AvroCodec.derived` encode at ~1.4–1.7× a hand-written
+  * direct-`BinaryEncoder` writer and ~9–18% more allocated, and hypothesised the intermediate
+  * `GenericData.Record` tree. Four routes split the cost so the remaining gap is attributable:
+  *
+  *   - `eo_encodeValue` — the full production route: kindlings A → Any, then the module's
+  *     `writeDatum`. Post-#119 `writeDatum` writes through per-thread reused buffer/encoder/writer
+  *     plumbing, so its allocation over `eo_encodeToAny` is (almost) exactly the returned `byte[]`.
+  *   - `naive_freshPlumbing` — the pre-#119 shape reproduced here as a baseline: a FRESH
+  *     `ByteArrayOutputStream` (32-byte start, doubling), `GenericDatumWriter` and
+  *     `BufferedBinaryEncoder` per call over the SAME prebuilt tree. The delta vs `eo_encodeValue`
+  *     is the per-call plumbing churn the fix removed — the dominant write-side allocator, not
+  *     `GenericDatumWriter` dispatch.
+  *   - `eo_encodeToAny` — the kindlings tree build alone (the reporter's hypothesis: what the
+  *     generic-record materialisation costs, isolated).
+  *   - `handwritten_stream` — the reporter's baseline: fields straight to a reused `BinaryEncoder`
+  *     in schema order, no tree. The gap to `eo_encodeValue` is tree + `GenericDatumWriter`
+  *     dispatch — the irreducible cost of going through avro's datum model, and what a true
+  *     streaming `A => Encoder => Unit` derivation would eventually remove.
+  *
+  * Run with the GC profiler — B/op is the metric, ns/op advises:
+  * {{{
+  *   sbt "benchmarks/Jmh/run -i 5 -wi 3 -f 3 -t 1 -prof gc .*AvroEncodeRouteBench.*"
+  * }}}
+  */
+@State(Scope.Benchmark)
+@BenchmarkMode(Array(Mode.AverageTime))
+@OutputTimeUnit(TimeUnit.NANOSECONDS)
+@Fork(3)
+@Warmup(iterations = 3, time = 1)
+@Measurement(iterations = 5, time = 1)
+class AvroEncodeRouteBench extends JmhDefaults:
+
+  import EncodeRouteImpls.*
+
+  var tree: Any = uninitialized
+  var out: ByteArrayOutputStream = uninitialized
+  var encoder: BinaryEncoder = uninitialized
+
+  @Setup(Level.Trial)
+  def init(): Unit =
+    tree = codec.encode(payload)
+    out = new ByteArrayOutputStream(16384)
+    encoder = EncoderFactory.get().binaryEncoder(out, null)
+    // sanity: the hand-written route must land on the codec's bytes, or the attribution is fiction
+    val eo = AvroCodec.encodeValue(payload)(using codec).getOrElse(null)
+    val hand = handwrittenToByteArray()
+    require(
+      eo != null && java.util.Arrays.equals(eo, hand),
+      "handwritten route diverged from eo encode"
+    )
+
+  @Benchmark def eo_encodeValue: Array[Byte] =
+    AvroCodec.encodeValue(payload)(using codec).fold(_ => null, identity)
+
+  @Benchmark def eo_encodeToAny: Any = codec.encode(payload)
+
+  /** The pre-#119 write plumbing, reconstructed: fresh BAOS + writer + encoder per call. */
+  @Benchmark def naive_freshPlumbing: Array[Byte] =
+    val o = new ByteArrayOutputStream()
+    val writer = new GenericDatumWriter[Any](schema)
+    val enc = EncoderFactory.get().binaryEncoder(o, null)
+    writer.write(tree, enc)
+    enc.flush()
+    o.toByteArray
+
+  @Benchmark def handwritten_stream: Array[Byte] = handwrittenToByteArray()
+
+  private def handwrittenToByteArray(): Array[Byte] =
+    out.reset()
+    encoder = EncoderFactory.get().binaryEncoder(out, encoder)
+    val e = encoder
+    e.writeString(payload.id)
+    e.writeString(payload.tenant)
+    e.writeString(payload.source)
+    e.writeLong(payload.ts)
+    e.writeLong(payload.seq)
+    e.writeDouble(payload.amount)
+    e.writeBoolean(payload.flag)
+    e.writeInt(payload.kind)
+    payload.metrics match
+      case Some(m) =>
+        e.writeIndex(metricsSomeIdx)
+        m.a1 match
+          case Some(v) => e.writeIndex(1 - metricsNullIdx(0)); e.writeDouble(v)
+          case None    => e.writeIndex(metricsNullIdx(0))
+        m.a2 match
+          case Some(v) => e.writeIndex(1 - metricsNullIdx(1)); e.writeDouble(v)
+          case None    => e.writeIndex(metricsNullIdx(1))
+        m.b1 match
+          case Some(v) => e.writeIndex(1 - metricsNullIdx(2)); e.writeLong(v)
+          case None    => e.writeIndex(metricsNullIdx(2))
+        m.s1 match
+          case Some(v) => e.writeIndex(1 - metricsNullIdx(3)); e.writeString(v)
+          case None    => e.writeIndex(metricsNullIdx(3))
+        m.i1 match
+          case Some(v) => e.writeIndex(1 - metricsNullIdx(4)); e.writeInt(v)
+          case None    => e.writeIndex(metricsNullIdx(4))
+        m.b2 match
+          case Some(v) => e.writeIndex(1 - metricsNullIdx(5)); e.writeBoolean(v)
+          case None    => e.writeIndex(metricsNullIdx(5))
+      case None => e.writeIndex(nullArm(metricsUnion))
+    payload.event match
+      case Event.Ev0(v) => e.writeIndex(0); e.writeLong(v)
+      case Event.Ev1(v) => e.writeIndex(1); e.writeLong(v)
+      case Event.Ev2(v) => e.writeIndex(2); e.writeLong(v)
+      case Event.Ev3(v) => e.writeIndex(3); e.writeLong(v)
+      case Event.Ev4(v) => e.writeIndex(4); e.writeLong(v)
+      case Event.Ev5(v) => e.writeIndex(5); e.writeLong(v)
+    e.flush()
+    out.toByteArray
+  end handwrittenToByteArray
+
+end AvroEncodeRouteBench