allow configuring web socket close timeout

mockwebserver/src/main/kotlin/mockwebserver3/MockWebServer.kt

@@ -808,6 +808,7 @@ class MockWebServer : Closeable {
         extensions = WebSocketExtensions.parse(webSocketResponse.headers),
         // Compress all messages if compression is enabled.
         minimumDeflateSize = 0L,
+        closeTimeoutMillis = RealWebSocket.CANCEL_AFTER_CLOSE_MILLIS,
       )
     val name = "MockWebServer WebSocket ${request.path!!}"
     webSocket.initReaderAndWriter(name, streams)

okhttp/api/okhttp.api

@@ -918,6 +918,7 @@ public class okhttp3/OkHttpClient : okhttp3/Call$Factory, okhttp3/WebSocket$Fact
 	public final fun callTimeoutMillis ()I
 	public final fun certificateChainCleaner ()Lokhttp3/internal/tls/CertificateChainCleaner;
 	public final fun certificatePinner ()Lokhttp3/CertificatePinner;
+	public final fun closeTimeoutMillis ()I
 	public final fun connectTimeoutMillis ()I
 	public final fun connectionPool ()Lokhttp3/ConnectionPool;
 	public final fun connectionSpecs ()Ljava/util/List;
@@ -961,6 +962,9 @@ public final class okhttp3/OkHttpClient$Builder {
 	public final fun callTimeout (Ljava/time/Duration;)Lokhttp3/OkHttpClient$Builder;
 	public final fun callTimeout-LRDsOJo (J)Lokhttp3/OkHttpClient$Builder;
 	public final fun certificatePinner (Lokhttp3/CertificatePinner;)Lokhttp3/OkHttpClient$Builder;
+	public final fun closeTimeout (JLjava/util/concurrent/TimeUnit;)Lokhttp3/OkHttpClient$Builder;
+	public final fun closeTimeout (Ljava/time/Duration;)Lokhttp3/OkHttpClient$Builder;
+	public final fun closeTimeout-LRDsOJo (J)Lokhttp3/OkHttpClient$Builder;
 	public final fun connectTimeout (JLjava/util/concurrent/TimeUnit;)Lokhttp3/OkHttpClient$Builder;
 	public final fun connectTimeout (Ljava/time/Duration;)Lokhttp3/OkHttpClient$Builder;
 	public final fun connectTimeout-LRDsOJo (J)Lokhttp3/OkHttpClient$Builder;

okhttp/src/main/kotlin/okhttp3/OkHttpClient.kt

@@ -252,6 +252,10 @@ open class OkHttpClient internal constructor(
   @get:JvmName("pingIntervalMillis")
   val pingIntervalMillis: Int = builder.pingInterval
 
+  /** Web socket close timeout (in milliseconds). */
+  @get:JvmName("closeTimeoutMillis")
+  val closeTimeoutMillis: Int = builder.closeTimeout
+
   /**
    * Minimum outbound web socket message size (in bytes) that will be compressed.
    * The default is 1024 bytes.
@@ -327,6 +331,7 @@ open class OkHttpClient internal constructor(
         // extensions is always null for clients:
         extensions = null,
         minimumDeflateSize = minWebSocketMessageToCompress,
+        closeTimeoutMillis = closeTimeoutMillis.toLong(),
       )
     webSocket.connect(this)
     return webSocket
@@ -572,6 +577,7 @@ open class OkHttpClient internal constructor(
     internal var readTimeout = 10_000
     internal var writeTimeout = 10_000
     internal var pingInterval = 0
+    internal var closeTimeout = RealWebSocket.CANCEL_AFTER_CLOSE_MILLIS.toInt()
     internal var minWebSocketMessageToCompress = RealWebSocket.DEFAULT_MINIMUM_DEFLATE_SIZE
     internal var routeDatabase: RouteDatabase? = null
     internal var taskRunner: TaskRunner? = null
@@ -606,6 +612,7 @@ open class OkHttpClient internal constructor(
       this.readTimeout = okHttpClient.readTimeoutMillis
       this.writeTimeout = okHttpClient.writeTimeoutMillis
       this.pingInterval = okHttpClient.pingIntervalMillis
+      this.closeTimeout = okHttpClient.closeTimeoutMillis
       this.minWebSocketMessageToCompress = okHttpClient.minWebSocketMessageToCompress
       this.routeDatabase = okHttpClient.routeDatabase
       this.taskRunner = okHttpClient.taskRunner
@@ -1276,6 +1283,49 @@ open class OkHttpClient internal constructor(
         pingInterval = checkDuration("duration", duration)
       }
 
+    /**
+     * Sets the close timeout for web socket connections. A value of 0 means no timeout, otherwise
+     * values must be between 1 and [Integer.MAX_VALUE] when converted to milliseconds.
+     *
+     * The close timeout is the maximum amount of time after the client calls [close] to wait
+     * for a graceful shutdown. If the server doesn't respond the web socket will be canceled.
+     * The default value is 60 seconds.
+     */
+    fun closeTimeout(
+      timeout: Long,
+      unit: TimeUnit,
+    ) = apply {
+      closeTimeout = checkDuration("timeout", timeout, unit)
+    }
+
+    /**
+     * Sets the close timeout for web socket connections. A value of 0 means no timeout, otherwise
+     * values must be between 1 and [Integer.MAX_VALUE] when converted to milliseconds.
+     *
+     * The close timeout is the maximum amount of time after the client calls [close] to wait
+     * for a graceful shutdown. If the server doesn't respond the web socket will be canceled.
+     * The default value is 60 seconds.
+     */
+    @SuppressLint("NewApi")
+    @IgnoreJRERequirement
+    fun closeTimeout(duration: Duration) =
+      apply {
+        closeTimeout(duration.toMillis(), MILLISECONDS)
+      }
+
+    /**
+     * Sets the close timeout for web socket connections. A value of 0 means no timeout, otherwise
+     * values must be between 1 and [Integer.MAX_VALUE] when converted to milliseconds.
+     *
+     * The close timeout is the maximum amount of time after the client calls [close] to wait
+     * for a graceful shutdown. If the server doesn't respond the web socket will be canceled.
+     * The default value is 60 seconds.
+     */
+    fun closeTimeout(duration: KotlinDuration) =
+      apply {
+        closeTimeout = checkDuration("duration", duration)
+      }
+
     /**
      * Sets minimum outbound web socket message size (in bytes) that will be compressed.
      *

okhttp/src/main/kotlin/okhttp3/internal/ws/RealWebSocket.kt

@@ -64,6 +64,7 @@ class RealWebSocket(
   private var extensions: WebSocketExtensions?,
   /** If compression is negotiated, outbound messages of this size and larger will be compressed. */
   private var minimumDeflateSize: Long,
+  private val closeTimeoutMillis: Long,
 ) : WebSocket, WebSocketReader.FrameCallback {
   private val key: String
 
@@ -469,7 +470,7 @@ class RealWebSocket(
     code: Int,
     reason: String?,
   ): Boolean {
-    return close(code, reason, CANCEL_AFTER_CLOSE_MILLIS)
+    return close(code, reason, closeTimeoutMillis)
   }
 
   @Synchronized fun close(
@@ -712,7 +713,7 @@ class RealWebSocket(
      * The maximum amount of time after the client calls [close] to wait for a graceful shutdown. If
      * the server doesn't respond the web socket will be canceled.
      */
-    private const val CANCEL_AFTER_CLOSE_MILLIS = 60L * 1000
+    const val CANCEL_AFTER_CLOSE_MILLIS = 60L * 1000
 
     /**
      * The smallest message that will be compressed. We use 1024 because smaller messages already

okhttp/src/test/java/okhttp3/OkHttpClientTest.kt

@@ -74,6 +74,7 @@ class OkHttpClientTest {
     assertThat(client.readTimeoutMillis).isEqualTo(10000)
     assertThat(client.writeTimeoutMillis).isEqualTo(10000)
     assertThat(client.pingIntervalMillis).isEqualTo(0)
+    assertThat(client.closeTimeoutMillis).isEqualTo(60_000)
   }
 
   @Test fun webSocketDefaults() {

okhttp/src/test/java/okhttp3/internal/ws/RealWebSocketTest.kt

@@ -460,6 +460,7 @@ class RealWebSocketTest {
             responseHeaders,
           ),
           RealWebSocket.DEFAULT_MINIMUM_DEFLATE_SIZE,
+          RealWebSocket.CANCEL_AFTER_CLOSE_MILLIS,
         )
       webSocket!!.initReaderAndWriter(name, this)
     }

okhttp/src/test/java/okhttp3/internal/ws/WebSocketHttpTest.kt

@@ -1214,6 +1214,7 @@ class WebSocketHttpTest {
         client.pingIntervalMillis.toLong(),
         null,
         0L,
+        client.closeTimeoutMillis.toLong(),
       )
     webSocket.connect(client)
     return webSocket