Improve resource cleanup on shutdown

Author: apuig

core/src/main/java/com/segment/analytics/kotlin/core/Analytics.kt

@@ -11,6 +11,8 @@ import com.segment.analytics.kotlin.core.platform.plugins.StartupQueue
 import com.segment.analytics.kotlin.core.platform.plugins.UserInfoPlugin
 import com.segment.analytics.kotlin.core.platform.plugins.logger.*
 import com.segment.analytics.kotlin.core.utilities.JsonAnySerializer
+import com.segment.analytics.kotlin.core.utilities.StorageImpl
+import com.segment.analytics.kotlin.core.utilities.FileEventStream
 import kotlinx.coroutines.*
 import kotlinx.serialization.*
 import kotlinx.serialization.json.Json
@@ -598,14 +600,38 @@ open class Analytics protected constructor(
      * Should only be called in containerized environments where you need to free resources like
      * CoroutineDispatchers and ExecutorService instances so they allow the container to shutdown
      * properly.
+     *
+     * @param waitForTasks if true, waits for all analyticsScope coroutines to complete before shutdown
      */
+    @JvmOverloads
     @OptIn(ExperimentalCoroutinesApi::class)
-    fun shutdown() {
+    fun shutdown(waitForTasks: Boolean = false) {
+        timeline.applyClosure {
+            if (it is com.segment.analytics.kotlin.core.platform.EventPipeline) {
+                it.stop()
+            }
+        }
+
+        val job = analyticsScope.coroutineContext[Job]
+        job?.cancel()
+        if (waitForTasks) {
+            runBlocking {
+                job?.join()
+            }
+        }
+
         (analyticsDispatcher as CloseableCoroutineDispatcher).close()
         (networkIODispatcher as CloseableCoroutineDispatcher).close()
         (fileIODispatcher as CloseableCoroutineDispatcher).close()
 
-        store.shutdown();
+        store.shutdown()
+
+        if (storage is StorageImpl) {
+            val s = storage as StorageImpl
+            if (s.eventStream is FileEventStream) {
+                (s.eventStream as FileEventStream).close()
+            }
+        }
     }
 
     /**
@@ -773,4 +799,4 @@ internal fun isAndroid(): Boolean {
     } catch (ignored: ClassNotFoundException) {
         false
     }
-}
+}

core/src/main/java/com/segment/analytics/kotlin/core/platform/EventPipeline.kt

@@ -53,8 +53,6 @@ open class EventPipeline(
 
         writeChannel = Channel(UNLIMITED)
         uploadChannel = Channel(UNLIMITED)
-
-        registerShutdownHook()
     }
 
     fun put(event: BaseEvent) {
@@ -202,13 +200,4 @@ open class EventPipeline(
 
         return shouldCleanup
     }
-
-    private fun registerShutdownHook() {
-        // close the stream if the app shuts down
-        Runtime.getRuntime().addShutdownHook(object : Thread() {
-            override fun run() {
-                this@EventPipeline.stop()
-            }
-        })
-    }
 }

core/src/main/java/com/segment/analytics/kotlin/core/utilities/EventStream.kt

@@ -160,7 +160,6 @@ open class FileEventStream(
 
     init {
         createDirectory(directory)
-        registerShutdownHook()
     }
 
     protected open var fs: FileOutputStream? = null
@@ -239,13 +238,4 @@ open class FileEventStream(
         val file = File(source)
         return if (file.exists()) FileInputStream(file) else null
     }
-
-    private fun registerShutdownHook() {
-        // close the stream if the app shuts down
-        Runtime.getRuntime().addShutdownHook(object : Thread() {
-            override fun run() {
-                fs?.close()
-            }
-        })
-    }
 }

core/src/main/java/com/segment/analytics/kotlin/core/utilities/EventsFileManager.kt

@@ -41,7 +41,6 @@ class EventsFileManager(
 
     init {
         createDirectory(directory)
-        registerShutdownHook()
     }
 
     private val fileIndexKey = if(subject == null) "segment.events.file.index.$writeKey" else "segment.events.file.index.$writeKey.$subject"
@@ -170,15 +169,6 @@ class EventsFileManager(
         curFile = null
     }
 
-    private fun registerShutdownHook() {
-        // close the stream if the app shuts down
-        Runtime.getRuntime().addShutdownHook(object : Thread() {
-            override fun run() {
-                os?.close()
-            }
-        })
-    }
-
     private suspend fun withLock(block: () -> Unit) {
         semaphore.acquire()
         block()