How to prevent memory leak with In-App Update Library

I want to implement the new In-App Update library in my app, but I've noticed that it trigger a memory leak in my activity when it's recreated/rotated.

Here's the only detail I have from LeakCanary:

Obviously, I've nothing if I remove the code from the In-App Update lib especially the addOnSuccessListener :

appUpdateManager.appUpdateInfo.addOnSuccessListener { appUpdateInfo ->
if (appUpdateInfo.updateAvailability() == UpdateAvailability.UPDATE_AVAILABLE
        && appUpdateInfo.isUpdateTypeAllowed(AppUpdateType.FLEXIBLE)){
            updateInfo.value = appUpdateInfo
            updateAvailable.value = true
        }else{
            updateInfo.value = null
            updateAvailable.value = false
        }
    }

According to this post, I have first used some LiveData, but the problem was the same, so I used a full class to handle the callback, with LiveData :

My Service class :

class AppUpdateService {

    val updateAvailable: MutableLiveData<Boolean> by lazy { MutableLiveData<Boolean>() }
    val updateDownloaded: MutableLiveData<Boolean> by lazy { MutableLiveData<Boolean>() }
    val updateInfo: MutableLiveData<AppUpdateInfo> by lazy { MutableLiveData<AppUpdateInfo>() }

    fun checkForUpdate(appUpdateManager: AppUpdateManager){
        appUpdateManager.appUpdateInfo.addOnSuccessListener { appUpdateInfo ->
            if (appUpdateInfo.updateAvailability() == UpdateAvailability.UPDATE_AVAILABLE
                    && appUpdateInfo.isUpdateTypeAllowed(AppUpdateType.FLEXIBLE)){
                updateInfo.value = appUpdateInfo
                updateAvailable.value = true
            }else{
                updateInfo.value = null
                updateAvailable.value = false
            }
        }
    }

    fun checkUpdateOnResume(appUpdateManager: AppUpdateManager){
        appUpdateManager.appUpdateInfo.addOnSuccessListener {
            updateDownloaded.value = (it.installStatus() == InstallStatus.DOWNLOADED)
        }
    }
}

My Activity simplified :

class MainActivity : BaseActivity(), InstallStateUpdatedListener {

    override fun contentViewID(): Int { return R.layout.activity_main }

    private val UPDATE_REQUEST_CODE = 8000

    private lateinit var appUpdateManager : AppUpdateManager

    private val appUpdateService = AppUpdateService()

    override fun onStateUpdate(state: InstallState?) {
        if(state?.installStatus() == InstallStatus.DOWNLOADED){ notifyUser() }
    }

    // Called in the onCreate()
    override fun setupView(){
        appUpdateManager = AppUpdateManagerFactory.create(this)
        appUpdateManager.registerListener(this)
        setupAppUpdateServiceObservers()
        // Check for Update
        appUpdateService.checkForUpdate(appUpdateManager)
    }

    private fun setupAppUpdateServiceObservers(){
        appUpdateService.updateAvailable.observe(this, Observer {
            if (it)
                requestUpdate(appUpdateService.updateInfo.value)
        })

        appUpdateService.updateDownloaded.observe(this, Observer {
            if (it)
                notifyUser()
        })
    }

    private fun requestUpdate(appUpdateInfo: AppUpdateInfo?){
        appUpdateManager.startUpdateFlowForResult(appUpdateInfo, AppUpdateType.FLEXIBLE, this, UPDATE_REQUEST_CODE)
    }

    private fun notifyUser(){
        showSnackbar(getString(R.string.updated_downloaded), getString(R.string.restart)) {
            appUpdateManager.completeUpdate()
            appUpdateManager.unregisterListener(this)
        }
    }

    override fun onActivityResult(requestCode: Int, resultCode: Int, data: Intent?) {
        super.onActivityResult(requestCode, resultCode, data)
        if (requestCode == UPDATE_REQUEST_CODE) {
            if (resultCode != RESULT_OK) {
                Timber.d("Update flow failed! Result code: $resultCode")
            }
        }
    }

    override fun onDestroy() {
        appUpdateManager.unregisterListener(this)
        super.onDestroy()
    }

    override fun onResume() {
        super.onResume()
        appUpdateService.checkUpdateOnResume(appUpdateManager)
    }

}

I don't really understand how to avoid the memory leak as the appUpdateManager has to be created with the context of the activity, and it looks to be the thing that causes the memory leak with the callback.

Does someone already implement it without having this issue?

Thanks to @Sina Farahzadi I searched and try a lot of things and figured that the problem was the appUpdateManager.appUdateInfo call with the Task object.

