add retry upload logic on EventPipeline

polbins · polbins · commit 6d1dbc348baf · 2025-03-21T16:23:56.000-07:00
- uses ExponentialBackoffRetryHandler to handle most of the logic
diff --git a/core/src/main/java/com/amplitude/core/platform/EventPipeline.kt b/core/src/main/java/com/amplitude/core/platform/EventPipeline.kt
@@ -1,10 +1,14 @@
 package com.amplitude.core.platform
 
 import com.amplitude.core.Amplitude
+import com.amplitude.core.Storage
 import com.amplitude.core.events.BaseEvent
+import com.amplitude.core.utilities.ExponentialBackoffRetryHandler
 import com.amplitude.core.utilities.http.HttpClient
 import com.amplitude.core.utilities.http.HttpClientInterface
+import com.amplitude.core.utilities.http.ResponseHandler
 import com.amplitude.core.utilities.logWithStackTrace
+import kotlinx.coroutines.CoroutineScope
 import kotlinx.coroutines.channels.Channel
 import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
 import kotlinx.coroutines.channels.consumeEach
@@ -17,21 +21,24 @@ import java.util.concurrent.atomic.AtomicInteger
 
 class EventPipeline(
     private val amplitude: Amplitude,
-) {
-    private val writeChannel: Channel<WriteQueueMessage>
-    private val uploadChannel: Channel<String>
-    private val eventCount: AtomicInteger = AtomicInteger(0)
+    private val eventCount: AtomicInteger = AtomicInteger(0),
     private val httpClient: HttpClientInterface = amplitude.configuration.httpClient
-        ?: HttpClient(amplitude.configuration)
-    private val storage get() = amplitude.storage
-    private val scope get() = amplitude.amplitudeScope
+        ?: HttpClient(amplitude.configuration),
+    private val retryUploadHandler: ExponentialBackoffRetryHandler =
+        ExponentialBackoffRetryHandler(),
+    private val storage: Storage = amplitude.storage,
+    private val scope: CoroutineScope = amplitude.amplitudeScope,
+    private val writeChannel: Channel<WriteQueueMessage> = Channel(UNLIMITED),
+    private var uploadChannel: Channel<String> = Channel(UNLIMITED),
+    overrideResponseHandler: ResponseHandler? = null,
+) {
 
     private var running: Boolean
     private var scheduled: Boolean
     var flushSizeDivider: AtomicInteger = AtomicInteger(1)
 
     private val responseHandler by lazy {
-        storage.getResponseHandler(
+        overrideResponseHandler ?: storage.getResponseHandler(
             this@EventPipeline,
             amplitude.configuration,
             scope,
@@ -40,16 +47,13 @@ class EventPipeline(
     }
 
     companion object {
-        internal const val UPLOAD_SIG = "#!upload"
+        private const val UPLOAD_SIG = "#!upload"
     }
 
     init {
         running = false
         scheduled = false
 
-        writeChannel = Channel(UNLIMITED)
-        uploadChannel = Channel(UNLIMITED)
-
         registerShutdownHook()
     }
 
@@ -107,7 +111,7 @@ class EventPipeline(
 
     private fun upload() =
         scope.launch(amplitude.networkIODispatcher) {
-            uploadChannel.consumeEach {
+            uploadChannel.consumeEach { signal ->
                 withContext(amplitude.storageIODispatcher) {
                     try {
                         storage.rollover()
@@ -118,25 +122,48 @@ class EventPipeline(
                     }
                 }
 
-                val eventsData = storage.readEventsContent()
-                for (events in eventsData) {
-                    try {
-                        val eventsString = storage.getEventsString(events)
-                        if (eventsString.isEmpty()) continue
+                uploadNextEventFile()
+            }
+        }
 
-                        val diagnostics = amplitude.diagnostics.extractDiagnostics()
-                        val response = httpClient.upload(eventsString, diagnostics)
-                        responseHandler.handle(response, events, eventsString)
-                    } catch (e: FileNotFoundException) {
-                        e.message?.let {
-                            amplitude.logger.warn("Event storage file not found: $it")
-                        }
-                    } catch (e: Exception) {
-                        e.logWithStackTrace(amplitude.logger, "Error when uploading event")
+    private suspend fun uploadNextEventFile() {
+        try {
+            // only get first event file, we want to upload them one by one, in order
+            val eventFile = storage.readEventsContent().firstOrNull() ?: return
+            val eventsString = storage.getEventsString(eventFile)
+            if (eventsString.isEmpty()) return
+
+            val diagnostics = amplitude.diagnostics.extractDiagnostics()
+            val response = httpClient.upload(eventsString, diagnostics)
+            val handled = responseHandler.handle(response, eventFile, eventsString)
+
+            when {
+                handled -> retryUploadHandler.reset()
+                !handled && retryUploadHandler.canRetry() -> {
+                    retryUploadHandler.retryWithDelay {
+                        uploadChannel.trySend(UPLOAD_SIG)
                     }
                 }
+                else -> {
+                    amplitude.logger.error(
+                        "Upload failed ${retryUploadHandler.maxRetryAttempt} times. " +
+                            "Cancel all uploads as we might be offline."
+                    )
+
+                    // cancel all previous signals in the upload channel and reset
+                    uploadChannel.cancel()
+                    uploadChannel = Channel(UNLIMITED)
+                    retryUploadHandler.reset()
+                }
+            }
+        } catch (e: FileNotFoundException) {
+            e.message?.let {
+                amplitude.logger.warn("Event storage file not found: $it")
             }
+        } catch (e: Exception) {
+            e.logWithStackTrace(amplitude.logger, "Error when uploading event")
         }
+    }
 
     private fun getFlushCount(): Int {
         val count = amplitude.configuration.flushQueueSize / flushSizeDivider.get()
diff --git a/core/src/test/kotlin/com/amplitude/core/platform/EventPipelineTest.kt b/core/src/test/kotlin/com/amplitude/core/platform/EventPipelineTest.kt
@@ -5,15 +5,30 @@ import com.amplitude.core.Configuration
 import com.amplitude.core.State
 import com.amplitude.core.events.BaseEvent
 import com.amplitude.core.utilities.ConsoleLoggerProvider
+import com.amplitude.core.utilities.ExponentialBackoffRetryHandler
 import com.amplitude.core.utilities.InMemoryStorageProvider
+import com.amplitude.core.utilities.http.BadRequestResponse
+import com.amplitude.core.utilities.http.FailedResponse
+import com.amplitude.core.utilities.http.HttpStatus
+import com.amplitude.core.utilities.http.PayloadTooLargeResponse
+import com.amplitude.core.utilities.http.ResponseHandler
+import com.amplitude.core.utilities.http.SuccessResponse
+import com.amplitude.core.utilities.http.TimeoutResponse
+import com.amplitude.core.utilities.http.TooManyRequestsResponse
 import com.amplitude.id.IMIdentityStorageProvider
+import io.mockk.coVerify
+import io.mockk.every
+import io.mockk.mockk
 import io.mockk.spyk
 import io.mockk.verify
 import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.Channel.Factory.UNLIMITED
 import kotlinx.coroutines.test.StandardTestDispatcher
 import kotlinx.coroutines.test.TestScope
 import kotlinx.coroutines.test.advanceUntilIdle
 import kotlinx.coroutines.test.runTest
+import org.junit.jupiter.api.Assertions.assertTrue
 import org.junit.jupiter.api.Test
 
 @ExperimentalCoroutinesApi
@@ -37,6 +52,7 @@ class EventPipelineTest {
         storageIODispatcher = testDispatcher,
         retryDispatcher = testDispatcher
     )
+    private val fakeResponseHandler: ResponseHandler = mockk(relaxed= true)
 
     @Test
     fun `should not flush when put and offline`() = runTest(testDispatcher) {
@@ -79,4 +95,69 @@ class EventPipelineTest {
 
         verify(exactly = 1) { eventPipeline.flush() }
     }
+
+    @Test
+    fun `should reset retry handler on successful upload`() = runTest(testDispatcher) {
+        amplitude.isBuilt.await()
+        val retryUploadHandler = spyk(ExponentialBackoffRetryHandler())
+        val eventPipeline = EventPipeline(
+            amplitude,
+            retryUploadHandler = retryUploadHandler,
+            overrideResponseHandler = fakeResponseHandler
+        )
+        every { fakeResponseHandler.handle(any(), any(), any()) } returns true
+        val event = BaseEvent().apply { eventType = "test_event" }
+
+        eventPipeline.start()
+        eventPipeline.put(event)
+        advanceUntilIdle()
+
+        verify { retryUploadHandler.reset() }
+    }
+
+    @Test
+    fun `should retry on failure`() = runTest(testDispatcher) {
+        amplitude.isBuilt.await()
+        val retryUploadHandler = spyk(ExponentialBackoffRetryHandler())
+        val eventPipeline = EventPipeline(
+            amplitude,
+            retryUploadHandler = retryUploadHandler,
+            overrideResponseHandler = fakeResponseHandler
+        )
+        val event = BaseEvent().apply { eventType = "test_event" }
+
+        eventPipeline.start()
+        eventPipeline.put(event)
+        advanceUntilIdle()
+
+        coVerify { retryUploadHandler.retryWithDelay(any<suspend () -> Unit>()) }
+    }
+
+    @Test
+    fun `should reset after max retry attempts`() = runTest(testDispatcher) {
+        amplitude.isBuilt.await()
+        val retryUploadHandler = spyk(
+            ExponentialBackoffRetryHandler(
+                maxRetryAttempt = 0
+            )
+        )
+        val uploadChannel = Channel<String>(UNLIMITED)
+        uploadChannel.trySend("test1")
+        uploadChannel.trySend("test2")
+        val eventPipeline = EventPipeline(
+            amplitude,
+            retryUploadHandler = retryUploadHandler,
+            overrideResponseHandler = fakeResponseHandler,
+            uploadChannel = uploadChannel
+        )
+        val event = BaseEvent().apply { eventType = "test_event" }
+
+        eventPipeline.start()
+        eventPipeline.put(event)
+        advanceUntilIdle()
+
+        verify { retryUploadHandler.reset() }
+        assertTrue(uploadChannel.isClosedForSend)
+        assertTrue(uploadChannel.isClosedForReceive)
+    }
 }
