Add receipt processing benchmark tests.

app/src/benchmarkShared/java/org/signal/benchmark/BenchmarkCommandReceiver.kt

@@ -13,8 +13,13 @@ import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
 import org.signal.benchmark.setup.Generator
 import org.signal.benchmark.setup.Harness
+import org.signal.benchmark.setup.OtherClient
 import org.signal.core.util.ThreadUtil
 import org.signal.core.util.logging.Log
+import org.thoughtcrime.securesms.database.SignalDatabase
+import org.thoughtcrime.securesms.database.TestDbUtils
+import org.thoughtcrime.securesms.groups.GroupId
+import org.thoughtcrime.securesms.recipients.Recipient
 import org.whispersystems.signalservice.internal.push.Envelope
 import org.whispersystems.signalservice.internal.websocket.BenchmarkWebSocketConnection
 import org.whispersystems.signalservice.internal.websocket.WebSocketRequestMessage
@@ -45,6 +50,8 @@ class BenchmarkCommandReceiver : BroadcastReceiver() {
     when (command) {
       "individual-send" -> handlePrepareIndividualSend()
       "group-send" -> handlePrepareGroupSend()
+      "group-delivery-receipt" -> handlePrepareGroupReceipts { client, timestamps -> client.generateInboundDeliveryReceipts(timestamps) }
+      "group-read-receipt" -> handlePrepareGroupReceipts { client, timestamps -> client.generateInboundReadReceipts(timestamps) }
       "release-messages" -> {
         BenchmarkWebSocketConnection.authInstance.startWholeBatchTrace = true
         BenchmarkWebSocketConnection.authInstance.releaseMessages()
@@ -113,6 +120,62 @@ class BenchmarkCommandReceiver : BroadcastReceiver() {
     BenchmarkWebSocketConnection.authInstance.addQueueEmptyMessage()
   }
 
+  private fun handlePrepareGroupReceipts(generateReceipts: (OtherClient, List<Long>) -> List<Envelope>) {
+    val clients = Harness.otherClients.take(5)
+
+    establishGroupSessions(clients)
+
+    val timestamps = getOutgoingGroupMessageTimestamps()
+    Log.i(TAG, "Found ${timestamps.size} outgoing message timestamps for receipts")
+
+    val allClientEnvelopes = clients.map { client ->
+      generateReceipts(client, timestamps).map { it.toWebSocketPayload() }
+    }
+
+    BenchmarkWebSocketConnection.authInstance.addPendingMessages(interleave(allClientEnvelopes))
+    BenchmarkWebSocketConnection.authInstance.addQueueEmptyMessage()
+  }
+
+  private fun establishGroupSessions(clients: List<OtherClient>) {
+    val encryptedEnvelopes = clients.map { it.encrypt(Generator.encryptedTextMessage(System.currentTimeMillis(), groupMasterKey = Harness.groupMasterKey)) }
+
+    runBlocking {
+      launch(Dispatchers.IO) {
+        BenchmarkWebSocketConnection.authInstance.run {
+          Log.i(TAG, "Sending initial group messages from clients to establish sessions.")
+          addPendingMessages(encryptedEnvelopes.map { it.toWebSocketPayload() })
+          releaseMessages()
+          ThreadUtil.sleep(1000)
+        }
+      }
+    }
+  }
+
+  private fun getOutgoingGroupMessageTimestamps(): List<Long> {
+    val groupId = GroupId.v2(Harness.groupMasterKey)
+    val groupRecipient = Recipient.externalGroupExact(groupId)
+    val threadId = SignalDatabase.threads.getOrCreateThreadIdFor(groupRecipient)
+    val selfId = Recipient.self().id.toLong()
+    return TestDbUtils.getOutgoingMessageTimestamps(threadId, selfId)
+  }
+
+  /**
+   * Interleaves lists so that items from different lists alternate:
+   * [[a1, a2], [b1, b2], [c1, c2]] -> [a1, b1, c1, a2, b2, c2]
+   */
+  private fun <T> interleave(lists: List<List<T>>): List<T> {
+    val result = mutableListOf<T>()
+    val maxSize = lists.maxOf { it.size }
+    for (i in 0 until maxSize) {
+      for (list in lists) {
+        if (i < list.size) {
+          result += list[i]
+        }
+      }
+    }
+    return result
+  }
+
   private fun Envelope.toWebSocketPayload(): WebSocketRequestMessage {
     return WebSocketRequestMessage(
       verb = "PUT",

app/src/benchmarkShared/java/org/signal/benchmark/BenchmarkSetupActivity.kt

@@ -6,7 +6,10 @@ import org.signal.benchmark.setup.TestMessages
 import org.signal.benchmark.setup.TestUsers
 import org.thoughtcrime.securesms.BaseActivity
 import org.thoughtcrime.securesms.database.SignalDatabase
+import org.thoughtcrime.securesms.database.TestDbUtils
+import org.thoughtcrime.securesms.mms.OutgoingMessage
 import org.thoughtcrime.securesms.recipients.Recipient
+import org.thoughtcrime.securesms.util.TextSecurePreferences
 
 class BenchmarkSetupActivity : BaseActivity() {
   override fun onCreate(savedInstanceState: Bundle?) {
@@ -17,6 +20,8 @@ class BenchmarkSetupActivity : BaseActivity() {
       "conversation-open" -> setupConversationOpen()
       "message-send" -> setupMessageSend()
       "group-message-send" -> setupGroupMessageSend()
+      "group-delivery-receipt" -> setupGroupReceipt(includeMsl = true)
+      "group-read-receipt" -> setupGroupReceipt(enableReadReceipts = true)
     }
 
     val textView: TextView = TextView(this).apply {
@@ -68,4 +73,40 @@ class BenchmarkSetupActivity : BaseActivity() {
     TestUsers.setupSelf()
     TestUsers.setupGroup()
   }
+
+  private fun setupGroupReceipt(includeMsl: Boolean = false, enableReadReceipts: Boolean = false) {
+    TestUsers.setupSelf()
+    val groupId = TestUsers.setupGroup()
+
+    val groupRecipient = Recipient.externalGroupExact(groupId)
+    val threadId = SignalDatabase.threads.getOrCreateThreadIdFor(groupRecipient)
+
+    val messageIds = mutableListOf<Long>()
+    val timestamps = mutableListOf<Long>()
+    val baseTimestamp = 2_000_000L
+
+    for (i in 0 until 100) {
+      val timestamp = baseTimestamp + i
+      val message = OutgoingMessage(
+        recipient = groupRecipient,
+        body = "Outgoing message $i",
+        timestamp = timestamp,
+        isSecure = true
+      )
+      val insert = SignalDatabase.messages.insertMessageOutbox(message, threadId, false, null)
+      SignalDatabase.messages.markAsSent(insert.messageId, true)
+      messageIds += insert.messageId
+      timestamps += timestamp
+    }
+
+    if (includeMsl) {
+      val selfId = Recipient.self().id
+      val memberRecipientIds = SignalDatabase.groups.getGroup(groupId).get().members.filter { it != selfId }
+      TestDbUtils.insertMessageSendLogEntries(messageIds, timestamps, memberRecipientIds)
+    }
+
+    if (enableReadReceipts) {
+      TextSecurePreferences.setReadReceiptsEnabled(this, true)
+    }
+  }
 }

app/src/benchmarkShared/java/org/signal/benchmark/setup/Generator.kt

@@ -18,6 +18,7 @@ import org.whispersystems.signalservice.internal.push.DataMessage
 import org.whispersystems.signalservice.internal.push.Envelope
 import org.whispersystems.signalservice.internal.push.GroupContextV2
 import org.whispersystems.signalservice.internal.push.OutgoingPushMessage
+import org.whispersystems.signalservice.internal.push.ReceiptMessage
 import java.util.Optional
 import java.util.UUID
 
@@ -45,6 +46,26 @@ object Generator {
     return EnvelopeContent.encrypted(content.build(), ContentHint.RESENDABLE, Optional.empty())
   }
 
+  fun encryptedDeliveryReceipt(now: Long, timestamps: List<Long>): EnvelopeContent {
+    return encryptedReceipt(ReceiptMessage.Type.DELIVERY, timestamps)
+  }
+
+  fun encryptedReadReceipt(now: Long, timestamps: List<Long>): EnvelopeContent {
+    return encryptedReceipt(ReceiptMessage.Type.READ, timestamps)
+  }
+
+  private fun encryptedReceipt(type: ReceiptMessage.Type, timestamps: List<Long>): EnvelopeContent {
+    val content = Content.Builder().apply {
+      receiptMessage(
+        ReceiptMessage.Builder().buildWith {
+          this.type = type
+          timestamp = timestamps
+        }
+      )
+    }
+    return EnvelopeContent.encrypted(content.build(), ContentHint.IMPLICIT, Optional.empty())
+  }
+
   fun OutgoingPushMessage.toEnvelope(timestamp: Long, destination: ServiceId): Envelope {
     val serverGuid = UUID.randomUUID()
     return Envelope.Builder()

app/src/benchmarkShared/java/org/signal/benchmark/setup/OtherClient.kt

@@ -62,6 +62,10 @@ class OtherClient(val serviceId: ServiceId, val e164: String, val identityKeyPai
 
   /** Inspired by SignalServiceMessageSender#getEncryptedMessage */
   fun encrypt(envelopeContent: EnvelopeContent): Envelope {
+    return encrypt(envelopeContent, envelopeContent.content.get().dataMessage!!.timestamp!!)
+  }
+
+  fun encrypt(envelopeContent: EnvelopeContent, timestamp: Long): Envelope {
     val cipher = SignalServiceCipher(serviceAddress, 1, aciStore, sessionLock, null)
 
     if (!aciStore.containsSession(getAliceProtocolAddress())) {
@@ -70,7 +74,7 @@ class OtherClient(val serviceId: ServiceId, val e164: String, val identityKeyPai
     }
 
     return cipher.encrypt(getAliceProtocolAddress(), getAliceUnidentifiedAccess(), envelopeContent)
-      .toEnvelope(envelopeContent.content.get().dataMessage!!.timestamp!!, getAliceServiceId())
+      .toEnvelope(timestamp, getAliceServiceId())
   }
 
   fun generateInboundEnvelopes(count: Int): List<Envelope> {
@@ -84,6 +88,24 @@ class OtherClient(val serviceId: ServiceId, val e164: String, val identityKeyPai
     return envelopes
   }
 
+  fun generateInboundDeliveryReceipts(messageTimestamps: List<Long>): List<Envelope> {
+    return generateInboundReceipts(messageTimestamps, Generator::encryptedDeliveryReceipt)
+  }
+
+  fun generateInboundReadReceipts(messageTimestamps: List<Long>): List<Envelope> {
+    return generateInboundReceipts(messageTimestamps, Generator::encryptedReadReceipt)
+  }
+
+  private fun generateInboundReceipts(messageTimestamps: List<Long>, receiptFactory: (Long, List<Long>) -> EnvelopeContent): List<Envelope> {
+    val envelopes = ArrayList<Envelope>(messageTimestamps.size)
+    var now = System.currentTimeMillis()
+    for (messageTimestamp in messageTimestamps) {
+      envelopes += encrypt(receiptFactory(now, listOf(messageTimestamp)), now)
+      now += 3
+    }
+    return envelopes
+  }
+
   fun generateInboundGroupEnvelopes(count: Int, groupMasterKey: GroupMasterKey): List<Envelope> {
     val envelopes = ArrayList<Envelope>(count)
     var now = System.currentTimeMillis()

app/src/benchmarkShared/java/org/signal/benchmark/setup/TestUsers.kt

@@ -21,6 +21,7 @@ import org.thoughtcrime.securesms.crypto.ProfileKeyUtil
 import org.thoughtcrime.securesms.database.SignalDatabase
 import org.thoughtcrime.securesms.database.model.databaseprotos.RestoreDecisionState
 import org.thoughtcrime.securesms.dependencies.AppDependencies
+import org.thoughtcrime.securesms.groups.GroupId
 import org.thoughtcrime.securesms.keyvalue.CertificateType
 import org.thoughtcrime.securesms.keyvalue.SignalStore
 import org.thoughtcrime.securesms.keyvalue.Skipped
@@ -170,7 +171,7 @@ object TestUsers {
     return others
   }
 
-  fun setupGroup() {
+  fun setupGroup(): GroupId.V2 {
     val members = setupTestClients(5)
     val self = Recipient.self()
 
@@ -202,6 +203,8 @@ object TestUsers {
     )
 
     SignalDatabase.recipients.setProfileSharing(Recipient.externalGroupExact(groupId!!).id, true)
+
+    return groupId
   }
 
   private fun member(aci: ACI, role: Member.Role = Member.Role.DEFAULT, joinedAt: Int = 0, labelEmoji: String = "", labelString: String = ""): DecryptedMember {

app/src/benchmarkShared/java/org/thoughtcrime/securesms/database/TestDbUtils.kt

@@ -2,6 +2,8 @@ package org.thoughtcrime.securesms.database
 
 import android.content.ContentValues
 import org.signal.core.util.SqlUtil.buildArgs
+import org.thoughtcrime.securesms.recipients.RecipientId
+import org.whispersystems.signalservice.internal.push.Content
 
 object TestDbUtils {
 
@@ -11,4 +13,58 @@ object TestDbUtils {
     contentValues.put(MessageTable.DATE_RECEIVED, timestamp)
     val rowsUpdated = database.update(MessageTable.TABLE_NAME, contentValues, DatabaseTable.ID_WHERE, buildArgs(messageId))
   }
+
+  fun getOutgoingMessageTimestamps(threadId: Long, selfRecipientId: Long): List<Long> {
+    val timestamps = mutableListOf<Long>()
+    SignalDatabase.messages.databaseHelper.signalReadableDatabase.query(
+      MessageTable.TABLE_NAME,
+      arrayOf(MessageTable.DATE_SENT),
+      "${MessageTable.THREAD_ID} = ? AND ${MessageTable.FROM_RECIPIENT_ID} = ?",
+      arrayOf(threadId.toString(), selfRecipientId.toString()),
+      null,
+      null,
+      "${MessageTable.DATE_SENT} ASC"
+    ).use { cursor ->
+      while (cursor.moveToNext()) {
+        timestamps += cursor.getLong(0)
+      }
+    }
+    return timestamps
+  }
+
+  fun insertMessageSendLogEntries(messageIds: List<Long>, timestamps: List<Long>, recipientIds: List<RecipientId>) {
+    val db = SignalDatabase.messages.databaseHelper.signalWritableDatabase
+    val dummyContent = Content.Builder().build().encode()
+
+    db.beginTransaction()
+    try {
+      for (i in messageIds.indices) {
+        val payloadValues = ContentValues().apply {
+          put("date_sent", timestamps[i])
+          put("content", dummyContent)
+          put("content_hint", 0)
+          put("urgent", 1)
+        }
+        val payloadId = db.insert("msl_payload", null, payloadValues)
+
+        val messageValues = ContentValues().apply {
+          put("payload_id", payloadId)
+          put("message_id", messageIds[i])
+        }
+        db.insert("msl_message", null, messageValues)
+
+        for (recipientId in recipientIds) {
+          val recipientValues = ContentValues().apply {
+            put("payload_id", payloadId)
+            put("recipient_id", recipientId.toLong())
+            put("device", 1)
+          }
+          db.insert("msl_recipient", null, recipientValues)
+        }
+      }
+      db.setTransactionSuccessful()
+    } finally {
+      db.endTransaction()
+    }
+  }
 }