Wake takers parked in MessageQueue.take() when a read fails

dadb/src/main/kotlin/dadb/AdbMessageQueue.kt

@@ -17,6 +17,8 @@
 
 package dadb
 
+import java.io.IOException
+
 internal class AdbMessageQueue(private val adbReader: AdbReader) : AutoCloseable, MessageQueue<AdbMessage>() {
 
     override fun readMessage() = adbReader.readMessage()
@@ -25,7 +27,12 @@ internal class AdbMessageQueue(private val adbReader: AdbReader) : AutoCloseable
 
     override fun getCommand(message: AdbMessage) = message.command
 
-    override fun close() = adbReader.close()
+    override fun close() {
+        // Wake parked takers before closing the reader, so they fail with this IOException instead of
+        // stranding in await() or waking to re-read the closed reader.
+        signalClosed(IOException("Connection closed"))
+        adbReader.close()
+    }
 
     override fun isCloseCommand(message: AdbMessage) = message.command == Constants.CMD_CLSE
 }

dadb/src/main/kotlin/dadb/MessageQueue.kt

@@ -18,6 +18,7 @@
 package dadb
 
 import org.jetbrains.annotations.TestOnly
+import java.io.IOException
 import java.util.*
 import java.util.concurrent.ConcurrentHashMap
 import java.util.concurrent.ConcurrentLinkedQueue
