From 899cb4af522b3464770b49c68f5d179815873c64 Mon Sep 17 00:00:00 2001
From: Martin Bonnin <martin@mbonnin.net>
Date: Tue, 23 Jul 2024 14:04:28 +0200
Subject: [PATCH] Add WebSocketEngine(WebSocket.Factory) and fix documentation

See #5862
---
 .../advanced/experimental-websockets.mdx      | 40 +++++++++++++------
 .../network/websocket/WebSocketEngine.jvm.kt  | 12 +++---
 2 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/docs/source/advanced/experimental-websockets.mdx b/docs/source/advanced/experimental-websockets.mdx
index 80f1fa12c83..0adf6fd529b 100644
--- a/docs/source/advanced/experimental-websockets.mdx
+++ b/docs/source/advanced/experimental-websockets.mdx
@@ -111,19 +111,35 @@ The above uses the default retry algorithm:
 * Wait until the network is available if you configured a [NetworkMonitor](network-connectivity).
 * Or use exponential backoff else.
 
-To customize the retry logic more, use `addRetryOnErrorInterceptor`:
+To customize the retry logic more, use `retryOnErrorInterceptor`:
 
 ```kotlin
 val apolloClient = ApolloClient.Builder()
-    .subscriptionNetworkTransport(
-        WebSocketNetworkTransport.Builder()
-            .serverUrl(url)
-            .build()
-    )
-    .addRetryOnErrorInterceptor { exception, attempt -> 
-      // retry logic here
-      
-      // return true to retry or false to terminate the Flow
-      true
-    }
+    .retryOnErrorInterceptor(MyRetryOnErrorInterceptor())
+    .build()
+
+class MyRetryOnErrorInterceptor : ApolloInterceptor {
+  object RetryException : Exception()
+
+  override fun <D : Operation.Data> intercept(request: ApolloRequest<D>, chain: ApolloInterceptorChain): Flow<ApolloResponse<D>> {
+    var attempt = 0
+    return request.ensureUniqueUuid()
+        .flatMapConcat { chain.proceed(it) }.onEach {
+          if (request.retryOnError == true && it.exception != null && it.exception is ApolloNetworkException) {
+            throw RetryException
+          } else {
+            attempt = 0
+          }
+        }.retryWhen { cause, _ ->
+          if (cause is RetryException) {
+            attempt++
+            delay(2.0.pow(attempt).seconds)
+            true
+          } else {
+            // Not a RetryException, probably a programming error, pass it through
+            false
+          }
+        }
+  }
+}
 ```
diff --git a/libraries/apollo-runtime/src/jvmCommonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketEngine.jvm.kt b/libraries/apollo-runtime/src/jvmCommonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketEngine.jvm.kt
index 020e254f000..2228b1e48af 100644
--- a/libraries/apollo-runtime/src/jvmCommonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketEngine.jvm.kt
+++ b/libraries/apollo-runtime/src/jvmCommonMain/kotlin/com/apollographql/apollo/network/websocket/WebSocketEngine.jvm.kt
@@ -1,10 +1,9 @@
 package com.apollographql.apollo.network.websocket
 
+import com.apollographql.apollo.annotations.ApolloExperimental
 import com.apollographql.apollo.api.http.HttpHeader
 import com.apollographql.apollo.exception.ApolloNetworkException
 import com.apollographql.apollo.network.defaultOkHttpClientBuilder
-import com.apollographql.apollo.network.websocket.WebSocket
-import com.apollographql.apollo.network.websocket.WebSocketListener
 import okhttp3.Headers
 import okhttp3.OkHttpClient
 import okhttp3.Request
@@ -16,9 +15,9 @@ import java.util.concurrent.atomic.AtomicReference
 import okhttp3.WebSocket as PlatformWebSocket
 import okhttp3.WebSocketListener as PlatformWebSocketListener
 
-internal class JvmWebSocketEngine(private val okHttpClient: OkHttpClient) : WebSocketEngine {
+internal class JvmWebSocketEngine(private val webSocketFactory: PlatformWebSocket.Factory) : WebSocketEngine {
   override fun newWebSocket(url: String, headers: List<HttpHeader>, listener: WebSocketListener): WebSocket {
-    return JvmWebSocket(okHttpClient, url, headers, listener)
+    return JvmWebSocket(webSocketFactory, url, headers, listener)
   }
 
   override fun close() {
@@ -107,4 +106,7 @@ internal class JvmWebSocket(
 }
 
 
-actual fun WebSocketEngine(): WebSocketEngine = JvmWebSocketEngine(defaultOkHttpClientBuilder.build())
\ No newline at end of file
+actual fun WebSocketEngine(): WebSocketEngine = JvmWebSocketEngine(defaultOkHttpClientBuilder.build())
+
+@ApolloExperimental
+fun WebSocketEngine(webSocketFactory: PlatformWebSocket.Factory): WebSocketEngine = JvmWebSocketEngine(webSocketFactory)
\ No newline at end of file