Serialize calls to Dns.Callback

okhttp-dnsoverhttps/src/main/kotlin/okhttp3/dnsoverhttps/internal/-DnsOverHttpsCall.kt

@@ -27,12 +27,34 @@ import okhttp3.Response
 import okhttp3.internal.OkHttpInternalApi
 import okhttp3.internal.testAndSet
 
-// TODO: honor resolvePrivateAddresses, resolvePublicAddresses
 // TODO: in-memory caching that uses timeToLive.
 // TODO: honor Https.priority and Https.targetName. Create new calls!
 
 /**
  * Implements [Dns.Call] by making multiple HTTPS calls.
+ *
+ * Concurrency
+ * -----------
+ *
+ * A few things conspire to make concurrency tricky:
+ *
+ *  * Each DNS record type is queried in parallel; [onResponse] and [onFailure] may be called
+ *    concurrently.
+ *  * Calls to [Dns.Callback] must be serialized.
+ *  * We don't want to use locks to guard access to [Dns.Callback] functions.
+ *
+ * Each time we receive data for the callback (in the form of records or an exception), we either
+ * immediately call the callback with that data (on a dispatcher thread), or queue it for the thread
+ * that's busy calling the callback.
+ *
+ * After calling a callback, the caller must check to see if there's more data queued to deliver,
+ * and deliver that also.
+ *
+ * The potentially surprising outcome of this strategy is the thread that performed the `TYPE_AAAA`
+ * DNS request may also deliver the `TYPE_A` records to the callback, or vice versa.
+ *
+ * If a thread is intending to call the callback, it sets [State.Running.lockHeld] to true while
+ * that call is executing.
  */
 @OkHttpInternalApi
 internal class DnsOverHttpsCall(
@@ -41,11 +63,16 @@ internal class DnsOverHttpsCall(
   private val canceledException: IOException?,
 ) : Dns.Call,
   Callback {
-  @Volatile private var canceled = false
+  @Volatile
+  private var canceled = false
   private val state = AtomicReference<State>(State.Idle)
 
   override fun enqueue(callback: Dns.Callback) {
-    val running = State.Running(callback, calls)
+    val running =
+      State.Running(
+        callback = callback,
+        runningCalls = calls,
+      )
 
     val previous = state.testAndSet(running) { it is State.Idle }
     check(previous is State.Idle) {
@@ -65,43 +92,10 @@ internal class DnsOverHttpsCall(
     call: Call,
     e: IOException,
   ) {
-    while (true) {
-      val previous =
-        state.get() as? State.Running
-          ?: return // Already complete or canceled; nothing to do.
-
-      val newRunningCalls = previous.runningCalls - call
-      val allFailures =
-        when {
-          canceledException != null -> listOf(canceledException)
-          else -> previous.delayedFailures + e
-        }
-      val next =
-        when {
-          newRunningCalls.isEmpty() -> {
-            State.Complete
-          }
-
-          else -> {
-            State.Running(
-              callback = previous.callback,
-              runningCalls = newRunningCalls,
-              delayedFailures = allFailures,
-            )
-          }
-        }
-
-      if (!state.compareAndSet(previous, next)) continue // Lost a race; retry.
-
-      if (next is State.Complete) {
-        previous.callback.onFailure(
-          call = this,
-          exceptions = allFailures,
-        )
-      }
-
-      return
-    }
+    updateStateAndCallCallbacks(
+      completedCall = call,
+      newException = e,
+    )
   }
 
   override fun onResponse(
@@ -112,7 +106,10 @@ internal class DnsOverHttpsCall(
       try {
         decodeResponse(response)
       } catch (e: IOException) {
-        return onFailure(call, e)
+        return updateStateAndCallCallbacks(
+          completedCall = call,
+          newException = e,
+        )
       }
 
     val dnsRecords =
@@ -144,47 +141,102 @@ internal class DnsOverHttpsCall(
         }
       }
 
+    updateStateAndCallCallbacks(
+      completedCall = call,
+      newRecords = dnsRecords,
+    )
+  }
+
+  private tailrec fun updateStateAndCallCallbacks(
+    completedCall: Call? = null,
+    newRecords: List<Dns.Record> = listOf(),
+    newException: IOException? = null,
+    lockHeldByThisThread: Boolean = false,
+  ) {
     while (true) {
       val previous =
         state.get() as? State.Running
           ?: return // Already complete or canceled; nothing to do.
 
-      val newRunningCalls = previous.runningCalls - call
-      val lastRunningCall = newRunningCalls.isEmpty()
+      val newRunningCalls =
+        when {
+          completedCall != null -> previous.runningCalls - completedCall
+          else -> previous.runningCalls
+        }
+
+      val allExceptions =
+        when {
+          canceledException != null -> listOf(canceledException)
+          newException != null -> previous.pendingExceptions + newException
+          else -> previous.pendingExceptions
+        }
+
+      val allRecords =
+        when {
+          newRecords.isNotEmpty() -> previous.pendingRecords + newRecords
+          else -> previous.pendingRecords
+        }
+
+      val last = newRunningCalls.isEmpty()
+      val lockHeldByAnotherThread = !lockHeldByThisThread && previous.lockHeld
+
+      // There's a few reasons why we might not call any callbacks:
+      //  - There's no more records or failures to emit immediately
+      //  - Another thread is already calling the callbacks.
+      // In such cases, hand off any new work to that other thread and be done.
+      if ((!last && allRecords.isEmpty()) || lockHeldByAnotherThread) {
+        val next =
+          State.Running(
+            callback = previous.callback,
+            runningCalls = newRunningCalls,
+            lockHeld = lockHeldByAnotherThread,
+            pendingRecords = allRecords,
+            pendingExceptions = allExceptions,
+          )
+        if (!state.compareAndSet(previous, next)) continue // Lost a race, retry.
+        return
+      }
+
+      // We need to call a callback. Take the lock and the records.
       val next =
         when {
-          lastRunningCall -> {
+          last -> {
             State.Complete
           }
 
           else -> {
             State.Running(
               callback = previous.callback,
               runningCalls = newRunningCalls,
-              delayedFailures = previous.delayedFailures,
+              lockHeld = true,
+              pendingRecords = listOf(),
+              pendingExceptions = allExceptions,
             )
           }
         }
+      if (!state.compareAndSet(previous, next)) continue // Lost a race, retry.
 
-      if (!state.compareAndSet(previous, next)) continue // Lost a race; retry.
-
-      val emitDelayedFailures = lastRunningCall && previous.delayedFailures.isNotEmpty()
-      val lastEvent = lastRunningCall && !emitDelayedFailures
-      if (dnsRecords.isNotEmpty() || lastEvent) {
+      val lastAndNoExceptions = last && allExceptions.isEmpty()
+      if (allRecords.isNotEmpty() || lastAndNoExceptions) {
         previous.callback.onRecords(
           call = this,
-          last = lastEvent,
-          records = dnsRecords,
+          last = lastAndNoExceptions,
+          records = allRecords,
         )
       }
-      if (emitDelayedFailures) {
+
+      if (last && allExceptions.isNotEmpty()) {
         previous.callback.onFailure(
           call = this,
-          exceptions = previous.delayedFailures,
+          exceptions = allExceptions,
         )
       }
 
-      return
+      // Success, attempt to release the held lock. This might also process more events if some
+      // were enqueued while the callback was executing.
+      return updateStateAndCallCallbacks(
+        lockHeldByThisThread = true,
+      )
     }
   }
 
@@ -204,9 +256,15 @@ internal class DnsOverHttpsCall(
 
     class Running(
       val callback: Dns.Callback,
+      val lockHeld: Boolean = false,
       val runningCalls: List<Call>,
-      val delayedFailures: List<IOException> = listOf(),
-    ) : State
+      val pendingRecords: List<Dns.Record> = listOf(),
+      val pendingExceptions: List<IOException> = listOf(),
+    ) : State {
+      init {
+        check(pendingRecords.isEmpty() || lockHeld)
+      }
+    }
 
     object Complete : State
   }

okhttp-dnsoverhttps/src/test/java/okhttp3/dnsoverhttps/DnsOverHttpsTest.kt

@@ -19,6 +19,7 @@ import app.cash.burst.Burst
 import assertk.assertThat
 import assertk.assertions.containsExactly
 import assertk.assertions.hasMessage
+import assertk.assertions.hasSize
 import assertk.assertions.isEqualTo
 import assertk.assertions.isInstanceOf
 import assertk.assertions.isNull
@@ -27,6 +28,11 @@ import java.io.EOFException
 import java.io.File
 import java.net.InetAddress
 import java.net.UnknownHostException
+import java.util.concurrent.CompletableFuture
+import java.util.concurrent.LinkedBlockingDeque
+import java.util.concurrent.SynchronousQueue
+import java.util.concurrent.ThreadPoolExecutor
+import java.util.concurrent.TimeUnit
 import kotlin.reflect.KClass
 import kotlin.test.assertFailsWith
 import mockwebserver3.MockResponse
@@ -93,9 +99,26 @@ class DnsOverHttpsTest(
         dispatcher = server.dispatcher
       }
 
+  private val completedRunnables = LinkedBlockingDeque<Runnable>()
+  private val executorService =
+    object : ThreadPoolExecutor(
+      0,
+      Int.MAX_VALUE,
+      1,
+      TimeUnit.SECONDS,
+      SynchronousQueue(),
+    ) {
+      override fun afterExecute(
+        r: Runnable,
+        t: Throwable?,
+      ) {
+        completedRunnables.add(r)
+      }
+    }
+
   private lateinit var dns: DnsOverHttps
   private val dispatcher =
-    Dispatcher()
+    Dispatcher(executorService)
       .apply {
         // Force calls to run sequentially for determinism.
         this.maxRequestsPerHost = 1
@@ -671,6 +694,90 @@ class DnsOverHttpsTest(
     assertThat(dnsEvents.take()).isInstanceOf<DnsEvent.Failure>()
   }
 
+  @Test
+  fun callbackIsCalledSequentiallyWhenHttpCallsAreParallel() {
+    callbackIsCalledSequentially()
+  }
+
+  @Test
+  fun callbackIsCalledSequentiallyWhenCallsFail() {
+    server.sequenceIndexToOverride[0] = overrideResponse("")
+    callbackIsCalledSequentially()
+  }
+
+  /**
+   * We'd like to confirm that calls into [Dns.Callback] are serialized.
+   *
+   * Asserting that is awkward, so we're asserting something stronger, advised by the
+   * implementation: if the first call to `onRecords()` waits for the other 2 HTTP calls to
+   * finish, then all 3 calls will happen on the same thread.
+   *
+   * (This works because we know the other 2 HTTP calls will finish before their data is emitted.)
+   */
+  private fun callbackIsCalledSequentially() {
+    assumeTrue(entryPoint == EntryPoint.NewCall)
+
+    dns = buildLocalhost(bootstrapClient, includeIPv6 = true, includeHttps = true)
+    server["lysine.dev"] =
+      listOf(
+        ResourceRecord.IpAddress(
+          name = "lysine.dev",
+          timeToLive = 5,
+          address = InetAddress.getByName("10.20.30.40"),
+        ),
+        ResourceRecord.IpAddress(
+          name = "lysine.dev",
+          timeToLive = 5,
+          address = InetAddress.getByName("1:2::3:4"),
+        ),
+        ResourceRecord.Https(
+          name = "lysine.dev",
+          timeToLive = 5,
+          alpnIds = listOf(Protocol.HTTP_2.toString()),
+        ),
+      )
+
+    // Track which threads receive callbacks.
+    val threadsFuture = CompletableFuture<Set<Thread>>()
+
+    dispatcher.maxRequestsPerHost = 3
+
+    val call = dns.newCall(Dns.Request("lysine.dev"))
+    call.enqueue(
+      object : Dns.Callback {
+        private val threads = mutableSetOf<Thread>()
+
+        override fun onRecords(
+          call: Dns.Call,
+          last: Boolean,
+          records: List<Dns.Record>,
+        ) {
+          onCallback(last)
+        }
+
+        override fun onFailure(
+          call: Dns.Call,
+          e: java.io.IOException,
+        ) {
+          onCallback(true)
+        }
+
+        private fun onCallback(last: Boolean) {
+          if (threads.isEmpty()) {
+            completedRunnables.take()
+            completedRunnables.take()
+          }
+          threads += Thread.currentThread()
+          if (last) {
+            threadsFuture.complete(threads)
+          }
+        }
+      },
+    )
+
+    assertThat(threadsFuture.get()).hasSize(1)
+  }
+
   private fun cacheEvents(): List<KClass<out CallEvent>> =
     eventRecorder
       .recordedEventTypes()