Add ability to timeout to SSE

Author: p1gp1g

okhttp-sse/src/main/kotlin/okhttp3/sse/EventSourceListener.kt

@@ -18,6 +18,11 @@ package okhttp3.sse
 import okhttp3.Response
 
 abstract class EventSourceListener {
+  /**
+   * Seconds elapsed between 2 events until connection failed. Doesn't timeout if null
+   */
+  open var idleTimeoutMillis: Long? = null
+
   /**
    * Invoked when an event source has been accepted by the remote peer and may begin transmitting
    * events.

okhttp-sse/src/main/kotlin/okhttp3/sse/internal/RealEventSource.kt

@@ -16,6 +16,8 @@
 package okhttp3.sse.internal
 
 import java.io.IOException
+import kotlin.time.Duration
+import kotlin.time.Duration.Companion.seconds
 import okhttp3.Call
 import okhttp3.Callback
 import okhttp3.Request
@@ -24,6 +26,8 @@ import okhttp3.ResponseBody
 import okhttp3.internal.stripBody
 import okhttp3.sse.EventSource
 import okhttp3.sse.EventSourceListener
+import okio.AsyncTimeout
+import okio.Timeout.Companion.timeout
 
 internal class RealEventSource private constructor(
   private val call: Call?,
@@ -38,6 +42,25 @@ internal class RealEventSource private constructor(
 
   @Volatile private var canceled = false
 
+  fun connect(callFactory: Call.Factory) {
+    call =
+      callFactory.newCall(request).apply {
+        enqueue(this@RealEventSource)
+      }
+  }
+
+  private fun updateTimeout(call: Call?, duration: Duration) {
+    if (call?.timeout() is AsyncTimeout) {
+      (call.timeout() as AsyncTimeout).apply {
+        if (this.timeoutNanos() > 0L) {
+          exit()
+        }
+        timeout(duration)
+        enter()
+      }
+    }
+  }
+
   override fun onResponse(
     call: Call,
     response: Response,
@@ -63,8 +86,11 @@ internal class RealEventSource private constructor(
         return
       }
 
-      // This is a long-lived response. Cancel full-call timeouts.
-      call?.timeout()?.cancel()
+      // This is a long-lived response. Cancel full-call timeouts if no timeout has been set
+      listener.idleTimeoutMillis?.let {
+        // We spend at most timeout seconds if set
+        updateTimeout(call, it.seconds)
+      } ?: call?.timeout()?.cancel()
 
       // Replace the body with a stripped one so the callbacks can't see real data.
       val response = response.stripBody()
@@ -74,6 +100,10 @@ internal class RealEventSource private constructor(
         if (!canceled) {
           listener.onOpen(this, response)
           while (!canceled && reader.processNextEvent()) {
+            listener.idleTimeoutMillis?.let {
+              // We spend at most timeout seconds if set
+              updateTimeout(call, it.seconds)
+            }
           }
         }
       } catch (e: Exception) {