Add tracing implementations of listeners

android-test-app/src/main/kotlin/okhttp/android/testapp/MainActivity.kt

@@ -19,18 +19,25 @@ import android.os.Bundle
 import androidx.activity.ComponentActivity
 import okhttp3.Call
 import okhttp3.Callback
+import okhttp3.ConnectionPool
 import okhttp3.HttpUrl.Companion.toHttpUrl
 import okhttp3.OkHttpClient
 import okhttp3.Request
 import okhttp3.Response
+import okhttp3.android.AndroidxTracingConnectionListener
+import okhttp3.android.AndroidxTracingInterceptor
 import okhttp3.internal.platform.AndroidPlatform
 import okio.IOException
 
 class MainActivity : ComponentActivity() {
   override fun onCreate(savedInstanceState: Bundle?) {
     super.onCreate(savedInstanceState)
 
-    val client = OkHttpClient()
+    val client =
+      OkHttpClient.Builder()
+        .connectionPool(ConnectionPool(connectionListener = AndroidxTracingConnectionListener()))
+        .addNetworkInterceptor(AndroidxTracingInterceptor())
+        .build()
 
     // Ensure we are compiling against the right variant
     println(AndroidPlatform.isSupported)

gradle/libs.versions.toml

@@ -1,4 +1,5 @@
 [versions]
+androidx-tracing = "1.1.0"
 # 7.0.0 is JDK 17+ https://github.com/bndtools/bnd/wiki/Changes-in-7.0.0
 biz-aQute-bnd = "6.4.0"
 checkStyle = "10.21.1"
@@ -28,6 +29,7 @@ androidx-junit = "androidx.test.ext:junit:1.2.1"
 androidx-lint-gradle = { module = "androidx.lint:lint-gradle", version.ref = "lintGradle" }
 androidx-startup-runtime = { module = "androidx.startup:startup-runtime", version.ref = "startupRuntime" }
 androidx-test-runner = "androidx.test:runner:1.6.2"
+androidx-tracing-ktx = { module = "androidx.tracing:tracing-ktx", version.ref = "androidx-tracing" }
 animalsniffer-annotations = "org.codehaus.mojo:animal-sniffer-annotations:1.24"
 aqute-resolve = { module = "biz.aQute.bnd:biz.aQute.resolve", version.ref = "biz-aQute-bnd" }
 assertk = "com.willowtreeapps.assertk:assertk:0.28.1"

okhttp-testing-support/src/main/kotlin/okhttp3/RecordingConnectionListener.kt

@@ -137,11 +137,13 @@ open class RecordingConnectionListener(
   }
 
   override fun connectStart(
+    connectionId: Long,
     route: Route,
     call: Call,
   ) = logEvent(ConnectionEvent.ConnectStart(System.nanoTime(), route, call))
 
   override fun connectFailed(
+    connectionId: Long,
     route: Route,
     call: Call,
     failure: IOException,

okhttp/api/android/okhttp.api

@@ -354,6 +354,7 @@ public final class okhttp3/CipherSuite$Companion {
 }
 
 public abstract interface class okhttp3/Connection {
+	public fun getId ()J
 	public abstract fun handshake ()Lokhttp3/Handshake;
 	public abstract fun protocol ()Lokhttp3/Protocol;
 	public abstract fun route ()Lokhttp3/Route;
@@ -364,8 +365,8 @@ public abstract class okhttp3/ConnectionListener {
 	public static final field Companion Lokhttp3/ConnectionListener$Companion;
 	public fun <init> ()V
 	public fun connectEnd (Lokhttp3/Connection;Lokhttp3/Route;Lokhttp3/Call;)V
-	public fun connectFailed (Lokhttp3/Route;Lokhttp3/Call;Ljava/io/IOException;)V
-	public fun connectStart (Lokhttp3/Route;Lokhttp3/Call;)V
+	public fun connectFailed (JLokhttp3/Route;Lokhttp3/Call;Ljava/io/IOException;)V
+	public fun connectStart (JLokhttp3/Route;Lokhttp3/Call;)V
 	public fun connectionAcquired (Lokhttp3/Connection;Lokhttp3/Call;)V
 	public fun connectionClosed (Lokhttp3/Connection;)V
 	public fun connectionReleased (Lokhttp3/Connection;Lokhttp3/Call;)V
@@ -1304,3 +1305,35 @@ public final class okhttp3/android/AndroidAsyncDns$Companion {
 	public final fun getIPv6 ()Lokhttp3/android/AndroidAsyncDns;
 }
 
+public final class okhttp3/android/AndroidxTracingConnectionListener : okhttp3/ConnectionListener {
+	public static final field Companion Lokhttp3/android/AndroidxTracingConnectionListener$Companion;
+	public fun <init> ()V
+	public fun <init> (Lokhttp3/ConnectionListener;Lkotlin/jvm/functions/Function1;)V
+	public synthetic fun <init> (Lokhttp3/ConnectionListener;Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public fun connectEnd (Lokhttp3/Connection;Lokhttp3/Route;Lokhttp3/Call;)V
+	public fun connectFailed (JLokhttp3/Route;Lokhttp3/Call;Ljava/io/IOException;)V
+	public fun connectStart (JLokhttp3/Route;Lokhttp3/Call;)V
+	public fun connectionAcquired (Lokhttp3/Connection;Lokhttp3/Call;)V
+	public fun connectionClosed (Lokhttp3/Connection;)V
+	public fun connectionReleased (Lokhttp3/Connection;Lokhttp3/Call;)V
+	public final fun getTraceLabel ()Lkotlin/jvm/functions/Function1;
+	public fun noNewExchanges (Lokhttp3/Connection;)V
+}
+
+public final class okhttp3/android/AndroidxTracingConnectionListener$Companion {
+	public final fun getDefaultTracingLabel (Lokhttp3/Route;)Ljava/lang/String;
+}
+
+public final class okhttp3/android/AndroidxTracingInterceptor : okhttp3/Interceptor {
+	public static final field Companion Lokhttp3/android/AndroidxTracingInterceptor$Companion;
+	public fun <init> ()V
+	public fun <init> (Lkotlin/jvm/functions/Function1;)V
+	public synthetic fun <init> (Lkotlin/jvm/functions/Function1;ILkotlin/jvm/internal/DefaultConstructorMarker;)V
+	public final fun getTraceLabel ()Lkotlin/jvm/functions/Function1;
+	public fun intercept (Lokhttp3/Interceptor$Chain;)Lokhttp3/Response;
+}
+
+public final class okhttp3/android/AndroidxTracingInterceptor$Companion {
+	public final fun getDefaultTracingLabel (Lokhttp3/Request;)Ljava/lang/String;
+}
+

okhttp/api/jvm/okhttp.api

@@ -354,6 +354,7 @@ public final class okhttp3/CipherSuite$Companion {
 }
 
 public abstract interface class okhttp3/Connection {
+	public fun getId ()J
 	public abstract fun handshake ()Lokhttp3/Handshake;
 	public abstract fun protocol ()Lokhttp3/Protocol;
 	public abstract fun route ()Lokhttp3/Route;
@@ -364,8 +365,8 @@ public abstract class okhttp3/ConnectionListener {
 	public static final field Companion Lokhttp3/ConnectionListener$Companion;
 	public fun <init> ()V
 	public fun connectEnd (Lokhttp3/Connection;Lokhttp3/Route;Lokhttp3/Call;)V
-	public fun connectFailed (Lokhttp3/Route;Lokhttp3/Call;Ljava/io/IOException;)V
-	public fun connectStart (Lokhttp3/Route;Lokhttp3/Call;)V
+	public fun connectFailed (JLokhttp3/Route;Lokhttp3/Call;Ljava/io/IOException;)V
+	public fun connectStart (JLokhttp3/Route;Lokhttp3/Call;)V
 	public fun connectionAcquired (Lokhttp3/Connection;Lokhttp3/Call;)V
 	public fun connectionClosed (Lokhttp3/Connection;)V
 	public fun connectionReleased (Lokhttp3/Connection;Lokhttp3/Call;)V

okhttp/build.gradle.kts

@@ -97,6 +97,7 @@ kotlin {
         compileOnly(libs.conscrypt.openjdk)
         implementation(libs.androidx.annotation)
         implementation(libs.androidx.startup.runtime)
+        implementation(libs.androidx.tracing.ktx)
       }
     }
 

okhttp/src/androidMain/kotlin/okhttp3/android/AndroidxTracingConnectionListener.kt

@@ -0,0 +1,90 @@
+/*
+ * Copyright (C) 2024 Block, 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 okhttp3.android
+
+import androidx.tracing.Trace
+import okhttp3.Call
+import okhttp3.Connection
+import okhttp3.ConnectionListener
+import okhttp3.Route
+import okhttp3.android.AndroidxTracingInterceptor.Companion.MAX_TRACE_LABEL_LENGTH
+import okio.IOException
+
+/**
+ * Tracing implementation of ConnectionListener that marks the lifetime of each connection
+ * in Perfetto traces.
+ */
+class AndroidxTracingConnectionListener(
+  private val delegate: ConnectionListener = NONE,
+  val traceLabel: (Route) -> String = { it.defaultTracingLabel },
+) : ConnectionListener() {
+  override fun connectStart(
+    connectionId: Long,
+    route: Route,
+    call: Call,
+  ) {
+    Trace.beginAsyncSection(labelForTrace(route), connectionId.toInt())
+    delegate.connectStart(connectionId, route, call)
+  }
+
+  override fun connectFailed(
+    connectionId: Long,
+    route: Route,
+    call: Call,
+    failure: IOException,
+  ) {
+    Trace.endAsyncSection(labelForTrace(route), connectionId.toInt())
+    delegate.connectFailed(connectionId, route, call, failure)
+  }
+
+  override fun connectEnd(
+    connection: Connection,
+    route: Route,
+    call: Call,
+  ) {
+    delegate.connectEnd(connection, route, call)
+  }
+
+  override fun connectionClosed(connection: Connection) {
+    Trace.endAsyncSection(labelForTrace(connection.route()), connection.id.toInt())
+    delegate.connectionClosed(connection)
+  }
+
+  private fun labelForTrace(route: Route): String = traceLabel(route).take(MAX_TRACE_LABEL_LENGTH)
+
+  override fun connectionAcquired(
+    connection: Connection,
+    call: Call,
+  ) {
+    delegate.connectionAcquired(connection, call)
+  }
+
+  override fun connectionReleased(
+    connection: Connection,
+    call: Call,
+  ) {
+    delegate.connectionReleased(connection, call)
+  }
+
+  override fun noNewExchanges(connection: Connection) {
+    delegate.noNewExchanges(connection)
+  }
+
+  companion object {
+    val Route.defaultTracingLabel: String
+      get() = this.address.url.host
+  }
+}

okhttp/src/androidMain/kotlin/okhttp3/android/AndroidxTracingInterceptor.kt

@@ -0,0 +1,43 @@
+/*
+ * Copyright (c) 2022 Block, 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 okhttp3.android
+
+import androidx.tracing.trace
+import okhttp3.Interceptor
+import okhttp3.Request
+import okhttp3.Response
+
+/**
+ * Tracing implementation of Interceptor that marks each Call in a Perfetto
+ * trace. Typically used as a network interceptor.
+ */
+class AndroidxTracingInterceptor(val traceLabel: (Request) -> String = { it.defaultTracingLabel }) : Interceptor {
+  override fun intercept(chain: Interceptor.Chain): Response {
+    return trace(traceLabel(chain.request()).take(MAX_TRACE_LABEL_LENGTH)) {
+      chain.proceed(chain.request())
+    }
+  }
+
+  companion object {
+    internal const val MAX_TRACE_LABEL_LENGTH = 127
+
+    val Request.defaultTracingLabel: String
+      get() {
+        return url.encodedPath
+      }
+  }
+}

okhttp/src/commonJvmAndroid/kotlin/okhttp3/Connection.kt

@@ -67,6 +67,10 @@ import java.net.Socket
  * been found. But only complete the stream once its data stream has been exhausted.
  */
 interface Connection {
+  /** Unique id of this connection, assigned at the time of the attempt. */
+  val id: Long
+    get() = 0L
+
   /** Returns the route used by this connection. */
   fun route(): Route
 

okhttp/src/commonJvmAndroid/kotlin/okhttp3/ConnectionListener.kt

@@ -30,6 +30,7 @@ abstract class ConnectionListener {
    * Invoked as soon as a call causes a connection to be started.
    */
   open fun connectStart(
+    connectionId: Long,
     route: Route,
     call: Call,
   ) {}
@@ -38,6 +39,7 @@ abstract class ConnectionListener {
    * Invoked when a connection fails to be established.
    */
   open fun connectFailed(
+    connectionId: Long,
     route: Route,
     call: Call,
     failure: IOException,

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/connection/CallConnectionUser.kt

@@ -36,18 +36,22 @@ internal class CallConnectionUser(
     call.client.routeDatabase.connected(route)
   }
 
-  override fun connectStart(route: Route) {
+  override fun connectStart(
+    connectionId: Long,
+    route: Route,
+  ) {
     eventListener.connectStart(call, route.socketAddress, route.proxy)
-    poolConnectionListener.connectStart(route, call)
+    poolConnectionListener.connectStart(connectionId, route, call)
   }
 
   override fun connectFailed(
+    connectionId: Long,
     route: Route,
     protocol: Protocol?,
     e: IOException,
   ) {
     eventListener.connectFailed(call, route.socketAddress, route.proxy, null, e)
-    poolConnectionListener.connectFailed(route, call, e)
+    poolConnectionListener.connectFailed(connectionId, route, call, e)
   }
 
   override fun secureConnectStart() {

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/connection/ConnectPlan.kt

@@ -24,6 +24,7 @@ import java.net.Socket
 import java.net.UnknownServiceException
 import java.security.cert.X509Certificate
 import java.util.concurrent.TimeUnit
+import java.util.concurrent.atomic.AtomicLong
 import javax.net.ssl.SSLPeerUnverifiedException
 import javax.net.ssl.SSLSocket
 import okhttp3.CertificatePinner
@@ -78,6 +79,8 @@ class ConnectPlan(
   internal val connectionSpecIndex: Int,
   internal val isTlsFallback: Boolean,
 ) : RoutePlanner.Plan, ExchangeCodec.Carrier {
+  private val id = idGenerator.incrementAndGet()
+
   /** True if this connect was canceled; typically because it lost a race. */
   @Volatile private var canceled = false
 
@@ -135,7 +138,7 @@ class ConnectPlan(
     // Tell the call about the connecting call so async cancels work.
     user.addPlanToCancel(this)
     try {
-      user.connectStart(route)
+      user.connectStart(id, route)
 
       connectSocket()
       success = true
@@ -149,7 +152,7 @@ class ConnectPlan(
           e,
         )
       }
-      user.connectFailed(route, null, e)
+      user.connectFailed(id, route, null, e)
       return ConnectResult(plan = this, throwable = e)
     } finally {
       user.removePlanToCancel(this)
@@ -231,6 +234,7 @@ class ConnectPlan(
           sink = sink,
           pingIntervalMillis = pingIntervalMillis,
           connectionListener = connectionPool.connectionListener,
+          id = id,
         )
       this.connection = connection
       connection.start()
@@ -240,7 +244,7 @@ class ConnectPlan(
       success = true
       return ConnectResult(plan = this)
     } catch (e: IOException) {
-      user.connectFailed(route, null, e)
+      user.connectFailed(id, route, null, e)
 
       if (!retryOnConnectionFailure || !retryTlsHandshake(e)) {
         retryTlsConnection = null
@@ -333,7 +337,7 @@ class ConnectPlan(
           ProtocolException(
             "Too many tunnel connections attempted: $MAX_TUNNEL_ATTEMPTS",
           )
-        user.connectFailed(route, null, failure)
+        user.connectFailed(id, route, null, failure)
         return ConnectResult(plan = this, throwable = failure)
       }
     }
@@ -565,5 +569,7 @@ class ConnectPlan(
   companion object {
     private const val NPE_THROW_WITH_NULL = "throw with null exception"
     private const val MAX_TUNNEL_ATTEMPTS = 21
+
+    private val idGenerator = AtomicLong(0)
   }
 }