@@ -31,20 +32,42 @@ internal abstract class MessageQueue<V> {
     private val queues = ConcurrentHashMap<Int, ConcurrentHashMap<Int, Queue<V>>>()
     private val openStreams = ConcurrentHashMap<Int, Boolean>().keySet(true)
 
+    // Terminal close marker, guarded by queueLock and set once by signalClosed(). A transient read
+    // failure does NOT set this — those stay per-call and retryable. Set only on close (when the
+    // socket also closes), so it stays in step with DadbImpl rebuilding the connection on the next op.
+    private var closedCause: IOException? = null
+
     fun take(localId: Int, command: Int): V {
         while (true) {
             queueLock.lock {
+                // Closed: fail fast rather than loop into a read on a dead reader. Fresh exception so
+                // each caller keeps its own stack.
+                closedCause?.let { throw IOException("MessageQueue closed", it) }
                 poll(localId, command)?.let { return it }
                 readLock.tryLock ({
-                    queueLock.unlock()
-                    read()
-                    queueLock.lock()
-                    queueCond.signalAll()
+                    queueLock.unlock() // released so peers can poll/queue while we read
+                    try {
+                        read()
+                    } finally {
+                        // Signal on every read() outcome, not just success: a throwing read must still
+                        // wake parked takers. The missing signal on this path was the lost-wakeup bug.
+                        // Each woken taker retries and sees the failure itself — per-call, not sticky.
+                        queueLock.lock()
+                        queueCond.signalAll()
+                    }
                 }) { queueCond.await() }
             }
         }
     }
 
+    /** Mark the queue closed and wake every parked taker so they fail fast. Call before closing the reader. */
+    protected fun signalClosed(cause: IOException) {
+        queueLock.lock {
+            if (closedCause == null) closedCause = cause
+            queueCond.signalAll()
+        }
+    }
+
     fun startListening(localId: Int) {
         openStreams.add(localId)
         queues.putIfAbsent(localId, ConcurrentHashMap())

dadb/src/test/kotlin/dadb/MessageQueueFailureContractTest.kt

@@ -0,0 +1,182 @@
+/*
+ * Copyright (c) 2021 mobile.dev inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package dadb
+
+import okio.Buffer
+import okio.Source
+import okio.Timeout
+import java.io.IOException
+import java.util.concurrent.ConcurrentLinkedQueue
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicInteger
+import kotlin.concurrent.thread
+import kotlin.test.Test
+import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
+import kotlin.test.assertIs
+import kotlin.test.assertTrue
+
+/**
+ * Pins the failure/close contract of [MessageQueue] that the lost-wakeup fix must preserve:
+ *
+ *  - a read failure wakes EVERY parked taker (not just the active reader), and surfaces as IOException;
+ *  - a read failure is per-call / non-sticky — a transient error does NOT poison the queue, so a later
+ *    take() still delivers a subsequently-read message (the regression the sticky-failure latch caused);
+ *  - close() wakes every parked taker with a clean IOException rather than stranding them.
+ *
+ * All deterministic and emulator-free.
+ */
+class MessageQueueFailureContractTest {
+
+    private companion object {
+        const val STREAM = 0
+        const val COMMAND = 0
+        const val TAKER_COUNT = 4
+    }
+
+    /** A queue whose single in-flight read() blocks until [gate] opens, then always fails. */
+    private class FailingQueue : MessageQueue<Int>() {
+        val readerEntered = CountDownLatch(1)
+        val gate = CountDownLatch(1)
+
+        override fun readMessage(): Int {
+            readerEntered.countDown()
+            gate.await()
+            throw IOException("simulated read failure")
+        }
+
+        override fun getLocalId(message: Int) = STREAM
+        override fun getCommand(message: Int) = COMMAND
+        override fun isCloseCommand(message: Int) = false
+    }
+
+    @Test
+    fun `a read failure wakes every parked taker, and each observes an IOException`() {
+        val queue = FailingQueue()
+        queue.startListening(STREAM)
+
+        val outcomes = ConcurrentLinkedQueue<Throwable>()
+        val done = CountDownLatch(TAKER_COUNT)
+        val takers = (0 until TAKER_COUNT).map { i ->
+            thread(name = "taker-$i") {
+                try {
+                    queue.take(STREAM, COMMAND)
+                } catch (t: Throwable) {
+                    outcomes.add(t)
+                } finally {
+                    done.countDown()
+                }
+            }
+        }
+
+        // One taker wins readLock and parks inside read(); the rest park in queueCond.await().
+        assertTrue(queue.readerEntered.await(2, TimeUnit.SECONDS), "no taker entered read()")
+        awaitAllParked(takers)
+
+        // The read fails. Every taker must be woken — the active reader by the throw, the rest by the
+        // signalAll() that now runs on failure; each then retries and fails for itself.
+        queue.gate.countDown()
+
+        val finished = done.await(3, TimeUnit.SECONDS)
+        takers.forEach { it.join(1000) }
+
+        assertTrue(finished, "a read failure stranded one or more parked takers (lost wakeup)")
+        assertEquals(TAKER_COUNT, outcomes.size, "every taker should have failed exactly once")
+        outcomes.forEach { assertIs<IOException>(it, "failures must surface as IOException") }
+    }
+
+    @Test
+    fun `a read failure is not sticky - the queue still delivers a later message`() {
+        val readCalls = AtomicInteger(0)
+        val queue = object : MessageQueue<Int>() {
+            // First read fails transiently; the second succeeds and yields a message for STREAM.
+            override fun readMessage(): Int {
+                if (readCalls.getAndIncrement() == 0) throw IOException("transient read failure")
+                return STREAM
+            }
+
+            override fun getLocalId(message: Int) = STREAM
+            override fun getCommand(message: Int) = COMMAND
+            override fun isCloseCommand(message: Int) = false
+        }
+        queue.startListening(STREAM)
+
+        // The transient failure surfaces to this caller...
+        assertFailsWith<IOException> { queue.take(STREAM, COMMAND) }
+
+        // ...but must NOT poison the queue: the next take() reads and returns the message.
+        val message = queue.take(STREAM, COMMAND)
+        assertEquals(STREAM, message)
+        assertEquals(2, readCalls.get(), "second take() should have re-attempted the read")
+    }
+
+    @Test
+    fun `closing the queue wakes every parked taker with an IOException`() {
+        // A real AdbReader over a source that blocks until closed, then fails — mirrors a wedged socket.
+        val closed = CountDownLatch(1)
+        val blockingSource = object : Source {
+            override fun read(sink: Buffer, byteCount: Long): Long {
+                closed.await()
+                throw IOException("source closed")
+            }
+
+            override fun close() {
+                closed.countDown()
+            }
+
+            override fun timeout(): Timeout = Timeout.NONE
+        }
+        val queue = AdbMessageQueue(AdbReader(blockingSource))
+        queue.startListening(STREAM)
+
+        val outcomes = ConcurrentLinkedQueue<Throwable>()
+        val done = CountDownLatch(TAKER_COUNT)
+        val takers = (0 until TAKER_COUNT).map { i ->
+            thread(name = "adb-taker-$i") {
+                try {
+                    queue.take(STREAM, Constants.CMD_OKAY)
+                } catch (t: Throwable) {
+                    outcomes.add(t)
+                } finally {
+                    done.countDown()
+                }
+            }
+        }
+        awaitAllParked(takers)
+
+        queue.close()
+
+        val finished = done.await(3, TimeUnit.SECONDS)
+        takers.forEach { it.join(1000) }
+
+        assertTrue(finished, "close() failed to wake one or more parked takers")
+        assertEquals(TAKER_COUNT, outcomes.size)
+        outcomes.forEach { assertIs<IOException>(it, "close() must surface as IOException") }
+    }
+
+    private fun awaitAllParked(threads: List<Thread>) {
+        val parked = setOf(Thread.State.WAITING, Thread.State.TIMED_WAITING)
+        val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2)
+        while (System.nanoTime() < deadline) {
+            if (threads.all { it.state in parked }) return
+            Thread.sleep(5)
+        }
+        throw AssertionError("not all takers parked; states=${threads.map { it.state }}")
+    }
+}

dadb/src/test/kotlin/dadb/MessageQueueLostWakeupTest.kt

@@ -0,0 +1,124 @@
+/*
+ * Copyright (c) 2021 mobile.dev inc.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ *
+ */
+
+package dadb
+
+import org.junit.jupiter.api.Assertions.assertFalse
+import org.junit.jupiter.api.Assertions.assertTrue
+import org.junit.jupiter.api.Test
+import java.io.IOException
+import java.util.concurrent.CountDownLatch
+import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicReference
+import kotlin.concurrent.thread
+
+/**
+ * Lost-wakeup hang when two threads share one [MessageQueue] (e.g. maestro tunnels gRPC over the
+ * dadb connection: a keep-alive reader thread + the test-runner thread):
+ *
+ *  1. Reader G wins `readLock` and parks in `read()`.
+ *  2. Consumer T fails `readLock.tryLock()` and parks in `queueCond.await()` (no timeout).
+ *  3. G's `read()` throws (SO_TIMEOUT / reset / EOF), but `signalAll()` ran only after a *successful*
+ *     read, so it is skipped.
+ *  4. T is never signaled and parks forever — neither the read nor the write timeout covers `await()`.
+ *
+ * Pure concurrency bug in `take()`, so it's reproduced directly with a controllable `readMessage()`
+ * — no emulator, fully deterministic.
+ */
+class MessageQueueLostWakeupTest {
+
+    private companion object {
+        const val READER_ID = 1
+        const val CONSUMER_ID = 2
+        const val ANY_COMMAND = 42
+    }
+
+    /** A queue whose single `read()` blocks until [gate] opens and then fails, like a wedged socket. */
+    private class WedgingQueue : MessageQueue<Int>() {
+        val readerEntered = CountDownLatch(1)
+        val gate = CountDownLatch(1)
+
+        override fun readMessage(): Int {
+            readerEntered.countDown()      // G is now inside read(), holding readLock
+            gate.await()                   // ...block here until the test releases the wedge
+            throw IOException("simulated wedged socket")
+        }
+
+        override fun getLocalId(message: Int) = message
+        override fun getCommand(message: Int) = ANY_COMMAND
+        override fun isCloseCommand(message: Int) = false
+    }
+
+    @Test
+    fun `a read failure must wake a thread parked in take(), not strand it forever`() {
+        val queue = WedgingQueue()
+        queue.startListening(READER_ID)
+        queue.startListening(CONSUMER_ID)
+
+        // Thread G: acquires readLock, enters read(), and will throw once the gate opens.
+        val reader = thread(name = "reader-G") {
+            try { queue.take(READER_ID, ANY_COMMAND) } catch (ignore: Throwable) {}
+        }
+        assertTrue(queue.readerEntered.await(2, TimeUnit.SECONDS), "reader G never entered read()")
+
+        // Thread T: fails to grab readLock (G holds it) and parks in queueCond.await().
+        val consumerOutcome = AtomicReference<Throwable?>()
+        val consumerReturned = CountDownLatch(1)
+        val consumer = thread(name = "consumer-T") {
+            try {
+                queue.take(CONSUMER_ID, ANY_COMMAND)
+            } catch (t: Throwable) {
+                consumerOutcome.set(t)
+            } finally {
+                consumerReturned.countDown()
+            }
+        }
+
+        // Wait until T is genuinely parked in await(). queueLock is free, and T uses tryLock (never
+        // blocking) on readLock, so the only thing that parks T is queueCond.await() — whether that
+        // await is untimed (WAITING) or bounded (TIMED_WAITING) is an implementation detail.
+        awaitParked(consumer)
+
+        // The wedge "times out": G's read() throws. In the buggy code signalAll() is skipped, so
+        // nothing wakes T even though readLock is now free.
+        queue.gate.countDown()
+
+        val consumerFinished = consumerReturned.await(3, TimeUnit.SECONDS)
+
+        // Clean up the parked thread so a failing run doesn't leak it.
+        if (!consumerFinished) consumer.interrupt()
+        reader.join(1000)
+        consumer.join(1000)
+
+        assertTrue(
+            consumerFinished,
+            "Consumer thread T was never woken after the reader's read() failed — it is stranded in " +
+                "queueCond.await() (lost wakeup). Neither the read nor the write timeout covers this."
+        )
+        assertFalse(consumer.isAlive, "consumer thread should have terminated")
+    }
+
+    private fun awaitParked(t: Thread) {
+        val parked = setOf(Thread.State.WAITING, Thread.State.TIMED_WAITING)
+        val deadline = System.nanoTime() + TimeUnit.SECONDS.toNanos(2)
+        while (System.nanoTime() < deadline) {
+            if (t.state in parked) return
+            Thread.sleep(5)
+        }
+        throw AssertionError("thread ${t.name} never parked in await(); state=${t.state}")
+    }
+}