聊聊Activity持有ViewModel的清理

在"单Activity"架构的App中，页面功能常常由Fragment承载，而Fragment之间的通信方式往往有两种：通过Fragment Result API或ViewModel；通过Fragment Result API的格式如下：

class FragmentA : Fragment() {
    override fun onCreate(savedInstanceState: Bundle?) {
        super.onCreate(savedInstanceState)
        // 此函数在 fragment-ktx 中
        setFragmentResultListener("requestKey") { requestKey, bundle ->
            val result = bundle.getString("bundleKey")
        }
    }
}

class FragmentB : Fragment() {
    override fun onViewCreated(view: View, savedInstanceState: Bundle?) {
        button.setOnClickListener {
            val result = "result"
            // 此函数在 fragment-ktx 中
            setFragmentResult("requestKey", bundleOf("bundleKey" to result))
        }
    }
}

FragmentA只有处于STARTED状态下才会接收到结果。此种方式适用于比较简单的参数通信，对于比较大的参数（比如Bitmap）或者比较复杂的通信逻辑，ViewModel才是更好的选择。Fragment通过ViewModel通信也比较简单，只需要在创建ViewModel时，指定ViewModelStore为Fragment的宿主Activity即可，这样多个Fragment中获取到的是同一个ViewModel，fragment-ktx库也提供了activityViewModels()函数可以直接获取宿主Activity持有的ViewModel。

上面的方式非常简单也非常方便，但也带来了一些问题，由于ViewModel存储在Activity的ViewModelStore中，当使用此ViewModel的Fragment全都被销毁时，ViewModel仍然不会被释放，事实上此ViewModel已经泄漏了，除此以外，当我们再次创建Fragment使用ViewModel时，ViewModel中仍然是上一次使用时产生的数据，由此可能会带来一些难以预料的问题。

为了解决上面的两个问题，比较常见的解决方案是在使用的Fragment销毁时，手动将ViewModel中的数据重置，如下：

class TestViewModel : ViewModel() {
    private val dataList: MutableList<String> = mutableListOf()
  
    fun addData(string: String) {
        dataList.add(string)
    }
  
    // 没有使用onCleared()是因为此函数无法在外部调用
    fun tearDown() {
        dataList.clear()
    }
}

然后在Fragment的onDestory()中调用TestViewModel#tearDown()即可。

这种手动调用的方式比较繁琐，可以借助LifecycleEventObserver以及ViewModel的onCleared()函数来自动进行“拆卸”ViewModel：

open class AutoTearDownViewModel(
    private val lifecycleOwner: LifecycleOwner
) : ViewModel(), LifecycleEventObserver {
    
    init {
        if (lifecycleOwner.lifecycle.currentState == Lifecycle.State.DESTROYED) {
            onCleared()
        } else {
            lifecycleOwner.lifecycle.addObserver(this)
        }
    }
    
    override fun onStateChanged(source: LifecycleOwner, event: Lifecycle.Event) {
        if (event == Lifecycle.Event.ON_DESTROY) {
            onCleared()
            lifecycleOwner.lifecycle.removeObserver(this)
        }
    }
}

class TestViewModel(val lifecycleOwner: LifecycleOwner) : AutoTearDownViewModel(lifecycleOwner) {
    private val dataList: MutableList<String> = mutableListOf()
    
    override fun onCleared() {
        dataList.clear()
    }
}

class AutoTearDownViewModelFactory(
    private val lifecycleOwner: LifecycleOwner
) : ViewModelProvider.Factory {
    override fun <T : ViewModel?> create(modelClass: Class<T>): T {
        return modelClass.getConstructor(LifecycleOwner::class.java).newInstance(lifecycleOwner)
    }
}

在Fragment中可以通过使用ViewModelProvider(requireActivity(), AutoTearDownViewModelFactory(this))创建ViewModel，在多个Fragment中使用时，创建的ViewModel会跟随第一个创建此ViewModel的Fragment的生命周期。

