New Response.peekTrailers() API (#8921)

* New Response.peekTrailers() API

This restores behavior we had in OkHttp 4.x that I
broke when I changed Response.trailers() to block
until trailers were available.

Closes: #8916

* apiDump

---------

Co-authored-by: Jesse Wilson <[email protected]>

Author: swankjesse

okhttp/api/android/okhttp.api

@@ -1114,6 +1114,7 @@ public final class okhttp3/Response : java/io/Closeable {
 	public final fun networkResponse ()Lokhttp3/Response;
 	public final fun newBuilder ()Lokhttp3/Response$Builder;
 	public final fun peekBody (J)Lokhttp3/ResponseBody;
+	public final fun peekTrailers ()Lokhttp3/Headers;
 	public final fun priorResponse ()Lokhttp3/Response;
 	public final fun protocol ()Lokhttp3/Protocol;
 	public final fun receivedResponseAtMillis ()J
@@ -1219,11 +1220,16 @@ public abstract interface class okhttp3/TrailersSource {
 	public static final field Companion Lokhttp3/TrailersSource$Companion;
 	public static final field EMPTY Lokhttp3/TrailersSource;
 	public abstract fun get ()Lokhttp3/Headers;
+	public fun peek ()Lokhttp3/Headers;
 }
 
 public final class okhttp3/TrailersSource$Companion {
 }
 
+public final class okhttp3/TrailersSource$DefaultImpls {
+	public static fun peek (Lokhttp3/TrailersSource;)Lokhttp3/Headers;
+}
+
 public abstract interface class okhttp3/WebSocket {
 	public abstract fun cancel ()V
 	public abstract fun close (ILjava/lang/String;)Z

okhttp/api/jvm/okhttp.api

@@ -1114,6 +1114,7 @@ public final class okhttp3/Response : java/io/Closeable {
 	public final fun networkResponse ()Lokhttp3/Response;
 	public final fun newBuilder ()Lokhttp3/Response$Builder;
 	public final fun peekBody (J)Lokhttp3/ResponseBody;
+	public final fun peekTrailers ()Lokhttp3/Headers;
 	public final fun priorResponse ()Lokhttp3/Response;
 	public final fun protocol ()Lokhttp3/Protocol;
 	public final fun receivedResponseAtMillis ()J
@@ -1219,11 +1220,16 @@ public abstract interface class okhttp3/TrailersSource {
 	public static final field Companion Lokhttp3/TrailersSource$Companion;
 	public static final field EMPTY Lokhttp3/TrailersSource;
 	public abstract fun get ()Lokhttp3/Headers;
+	public fun peek ()Lokhttp3/Headers;
 }
 
 public final class okhttp3/TrailersSource$Companion {
 }
 
+public final class okhttp3/TrailersSource$DefaultImpls {
+	public static fun peek (Lokhttp3/TrailersSource;)Lokhttp3/Headers;
+}
+
 public abstract interface class okhttp3/WebSocket {
 	public abstract fun cancel ()V
 	public abstract fun close (ILjava/lang/String;)Z

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

@@ -192,6 +192,26 @@ class Response internal constructor(
   @Throws(IOException::class)
   fun trailers(): Headers = trailersSource.get()
 
+  /**
+   * Returns the trailers after the HTTP response, if they are available to read immediately. Unlike
+   * [trailers], this doesn't block if the trailers are not immediately available, and instead
+   * returns null.
+   *
+   * This will typically return null until [ResponseBody.source] has buffered the last byte of the
+   * response body. Call `body.source().request(1024 * 1024)` to block until either that's done, or
+   * 1 MiB of response data is loaded into memory. (You could use any size here, though large values
+   * risk exhausting memory.)
+   *
+   * This returns an empty value if the trailers are available, but have no data.
+   *
+   * It is not safe to call this concurrently with code that is processing the response body.
+   *
+   * @throws IOException if the trailers cannot be loaded, such as if the network connection is
+   *     dropped.
+   */
+  @Throws(IOException::class)
+  fun peekTrailers(): Headers? = trailersSource.peek()
+
   /**
    * Peeks up to [byteCount] bytes from the response body and returns them as a new response
    * body. If fewer than [byteCount] bytes are in the response body, the full response body is

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

@@ -27,12 +27,20 @@ import okio.IOException
  * This interface is for test and production code that creates [Response] instances without making
  * an HTTP call to a remote server.
  */
-fun interface TrailersSource {
+interface TrailersSource {
+  @Throws(IOException::class)
+  fun peek(): Headers? = null
+
   @Throws(IOException::class)
   fun get(): Headers
 
   companion object {
     @JvmField
-    val EMPTY: TrailersSource = TrailersSource { Headers.EMPTY }
+    val EMPTY: TrailersSource =
+      object : TrailersSource {
+        override fun peek() = Headers.EMPTY
+
+        override fun get() = Headers.EMPTY
+      }
   }
 }

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

@@ -141,7 +141,7 @@ class Exchange(
   }
 
   @Throws(IOException::class)
-  fun trailers(): Headers = codec.trailers()
+  fun peekTrailers(): Headers? = codec.peekTrailers()
 
   @Throws(SocketException::class)
   fun newWebSocketStreams(): RealWebSocket.Streams {

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http/CallServerInterceptor.kt

@@ -17,8 +17,10 @@ package okhttp3.internal.http
 
 import java.io.IOException
 import java.net.ProtocolException
+import okhttp3.Headers
 import okhttp3.Interceptor
 import okhttp3.Response
+import okhttp3.TrailersSource
 import okhttp3.internal.connection.Exchange
 import okhttp3.internal.http2.ConnectionShutdownException
 import okhttp3.internal.skipAll
@@ -134,13 +136,19 @@ class CallServerInterceptor(
           response
             .newBuilder()
             .body(responseBody)
-            .trailers {
-              val source = responseBody.source()
-              if (source.isOpen) {
-                source.skipAll()
-              }
-              exchange.trailers()
-            }.build()
+            .trailers(
+              object : TrailersSource {
+                override fun peek() = exchange.peekTrailers()
+
+                override fun get(): Headers {
+                  val source = responseBody.source()
+                  if (source.isOpen) {
+                    source.skipAll()
+                  }
+                  return peek() ?: error("null trailers after exhausting response body?!")
+                }
+              },
+            ).build()
         }
       if ("close".equals(response.request.header("Connection"), ignoreCase = true) ||
         "close".equals(response.header("Connection"), ignoreCase = true)

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http/ExchangeCodec.kt

@@ -66,9 +66,9 @@ interface ExchangeCodec {
   @Throws(IOException::class)
   fun openResponseBodySource(response: Response): Source
 
-  /** Returns the trailers after the HTTP response. May be empty. */
+  /** Returns the trailers after the HTTP response if they're ready. May be empty. */
   @Throws(IOException::class)
-  fun trailers(): Headers
+  fun peekTrailers(): Headers?
 
   /**
    * Cancel this stream. Resources held by this stream will be cleaned up, though not synchronously.

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http1/Http1ExchangeCodec.kt

@@ -147,11 +147,14 @@ class Http1ExchangeCodec(
       }
     }
 
-  override fun trailers(): Headers {
+  override fun peekTrailers(): Headers? {
     if (trailers === TRAILERS_RESPONSE_BODY_TRUNCATED) {
       throw IOException("Trailers cannot be read because the response body was truncated")
     }
-    return trailers ?: error("state: $state")
+    check(state == STATE_READING_RESPONSE_BODY || state == STATE_CLOSED) {
+      "Trailers cannot be read because the state is $state"
+    }
+    return trailers
   }
 
   override fun flushRequest() {

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http2/Http2ExchangeCodec.kt

@@ -115,7 +115,7 @@ class Http2ExchangeCodec(
 
   override fun openResponseBodySource(response: Response): Source = stream!!.source
 
-  override fun trailers(): Headers = stream!!.trailers()
+  override fun peekTrailers(): Headers? = stream!!.peekTrailers()
 
   override fun cancel() {
     canceled = true

okhttp/src/commonJvmAndroid/kotlin/okhttp3/internal/http2/Http2Stream.kt

@@ -168,19 +168,18 @@ class Http2Stream internal constructor(
   }
 
   /**
-   * Returns the trailers. It is only safe to call this once the source stream has been completely
-   * exhausted.
+   * Returns the trailers if they're immediately available.
    */
   @Throws(IOException::class)
-  fun trailers(): Headers {
+  fun peekTrailers(): Headers? {
     withLock {
       if (source.finished && source.receiveBuffer.exhausted() && source.readBuffer.exhausted()) {
         return source.trailers ?: Headers.EMPTY
       }
       if (errorCode != null) {
         throw errorException ?: StreamResetException(errorCode!!)
       }
-      throw IllegalStateException("too early; can't read the trailers yet")
+      return null
     }
   }