feat: Make autobatch queue denote batches by payload size, add aggregation delay (#137)

* Make autobatch queue denote batches by payload size, add aggregation delay

* Fully drain queue when batch delay is hit

* Feedback on locks

Author: jacoblee93

AGENTS.md

@@ -233,7 +233,12 @@ Tests skip gracefully via `assumeTrue` if keys are missing.
 
 ## Code style
 
-- For cross-method concurrency coordination, prefer an explicit named `ReentrantLock` with `lock()` / `try` / `finally { unlock() }` over `synchronized` when review clarity matters. Keep the locked section minimal and do slow/blocking work outside the lock.
+- For cross-method concurrency coordination, prefer an explicit named `ReentrantLock` over `synchronized` when review clarity matters. Use Kotlin's `withLock { ... }` extension instead of manual `lock()` / `try` / `finally { unlock() }` unless explicit lock management is required. Keep the locked section minimal and do slow/blocking work outside the lock.
+- Choose the simplest concurrency primitive that fits the state being protected:
+  - Use atomic types (`AtomicBoolean`, `AtomicInteger`, etc.) for simple flags, counters, and compare-and-set state.
+  - Use `synchronized` only for small, local critical sections where a named lock would not improve clarity.
+  - Use `ReentrantReadWriteLock` when reads are frequent, writes are infrequent, and concurrent reads materially help.
+  - Use coroutine `Mutex.withLock { ... }` for coroutine-based concurrency instead of blocking thread locks.
 - `toString()` should be single-line, following the `ClassName{field=value, field=value}` convention used by the rest of the SDK.
 - Avoid `@Suppress("UNCHECKED_CAST")` — restructure code to use safe patterns (`as? String`, `is Map<*, *>` with `entries.associate`, etc). When unavoidable (e.g. generic type erasure after an `is` check), add a comment explaining why the cast is safe.
 - Use named arguments for constructor/function calls with 2+ parameters, especially when types could be confused:

langsmith-java-core/src/main/kotlin/com/langchain/smith/client/AutoBatchIngestLimits.kt

@@ -0,0 +1,19 @@
+package com.langchain.smith.client
+
+import com.langchain.smith.models.annotationqueues.info.InfoListResponse.BatchIngestConfig
+import kotlin.jvm.optionals.getOrNull
+
+internal data class AutoBatchIngestLimits(
+    val batchSizeLimit: Int = AutoBatchQueue.DEFAULT_BATCH_SIZE_LIMIT,
+    val batchSizeLimitBytes: Int = AutoBatchQueue.DEFAULT_BATCH_SIZE_LIMIT_BYTES,
+)
+
+internal fun BatchIngestConfig?.toAutoBatchIngestLimits(): AutoBatchIngestLimits =
+    AutoBatchIngestLimits(
+        batchSizeLimit =
+            this?.sizeLimit()?.getOrNull()?.takeIf { it > 0 && it <= Int.MAX_VALUE }?.toInt()
+                ?: AutoBatchQueue.DEFAULT_BATCH_SIZE_LIMIT,
+        batchSizeLimitBytes =
+            this?.sizeLimitBytes()?.getOrNull()?.takeIf { it > 0 && it <= Int.MAX_VALUE }?.toInt()
+                ?: AutoBatchQueue.DEFAULT_BATCH_SIZE_LIMIT_BYTES,
+    )

langsmith-java-core/src/main/kotlin/com/langchain/smith/client/AutoBatchQueue.kt

@@ -17,10 +17,12 @@ import java.util.concurrent.Executors
 import java.util.concurrent.Phaser
 import java.util.concurrent.RejectedExecutionException
 import java.util.concurrent.ScheduledExecutorService
+import java.util.concurrent.ScheduledFuture
 import java.util.concurrent.TimeUnit
 import java.util.concurrent.atomic.AtomicBoolean
 import java.util.concurrent.atomic.AtomicInteger
 import java.util.concurrent.locks.ReentrantLock
+import kotlin.concurrent.withLock
 import kotlin.jvm.optionals.getOrNull
 import org.slf4j.LoggerFactory
 
@@ -35,20 +37,28 @@ import org.slf4j.LoggerFactory
  *
  * @param sendBatch sends a batch and completes when the send has finished
  * @param batchSizeLimit max operations before auto-flush (default 100)
- * @param aggregationDelayMs delay before timer-based flush (default 250ms)
+ * @param aggregationDelayMs inactivity delay before timer-based flush (default 250ms, matching JS
+ *   auto-batch aggregation and Python's initial queue wait)
  * @param sendParallelism max number of batch requests to send concurrently (default 4)
+ * @param maxAggregationDelayMs max delay before timer-based flush after first item in a burst
+ *   (default 500ms, matching Python's background flush interval)
+ * @param batchSizeLimitBytes max serialized JSON body size per batch (default 20 MiB)
  */
 class AutoBatchQueue(
     private val sendBatch: (RunIngestBatchParams, RequestOptions) -> CompletionStage<Void?>,
     private val batchSizeLimit: Int = DEFAULT_BATCH_SIZE_LIMIT,
     private val aggregationDelayMs: Long = DEFAULT_AGGREGATION_DELAY_MS,
     private val sendParallelism: Int = DEFAULT_SEND_PARALLELISM,
+    private val maxAggregationDelayMs: Long = DEFAULT_MAX_AGGREGATION_DELAY_MS,
+    private val batchSizeLimitBytes: Int,
 ) {
     private val items = ConcurrentLinkedQueue<BatchItem>()
     private val queuedCount = AtomicInteger(0)
     private val shutdown = AtomicBoolean(false)
-    private val delayedFlushScheduled = AtomicBoolean(false)
     private val enqueueShutdownLock = ReentrantLock()
+    private val delayedFlushLock = ReentrantLock()
+    private var delayedFlushFuture: ScheduledFuture<*>? = null
+    private var firstQueuedAtNanos: Long? = null
     private val activeSends =
         object : Phaser(0) {
             override fun onAdvance(phase: Int, registeredParties: Int): Boolean = false
@@ -91,6 +101,7 @@ class AutoBatchQueue(
      * Safe to call from any thread. No-op if the queue is empty.
      */
     fun flush() {
+        cancelDelayedFlush()
         while (true) {
             if (queuedCount.get() > 0 && !drainOnCoordinator()) {
                 return
@@ -112,14 +123,13 @@ class AutoBatchQueue(
      * After calling this, the queue will no longer accept new operations.
      */
     fun shutdown() {
-        enqueueShutdownLock.lock()
-        try {
-            // Serialize with enqueue's check-and-add so flush cannot miss an item that observed
-            // shutdown=false but has not yet been queued.
-            if (!shutdown.compareAndSet(false, true)) return
-        } finally {
-            enqueueShutdownLock.unlock()
-        }
+        val startedShutdown =
+            enqueueShutdownLock.withLock {
+                // Serialize with enqueue's check-and-add so flush cannot miss an item that observed
+                // shutdown=false but has not yet been queued.
+                shutdown.compareAndSet(false, true)
+            }
+        if (!startedShutdown) return
 
         flush()
         coordinator.shutdown()
@@ -139,6 +149,20 @@ class AutoBatchQueue(
         }
     }
 
+    constructor(
+        sendBatch: (RunIngestBatchParams, RequestOptions) -> CompletionStage<Void?>,
+        batchSizeLimit: Int = DEFAULT_BATCH_SIZE_LIMIT,
+        aggregationDelayMs: Long = DEFAULT_AGGREGATION_DELAY_MS,
+        sendParallelism: Int = DEFAULT_SEND_PARALLELISM,
+    ) : this(
+        sendBatch = sendBatch,
+        batchSizeLimit = batchSizeLimit,
+        aggregationDelayMs = aggregationDelayMs,
+        sendParallelism = sendParallelism,
+        maxAggregationDelayMs = DEFAULT_MAX_AGGREGATION_DELAY_MS,
+        batchSizeLimitBytes = DEFAULT_BATCH_SIZE_LIMIT_BYTES,
+    )
+
     /** Returns the number of queued operations (for testing). */
     internal fun size(): Int = queuedCount.get()
 
@@ -149,9 +173,8 @@ class AutoBatchQueue(
         queryParams: QueryParams,
         requestOptions: RequestOptions,
     ) {
-        val count = run {
-            enqueueShutdownLock.lock()
-            try {
+        val count =
+            enqueueShutdownLock.withLock {
                 check(!shutdown.get()) { "AutoBatchQueue is shut down" }
                 items.add(
                     BatchItem(
@@ -163,10 +186,7 @@ class AutoBatchQueue(
                     )
                 )
                 queuedCount.incrementAndGet()
-            } finally {
-                enqueueShutdownLock.unlock()
             }
-        }
 
         afterEnqueue(count)
     }
@@ -180,31 +200,56 @@ class AutoBatchQueue(
     }
 
     private fun scheduleFlush() {
-        if (!delayedFlushScheduled.compareAndSet(false, true)) return
+        delayedFlushLock.withLock {
+            try {
+                val now = System.nanoTime()
+                val firstQueuedAt = firstQueuedAtNanos ?: now.also { firstQueuedAtNanos = it }
+                val elapsedMs = TimeUnit.NANOSECONDS.toMillis(now - firstQueuedAt)
+                val remainingMaxDelayMs = (maxAggregationDelayMs - elapsedMs).coerceAtLeast(0)
+                val delayMs = minOf(aggregationDelayMs, remainingMaxDelayMs)
+
+                delayedFlushFuture?.cancel(false)
+                delayedFlushFuture =
+                    coordinator.schedule(
+                        {
+                            clearDelayedFlushState()
+                            drainAndSubmitSends()
+                        },
+                        delayMs,
+                        TimeUnit.MILLISECONDS,
+                    )
+            } catch (e: RejectedExecutionException) {
+                delayedFlushFuture = null
+                firstQueuedAtNanos = null
+                logger.warn("Batch queue coordinator rejected delayed flush", e)
+            }
+        }
+    }
 
-        try {
-            coordinator.schedule(
-                {
-                    delayedFlushScheduled.set(false)
-                    drainAndSubmitSends()
-                },
-                aggregationDelayMs,
-                TimeUnit.MILLISECONDS,
-            )
-        } catch (e: RejectedExecutionException) {
-            delayedFlushScheduled.set(false)
-            logger.warn("Batch queue coordinator rejected delayed flush", e)
+    private fun clearDelayedFlushState() {
+        delayedFlushLock.withLock {
+            delayedFlushFuture = null
+            firstQueuedAtNanos = null
         }
     }
 
     private fun triggerFlush() {
+        cancelDelayedFlush()
         try {
             coordinator.execute { drainAndSubmitSends() }
         } catch (e: RejectedExecutionException) {
             logger.warn("Batch queue coordinator rejected flush", e)
         }
     }
 
+    private fun cancelDelayedFlush() {
+        delayedFlushLock.withLock {
+            delayedFlushFuture?.cancel(false)
+            delayedFlushFuture = null
+            firstQueuedAtNanos = null
+        }
+    }
+
     private fun drainOnCoordinator(): Boolean {
         val drainFuture =
             try {
@@ -225,10 +270,14 @@ class AutoBatchQueue(
     }
 
     private fun drainAndSubmitSends() {
-        val batches = drainUpTo(batchSizeLimit)
-        if (batches.isEmpty()) return
+        var remainingToDrain = queuedCount.get()
+        while (remainingToDrain > 0) {
+            val drainResult = drainUpTo(minOf(batchSizeLimit, remainingToDrain))
+            if (drainResult.itemCount == 0) break
 
-        batches.forEach(::submitBatch)
+            drainResult.batches.forEach(::submitBatch)
+            remainingToDrain -= drainResult.itemCount
+        }
 
         when {
             queuedCount.get() >= batchSizeLimit -> triggerFlush()
@@ -241,30 +290,36 @@ class AutoBatchQueue(
      * Drains up to [maxItems] queued operations and returns batch params grouped by request
      * options.
      *
-     * TODO: Also flush/split batches based on serialized payload size, not just operation count.
      * TODO: Support multipart ingest endpoint for large payloads with attachments.
      * TODO: Support gzip compression for batch requests.
      */
-    private fun drainUpTo(maxItems: Int): List<Batch> {
-        val groups = linkedMapOf<RequestOptionsKey, BatchGroup>()
+    private fun drainUpTo(maxItems: Int): DrainResult {
+        val openGroups = linkedMapOf<RequestOptionsKey, BatchGroup>()
         var drained = 0
 
-        while (drained < maxItems) {
-            val item = items.poll() ?: break
-            queuedCount.decrementAndGet()
-            drained++
-
-            val group =
-                groups.getOrPut(item.requestOptions.key()) { BatchGroup(item.requestOptions) }
-            when (item.op) {
-                BatchOp.Post -> group.posts.add(item.run)
-                BatchOp.Patch -> group.patches.add(item.run)
+        val batches = buildList {
+            while (drained < maxItems) {
+                val item = items.poll() ?: break
+                queuedCount.decrementAndGet()
+                drained++
+
+                val key = item.requestOptions.key()
+                val itemSerializedSize = objectMapper.writeValueAsBytes(item.run).size
+                var group = openGroups.getOrPut(key) { BatchGroup(item.requestOptions) }
+                if (
+                    group.isNotEmpty() &&
+                        group.serializedSizeWith(item, itemSerializedSize) > batchSizeLimitBytes
+                ) {
+                    add(group.toBatch())
+                    group = BatchGroup(item.requestOptions)
+                    openGroups[key] = group
+                }
+                group.add(item, itemSerializedSize)
             }
-            group.headers.putAll(item.headers)
-            group.queryParams.putAll(item.queryParams)
-        }
 
-        return groups.values.map { it.toBatch() }
+            addAll(openGroups.values.map { it.toBatch() })
+        }
+        return DrainResult(batches = batches, itemCount = drained)
     }
 
     private fun submitBatch(batch: Batch) {
@@ -314,15 +369,71 @@ class AutoBatchQueue(
         Patch,
     }
 
+    private data class DrainResult(val batches: List<Batch>, val itemCount: Int)
+
     private data class Batch(val params: RunIngestBatchParams, val requestOptions: RequestOptions)
 
-    private data class BatchGroup(
+    private class BatchGroup(
         val requestOptions: RequestOptions,
         val posts: MutableList<Run> = mutableListOf(),
         val patches: MutableList<Run> = mutableListOf(),
         val headers: Headers.Builder = Headers.builder(),
         val queryParams: QueryParams.Builder = QueryParams.builder(),
     ) {
+        private var postSerializedSizeBytes = 0
+        private var patchSerializedSizeBytes = 0
+
+        fun add(item: BatchItem, serializedSizeBytes: Int) {
+            when (item.op) {
+                BatchOp.Post -> {
+                    posts.add(item.run)
+                    postSerializedSizeBytes += serializedSizeBytes
+                }
+                BatchOp.Patch -> {
+                    patches.add(item.run)
+                    patchSerializedSizeBytes += serializedSizeBytes
+                }
+            }
+            headers.putAll(item.headers)
+            queryParams.putAll(item.queryParams)
+        }
+
+        fun isNotEmpty(): Boolean = posts.isNotEmpty() || patches.isNotEmpty()
+
+        fun serializedSizeWith(item: BatchItem, itemSerializedSizeBytes: Int): Int =
+            estimatedBodySizeBytes(
+                postCount = posts.size + if (item.op == BatchOp.Post) 1 else 0,
+                postItemsSizeBytes =
+                    postSerializedSizeBytes +
+                        if (item.op == BatchOp.Post) itemSerializedSizeBytes else 0,
+                patchCount = patches.size + if (item.op == BatchOp.Patch) 1 else 0,
+                patchItemsSizeBytes =
+                    patchSerializedSizeBytes +
+                        if (item.op == BatchOp.Patch) itemSerializedSizeBytes else 0,
+            )
+
+        private fun estimatedBodySizeBytes(
+            postCount: Int,
+            postItemsSizeBytes: Int,
+            patchCount: Int,
+            patchItemsSizeBytes: Int,
+        ): Int {
+            var size = 2 // Opening and closing braces.
+            var fieldCount = 0
+            if (patchCount > 0) {
+                size += fieldSizeBytes("patch", patchCount, patchItemsSizeBytes)
+                fieldCount++
+            }
+            if (postCount > 0) {
+                if (fieldCount > 0) size++ // Comma between top-level fields.
+                size += fieldSizeBytes("post", postCount, postItemsSizeBytes)
+            }
+            return size
+        }
+
+        private fun fieldSizeBytes(name: String, count: Int, itemsSizeBytes: Int): Int =
+            name.length + 5 + itemsSizeBytes + (count - 1) // `"name":[]` + item commas.
+
         fun toBatch(): Batch {
             val mergeResult = mergePostsAndPatches()
             val builder = RunIngestBatchParams.builder()
@@ -394,7 +505,9 @@ class AutoBatchQueue(
         private val objectMapper = jsonMapper()
 
         const val DEFAULT_BATCH_SIZE_LIMIT = 100
+        const val DEFAULT_BATCH_SIZE_LIMIT_BYTES = 20_971_520
         const val DEFAULT_AGGREGATION_DELAY_MS = 250L
+        const val DEFAULT_MAX_AGGREGATION_DELAY_MS = 500L
         const val DEFAULT_SEND_PARALLELISM = 4
     }
 }