此种方式虽然不用我们手动去调用清理函数，但是仍然无法解决ViewModel泄漏的问题，并且还是依赖开发者手动在onCleared()中清理或重置数据，在开发时很容易忘记。

那如何能够让不再使用的ViewModel能够自动从Activity的ViewModelStore中消失呢？在上面我们通过LifecycleEventObserver实现了当Fragment销毁时自动清理ViewModel，剩下需要做的就是当Fragment销毁时，将ViewModel从Activity的ViewModelStore中删掉。通过查看ViewModelStore的源码不难发现，ViewModel是存储在HashMap<String, ViewModel>中，此字段是private的，所以我们需要通过反射拿到这个Map，从而将ViewModel删除，所以我们将AutoTearDownViewModel进行以下改造：

// AutoTearDownViewModel.kt
const val VIEW_MODEL_KEY = "auto_tear_down_view_model"

fun <T> getVmKey(clazz: Class<T>): String {
    return VIEW_MODEL_KEY + ":" + clazz.canonicalName
}

open class AutoTearDownViewModel(
    private val fragment: Fragment
) : ViewModel(), LifecycleEventObserver {

    init {
        if (fragment.lifecycle.currentState == Lifecycle.State.DESTROYED) {
            onCleared()
        } else {
            fragment.lifecycle.addObserver(this)
        }
    }

    override fun onStateChanged(source: LifecycleOwner, event: Lifecycle.Event) {
        if (event == Lifecycle.Event.ON_DESTROY) {
            val mapField = ViewModelStore::class.java.getDeclaredField("mMap")
            mapField.isAccessible = true
            val viewModelMap = mapField.get(fragment.requireActivity().viewModelStore) as HashMap<String, ViewModel>
            viewModelMap.remove(getVmKey(this::class.java))?.let {
                val method = ViewModel::class.java.getDeclaredMethod("clear")
                method.isAccessible = true
                method.invoke(it)
            }
            fragment.lifecycle.removeObserver(this)
        }
    }
}

class AutoTearDownViewModelFactory(
    private val fragment: Fragment,
) : ViewModelProvider.Factory {
    override fun <T : ViewModel?> create(modelClass: Class<T>): T {
        return modelClass.getConstructor(Fragment::class.java).newInstance(fragment) as T
    }
}

通过上面的改造后，ViewModel会跟随第一个创建它的Fragment的生命周期，当此Fragment销毁时，ViewModel也将被清理掉，下次再进入到此Fragment时，会创建一个新的ViewModel。

需要注意的是，上面的代码在remove了ViewModel后，会调用一次ViewModel的clear()方法，而此方法是在androidx.lifecycle:lifecycle-viewmodel:2.1.0-alpha01版本之后才被添加进去，如果是之前的版本，则需要调用onCleared()方法。

为了使用方便，我们可以借助Kotlin扩展函数以及委派，实现一个类似activityViewModels()的函数：

// AutoTearDownViewModelProvider.kt
inline fun <reified VM : ViewModel> Fragment.autoTearDownViewModel(): Lazy<VM> {
    return AutoTearDownViewModelLazy(VM::class, this)
}

class AutoTearDownViewModelLazy<VM : ViewModel>(
    private val viewModelClass: KClass<VM>,
    private val fragment: Fragment
) : Lazy<VM> {
    private var cached: VM? = null

    override val value: VM
        get() {
            return cached
                ?: ViewModelProvider(
                    fragment.requireActivity(),
                    AutoTearDownViewModelFactory(fragment)
                ).get(getVmKey(viewModelClass.java), viewModelClass.java).also {
                    cached = it
                }
        }

    override fun isInitialized() = cached != null
}

上述提到的解决方案也并非完美，由于用到了反射，并且反射的字段和方法非public，所以在之后的版本中如果ViewModel或ViewModelStore修改了方法声明可能会导致失效；对于配置更改的情况（比如旋转屏幕），ViewModel仍然会被重建，这其实是和ViewModel的初衷相背的，但如果App限制了只能使用竖屏，那么此种方案也不失为一种比较好的方式。