How to cache room data in View Model?

In my Android app, I fetch data from the room database and display it on the UI.

I am caching mapped data. Is this the correct way to cache data in the View Model?

I am new to Android, please correct me if I am wrong.

I am switching dispatcher using withContext, do we need to switch?

@HiltViewModel
class TrainVM @Inject constructor(
    private val repo: Repository,  // Inject the repository
private val apiRepo: KtorRepo

) : ViewModel() {

   
    private var localSchCache = mutableMapOf<String, TrainSchStat>()
    private var trainSchCache = mutableMapOf<String, TrainWithNew>()

    private val _trainStatus =
        MutableStateFlow<ViewState<TrainWithNew>>(value = ViewState.Loading)
    val trainStatus = _trainStatus.asStateFlow()

    private val _isLoading = MutableStateFlow(true)
    val isLoading = _isLoading.asStateFlow()

    fun fetchTrain(train: String, isLive: Boolean, date: LocalDate?) {
        _isLoading.update { true }
        _trainStatus.update { ViewState.Loading }
        if (isLive) {
            loadTrainStatus(train, date = date) //loads live status using API data
        } else {
            loadSchedule(train = train) // loads data from room
        }
    }

    //////////////////////////////////////////

    private fun loadSchedule(
        train: String,
    ) {
        viewModelScope.launch {
            try {
                withContext(Dispatchers.IO) {
                    val localSch = localSchCache[train] ?: repo.getSch(train)
                        .also { localSchCache[train] = it }

                    val trainSch = trainSchCache[train]
                        ?: mappedTrainToStation(localSch).also {
                            trainSchCache[train] = it
                        }

                    _trainStatus.update { ViewState.Success(data = trainSch) }
                }
            } catch (exception: Exception) {
                _trainStatus.update {
                    ViewState.Error(
                        message = exception.message ?: "An unknown error occurred"
                    )
                }
            } finally {
                _isLoading.update { false }
            }
        }
    }
private fun loadTrainStatus(
    train: String,
    date: LocalDate? = null,
) {

   viewModelScope.launch {
        try {
            val localData = localSchCache[train] ?: repo.getSch(train)
                .also { localSchCache[train] = it }

            apiRepo.fetchTrainStatus(train).onSuccess { apiData ->
                val mappedData = mappedTrainToStation(localData, apiData)
        _trainStatus.update { ViewState.Success(data = mappedData) }
            }.onFailure { exception ->
                _trainStatus.update {
                    ViewState.Error(
                        message = exception.message ?: "An unknown error"
                    )
                }
            }
        } catch (exception: Exception) {
            _trainStatus.update {
                ViewState.Error(
                    message = exception.message ?: "An unknown error occurred"
                )
            }
        } finally {
            _isLoading.update { false }
        }
    }
}

The code works fine as desired, I want to ensure I am doing the correct thing.

What you are doing is absolutely correct but there are many ways it can be optimized and everyone has there way to cache data.You are using localSchCache and trainSchCache to cache data, which is fine. However, make sure to handle memory leaks and lifecycle changes carefully. ViewModel lives as long as the UI lifecycle, so this approach will work as long as the ViewModel is alive.

A suggestion for thread Safety and cache eviction.

private val localSchCache = object : LinkedHashMap<String, TrainSchStat>(MAX_CACHE_SIZE, 0.75f, true) {
    override fun removeEldestEntry(eldest: Map.Entry<String, TrainSchStat>): Boolean {
        return size > MAX_CACHE_SIZE
    }
}