The way I found to solve the memory leak is to use the applicationContext instead of the context of the activity. I'm not sure it's the best solution, but it's the one I've found for now. I've exported all in my service class so here's my code :

AppUpdateService.kt :

class AppUpdateService : InstallStateUpdatedListener {

    val updateAvailable: MutableLiveData<Boolean> by lazy { MutableLiveData<Boolean>() }
    val updateDownloaded: MutableLiveData<Boolean> by lazy { MutableLiveData<Boolean>() }
    val notifyUser: MutableLiveData<Boolean> by lazy { MutableLiveData<Boolean>() }
    val updateInfo: MutableLiveData<AppUpdateInfo> by lazy { MutableLiveData<AppUpdateInfo>() }

    private var appUpdateManager : AppUpdateManager? = null
    private var appUpdateInfoTask: Task<AppUpdateInfo>? = null

    override fun onStateUpdate(state: InstallState?) {
        notifyUser.value =  (state?.installStatus() == InstallStatus.DOWNLOADED)
    }

    fun setupAppUpdateManager(context: Context){
        appUpdateManager = AppUpdateManagerFactory.create(context)
        appUpdateManager?.registerListener(this)
        checkForUpdate()
    }

    fun onStopCalled(){
        appUpdateManager?.unregisterListener(this)
        appUpdateInfoTask = null
        appUpdateManager = null
    }

    fun checkForUpdate(){
        appUpdateInfoTask = appUpdateManager?.appUpdateInfo
        appUpdateInfoTask?.addOnSuccessListener { appUpdateInfo ->
            if (appUpdateInfo.updateAvailability() == UpdateAvailability.UPDATE_AVAILABLE
                    && appUpdateInfo.isUpdateTypeAllowed(AppUpdateType.FLEXIBLE)){
                updateInfo.value = appUpdateInfo
                updateAvailable.value = true
            }else{
                updateInfo.value = null
                updateAvailable.value = false
            }
        }
    }

    fun startUpdate(activity: Activity, code: Int){
        appUpdateManager?.startUpdateFlowForResult(updateInfo.value, AppUpdateType.FLEXIBLE, activity, code)
    }

    fun updateComplete(){
        appUpdateManager?.completeUpdate()
        appUpdateManager?.unregisterListener(this)
    }

    fun checkUpdateOnResume(){
        appUpdateManager?.appUpdateInfo?.addOnSuccessListener {
            updateDownloaded.value = (it.installStatus() == InstallStatus.DOWNLOADED)
        }
    }

}

MainActivity simplified :

class MainActivity : BaseActivity(){

    override fun contentViewID(): Int { return R.layout.activity_main }

    private val UPDATE_REQUEST_CODE = 8000

    private var appUpdateService: AppUpdateService? = AppUpdateService()

    /**
     * Setup the view of the activity (navigation and menus)
     */
    override fun setupView(){
        val contextWeakReference = WeakReference<Context>(applicationContext)
        contextWeakReference.get()?.let {weakContext ->
            appUpdateService?.setupAppUpdateManager(weakContext)
        }
    }

    private fun setupAppUpdateServiceObservers(){
        appUpdateService?.updateAvailable?.observe(this, Observer {
            if (it)
                requestUpdate()
        })

        appUpdateService?.updateDownloaded?.observe(this, Observer {
            if (it)
                notifyUser()
        })

        appUpdateService?.notifyUser?.observe(this, Observer {
            if (it)
                notifyUser()
        })
    }

    private fun removeAppUpdateServiceObservers(){
        appUpdateService?.updateAvailable?.removeObservers(this)
        appUpdateService?.updateDownloaded?.removeObservers(this)
        appUpdateService?.notifyUser?.removeObservers(this)
    }

    private fun requestUpdate(){
        appUpdateService?.startUpdate(this, UPDATE_REQUEST_CODE)
    }

    private fun notifyUser(){
        showSnackbar(getString(R.string.updated_downloaded), getString(R.string.restart)) {
            appUpdateService?.updateComplete()
        }
    }

    override fun onActivityResult(requestCode: Int, resultCode: Int, data: Intent?) {
        super.onActivityResult(requestCode, resultCode, data)
        if (requestCode == UPDATE_REQUEST_CODE) {
            if (resultCode != RESULT_OK) {
                Timber.d("Update flow failed! Result code: $resultCode")
            }
        }
    }

    override fun onStop() {
        appUpdateService?.onStopCalled()
        removeAppUpdateServiceObservers()
        appUpdateService = null
        super.onStop()
    }

    override fun onResume() {
        super.onResume()
        setupAppUpdateServiceObservers()
        appUpdateService?.checkUpdateOnResume()
    }

}

For now, I will keep it that way and continue to search for another way to do it. Let me know if someone has a better way to do it.