androidinterview.com

Low Level Design (LLD) Interview Questions

Implement your own ViewModel

Tier: EssentialDifficulty: MediumAsked of: Mid, Senior

The whole thing is a map, and one object that holds the map for slightly longer than the screen does. A ViewModel survives a rotation because the map it lives in is not owned by the Activity object at all. It is owned by something in the process, which hands the same map to the next Activity object a moment later. Once you have said that sentence, the base class, the provider, the factory and the scope are all small.

The mistake almost everyone makes is saying that the ViewModel is retained. Nothing retains a ViewModel. Something retains the store, and the store happens to hold ViewModels. Say it the second way and every follow up in this interview answers itself.

What this really tests

Whether you know what actually happens on a rotation, or whether you have only used the API. There are three questions hiding inside this one. Who owns the map, who decides when it is emptied, and how does anything hanging off a model get shut down at that same moment. The weak answer is a static map keyed by class. It does survive rotation, which is why it feels right for about ten seconds, and it also survives the screen closing, so it leaks every screen the user ever opened.

What to clarify first

Ask these before you write a class name. Two of them cut the problem in half.

  • Am I building the retention mechanism, or only the base class. These are different questions. The base class is tiny. The retention is the interesting part, and it is almost always what is being asked.
  • Does it have to survive process death. The answer is no, and saying so early is worth more than any class you draw. Rotation and process death are different problems with different mechanisms.
  • Is the owner an Activity only, or fragments as well. Fragments change the design, because they need an owner of their own that lives inside the Activity's.
  • Do models take constructor arguments. If yes, creation has to be a separate object from lookup, which is the reason a factory exists at all.
  • Is there background work to cancel. If yes, the base class needs a general way to hold a resource and shut it down, not a coroutine shaped special case.
  • How much of the framework may I assume. Assume nothing. Model a tiny lifecycle and a tiny activity yourself, because the answer lives in the seam between them and you cannot show that seam if you hand wave the framework.

State the out of scope list too. Saved state, dependency injection, the main thread, and anything to do with views.

Why rotation is survived and process death is not

On a rotation the framework destroys the Activity object and builds a brand new one. The object is genuinely gone. If the store were a field on it, the store would be gone too, so something outside the object has to be holding it during the gap.

That something is a retained slot on the framework side. On the way out the Activity is asked for anything it wants to keep, it hands back a small holder carrying the store, and the framework parks it. A moment later the new Activity is created, is given that holder, and adopts the store inside it. The new screen then asks the provider for a ViewModel, the provider looks in the store, finds the instance the old screen created, and returns it. Nothing was serialised and nothing was copied. It is the same object at the same address.

Now say the second half. That slot is memory in your process. When the system kills the process for memory, every page of it goes, including the slot, including the store, including your models. No destroy callback runs and onCleared is never called, because there is nobody left to call it. That is not a bug you can fix inside this design. Anything that must come back after process death has to have been written somewhere that is not memory, which is a completely separate mechanism.

Two designs to reject out loud, because an interviewer will offer you one of them. A static map keyed by class survives rotation and also survives the screen closing, so it is a leak with extra steps. Putting the models on the Application object is the same leak wearing a nicer name. The point of the retained slot is that it is scoped to one screen and it disappears when that screen finishes.

The classes

Fourteen types, and none of them is large. The interesting part is which one owns which decision.

  • ViewModel is the base class. It has one hook, onCleared, that runs exactly once at the end of its life, and a bag of AutoCloseable resources that are closed at that same moment. Its clear method is package private in Java and internal in Kotlin, because only the store that holds a model is allowed to end it.
  • ViewModelStore is the map from key to instance, and that is genuinely all it is. It can put, get and clear. It is not a singleton and it is not static, because there is one per screen and its life is the screen's life.
  • ViewModelStoreOwner is one method that hands back a store. It is the seam the design turns on. Anything that can produce a store can own ViewModels, so the provider never learns whether it is talking to an Activity, a fragment or a test double.
  • ViewModelProvider looks up by key and creates only when the lookup misses. It holds no cache of its own, which is why building a fresh provider on every call is free and why asking twice can never produce two instances.
  • Factory is the one method that builds a model. Creation is a separate job from lookup, because a model with constructor arguments still has to be found in the store on the way back. NewInstanceFactory is the fallback for a model that needs nothing.
  • Lifecycle is the smallest lifecycle that makes the point, a current state and a list of observers. Real Android has more states and none of them change the answer.
  • ComponentActivity owns a store, hands it over on the way out as NonConfigurationInstances, and clears it on destroy only when it is not being rebuilt. That one condition is the entire question.
  • ActivityThread is the framework side. It lives in the process rather than in a screen, so it can hold the retained slot across the gap. It also gives us an honest way to model process death, which is a method that drops everything without running a single callback.
  • Fragment is a second owner living inside the first, and FragmentStoreHolder is the trick that makes it work. Fragments have no retained slot of their own, so their stores are kept inside a ViewModel that lives in the Activity's store. They ride the same rotation and they are cleared by the same call.
  • CloseableTaskScope is a scope of background work that shuts down in one call, and ViewModelScope is the property that attaches one to a model under a known key.

Say the mapping out loud so the interviewer can follow you. ViewModel, ViewModelStore, ViewModelStoreOwner, ViewModelProvider, Factory, NewInstanceFactory, Lifecycle, ComponentActivity and NonConfigurationInstances are the real Jetpack names, used here unchanged. Three are renamed. CloseableTaskScope stands in for CloseableCoroutineScope, FragmentStoreHolder stands in for FragmentManagerViewModel, and ActivityThread is doing the job the real framework does across several classes.

ViewModel class diagramClasses ViewModelStoreOwner, ViewModelProvider, Factory, ComponentActivity, ViewModelStore, ViewModel, ActivityThread, Fragment, CloseableTaskScope. ViewModelProvider asks ViewModelStoreOwner. ViewModelProvider creates via Factory. ViewModelProvider looks up in 1 ViewModelStore. ComponentActivity implements ViewModelStoreOwner. ComponentActivity is composed of 1 ViewModelStore. ViewModelStore is composed of 0..* ViewModel. ActivityThread retains store ComponentActivity. Fragment host ComponentActivity. ViewModel aggregates 0..* CloseableTaskScope.
ViewModel class diagram, a UML class diagram of ViewModelStoreOwner, ViewModelProvider, Factory, ComponentActivity, ViewModelStore, ViewModel, ActivityThread, Fragment, CloseableTaskScope
The accented edge is the whole answer. Every other line is ordinary object wiring, and only that one explains why an instance outlives the screen that asked for it.

A rotation, walked through

Follow one model through a device turning sideways. This is the walk to do at the whiteboard.

  1. The screen is created, and it asks a provider for a CartViewModel. It builds the provider fresh, because a provider caches nothing.
  2. The provider asks the owner for its store. The Activity has none yet, so it makes one now. A screen that never asks for a model never allocates a store at all.
  3. The provider builds a key from the class name, looks it up, and misses. It calls the factory, gets an instance, puts it in the store under that key, and returns it.
  4. The user turns the phone. Before anything is destroyed, the framework asks the Activity for what it wants to keep. The Activity wraps its store in a NonConfigurationInstances and hands it over.
  5. The Activity is destroyed with its changing configurations flag set. Its lifecycle reaches destroyed, the observer that clears the store looks at that flag, and does nothing. This is the single most important line in the design.
  6. A new Activity object is built. The framework passes it the holder from step four, and it adopts the store inside rather than creating one.
  7. The new screen asks a new provider for a CartViewModel. Same key, and this time the lookup hits. The instance from step three comes back, with its data, its in flight work and its scope all intact. The factory is never called.
  8. The user presses back. The Activity is destroyed with the flag clear, so the observer calls clear on the store. Every model in it has its closeables closed and then its onCleared run, once.

Then say what step eight would look like on a process kill. Nothing. No destroy, no clear, no onCleared. The map is simply not there any more, which is why cleanup in onCleared is a best effort and never a guarantee.

Keys, and one model shared by two fragments

The key is why the provider can promise you one instance and still let you have two. Ask for a type and the key is derived from the class name, so one screen gets one model of each type without anybody naming anything. Ask with an explicit key and two models of the same class sit side by side in the same store, which is what you need for a list of tabs that each own the same kind of state. Getting the same key twice always returns the same instance, and getting a key that already holds a different type is an error rather than a silent replacement.

Sharing between fragments is not a feature, it is a choice of which owner you ask. Ask the fragment and you get its own store, which lives in the FragmentStoreHolder inside the Activity's store and dies when the fragment is popped. Ask the Activity and you get the Activity's store, which every fragment on that screen reaches. Two fragments asking the Activity for the same type get the same instance, because there is one store and one key. That is the whole of the shared ViewModel pattern, and it is worth saying that it is also its limitation. The sharing is by identity in a map, so two fragments in two different Activities share nothing.

How viewModelScope is wired

A ViewModel that starts background work has to stop it at exactly the moment the model dies, and the base class must not know what kind of work it is. So the base class holds AutoCloseable and nothing more specific. A scope is registered on the model under a known key, and clear closes everything in the bag before it calls onCleared.

The scope itself is an object you can launch work into and close once. Closing cancels everything, and cancellation is cooperative, so work that has not started never runs and work that has already started sees its job go inactive at the next check. That is the same contract a coroutine has, without any of the machinery.

In Kotlin this reads as an extension property on ViewModel, which is exactly how the real viewModelScope is written. It looks for a scope under a fixed key, creates one if there is none, and registers it. The registration is first caller wins, so two threads touching the property at the same moment still end up sharing one scope. In Java the same thing is a static helper, because Java has no extension properties, and the mechanism underneath is identical.

What SavedStateHandle adds, and why it is a different mechanism

If the interviewer asks for process death, do not try to patch this design. Say that the store is memory and memory does not survive, then describe the second mechanism.

SavedStateHandle is a map the system writes into a bundle when the process is about to be killed, and reads back when the user returns. Two things follow from that and they are the reason it is a separate object rather than a feature of the store. Everything in it has to be serialisable, because it crosses a process boundary, so it holds a screen's worth of identifiers and not a screen's worth of loaded data. And it is written by the framework at save time rather than kept alive, so it is a snapshot and never the live instance.

The wiring is worth one sentence. The handle is created outside the ViewModel and passed into its constructor, which means the factory is the piece that changes, not the store and not the provider. That is a good sign about the design. A new way to build models slots into the one class whose job is building models.

Patterns actually used

Two patterns carry this, and the interesting part is how few there are.

  • A registry, sometimes called an identity map. The store is a keyed cache with an ownership rule, which is that entries are created on miss and destroyed together. Calling it a registry rather than a HashMap is the difference between a candidate who has named the concept and one who has not.
  • Factory, and it earns its place here. The provider must be able to find a model it did not build, so lookup and creation cannot be the same method. The moment a model needs a repository in its constructor, the factory is the only thing that changes.

Three to leave out, and say why.

  • No singleton anywhere. It is the reflex answer to make anything survive and it is wrong here in the exact way that matters. A singleton store survives the screen closing as well, so the leak is the feature you asked for.
  • No observer on the store. The lifecycle already publishes destroyed, and adding a second notification path so models can hear about their own removal is a loop waiting to happen. The store calls clear directly.
  • No service locator around the provider. A provider is cheap, stateless and built at the call site. Caching one in a global would only reintroduce the lifetime problem this design just solved.

There is a defensible argument for skipping the Factory interface entirely and taking a lambda that returns a model. That is what most Kotlin code does in practice. Name the tradeoff rather than defending the interface by reflex. The interface earns its keep when one factory has to serve several types, which is the case for a saved state factory, and a lambda is better everywhere else.

The implementation

Both trees are the same design, and both compile as plain JVM code with no Android and no coroutines. The lifecycle, the Activity and the framework side are modelled in a few dozen lines each, because the seam between them is where the answer lives.

The Kotlin is written as Kotlin. The owner exposes a property rather than a getter, the store carries get and set operators so a lookup reads like an index, the key is a default argument instead of an overload, and reified generics remove the Class argument the Java side has to pass around. It also carries a viewModels delegate and an activityViewModels delegate, so a screen declares its model as one property, which is how anybody actually writes this.

Java

com.androidinterview.viewmodel.ViewModel.java

package com.androidinterview.viewmodel;

import java.util.ArrayList;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;

// The base class. Two things make it a ViewModel rather than a plain object.
// It has one hook that runs exactly once, at the end of its life, and a bag of
// AutoCloseable resources closed at that same moment. Every scope and every
// subscription a real ViewModel holds is one of those closeables, so there is a
// single place that has to be right.
public abstract class ViewModel {

    private final Map<String, AutoCloseable> closeables = new LinkedHashMap<>();
    private int anonymous;
    private boolean cleared;

    // Called once, when the store that owns this model lets it go. Never on a
    // rotation. Override it to drop anything that is not already a closeable.
    protected void onCleared() { }

    public final void addCloseable(AutoCloseable closeable) {
        addCloseable("anonymous:" + anonymous++, closeable);
    }

    // The keyed form, and the one an extension like viewModelScope uses. First
    // caller wins, so a scope is created once even if two threads ask at once.
    @SuppressWarnings("unchecked")
    public final <T extends AutoCloseable> T addCloseable(String key, T closeable) {
        synchronized (this) {
            if (!cleared) {
                AutoCloseable existing = closeables.get(key);
                if (existing != null) {
                    return (T) existing;
                }
                closeables.put(key, closeable);
                return closeable;
            }
        }
        // Registering on an already cleared model is a race, not a bug. Close
        // the resource straight away rather than leak it.
        closeQuietly(closeable);
        return closeable;
    }

    public final synchronized AutoCloseable getCloseable(String key) {
        return closeables.get(key);
    }

    // Package private on purpose. Only the store holding this model may end its
    // life, so no screen can clear a model another screen is still using.
    final void clear() {
        List<AutoCloseable> toClose;
        synchronized (this) {
            if (cleared) {
                return;
            }
            cleared = true;
            toClose = new ArrayList<>(closeables.values());
            closeables.clear();
        }
        toClose.forEach(ViewModel::closeQuietly);
        onCleared();
    }

    private static void closeQuietly(AutoCloseable closeable) {
        try {
            closeable.close();
        } catch (Exception ignored) {
            // One resource failing to close must not strand the rest.
        }
    }
}

com.androidinterview.viewmodel.ViewModelProvider.java

package com.androidinterview.viewmodel;

// Look up by key, and only build when the lookup misses. That is the entire
// class. It owns no models, it holds no cache of its own, and constructing a
// second provider over the same owner changes nothing, which is why calling it
// in onCreate every time is safe.
public final class ViewModelProvider {

    private static final String DEFAULT_KEY = "com.androidinterview.viewmodel.ViewModelProvider.DefaultKey";

    private final ViewModelStore store;
    private final Factory factory;

    public ViewModelProvider(ViewModelStoreOwner owner) {
        this(owner, new NewInstanceFactory());
    }

    public ViewModelProvider(ViewModelStoreOwner owner, Factory factory) {
        this.store = owner.viewModelStore();
        this.factory = factory;
    }

    // The default key is derived from the class, so one screen gets one model
    // of each type without anybody naming anything.
    public <T extends ViewModel> T get(Class<T> type) {
        return get(DEFAULT_KEY + ":" + type.getName(), type);
    }

    public <T extends ViewModel> T get(String key, Class<T> type) {
        ViewModel existing = store.get(key);
        if (existing != null) {
            if (type.isInstance(existing)) {
                return type.cast(existing);
            }
            throw new IllegalArgumentException(
                    "key " + key + " already holds a " + existing.getClass().getName());
        }
        T created = factory.create(type);
        store.put(key, created);
        return created;
    }

    // Creation is a separate job from lookup, because a model with constructor
    // arguments still has to be found in the store on the way back.
    public interface Factory {

        <T extends ViewModel> T create(Class<T> type);
    }

    // The fallback for a model that needs nothing. Anything with dependencies
    // gets its own factory, which is a lambda in practice.
    public static final class NewInstanceFactory implements Factory {

        @Override
        public <T extends ViewModel> T create(Class<T> type) {
            try {
                return type.getDeclaredConstructor().newInstance();
            } catch (ReflectiveOperationException failure) {
                throw new IllegalArgumentException(
                        type.getName() + " needs a public no argument constructor", failure);
            }
        }
    }
}

com.androidinterview.viewmodel.ViewModelStore.java

package com.androidinterview.viewmodel;

import java.util.HashMap;
import java.util.Map;

// A map from key to instance, and nothing else. This is the whole of the
// mechanism people think is magic. It is deliberately not a singleton and not
// static, because there is one of these per screen and its life is the screen's
// life, not the process's.
public final class ViewModelStore {

    private final Map<String, ViewModel> models = new HashMap<>();

    public ViewModel get(String key) {
        return models.get(key);
    }

    public void put(String key, ViewModel model) {
        ViewModel previous = models.put(key, model);
        if (previous != null) {
            // Replacing under a live key would silently leak the old model's
            // resources, so it is ended here rather than left to the collector.
            previous.clear();
        }
    }

    // The only way a model ever dies. Called when the owner is finishing for
    // real, never when it is being recreated.
    public void clear() {
        for (ViewModel model : models.values()) {
            model.clear();
        }
        models.clear();
    }
}

com.androidinterview.viewmodel.ViewModelStoreOwner.java

package com.androidinterview.viewmodel;

// One method, and it is the seam the whole design turns on. Anything that can
// hand back a store can own ViewModels, so the provider never has to know
// whether it is talking to an activity, a fragment or a test double.
public interface ViewModelStoreOwner {

    ViewModelStore viewModelStore();
}

com.androidinterview.viewmodel.host.ActivityThread.java

package com.androidinterview.viewmodel.host;

import java.util.HashMap;
import java.util.Map;
import java.util.function.Supplier;

// The framework side, and the reason any of this survives a rotation. It lives
// in the process rather than in a screen, so it can hold the retained slot in
// the gap where the old activity object is gone and the new one does not exist
// yet. It is also the whole answer to process death. Kill the process and this
// map goes with it, no callback runs, and nothing was ever written down.
public final class ActivityThread {

    private static final class Record {

        final Supplier<ComponentActivity> builder;
        ComponentActivity current;

        Record(Supplier<ComponentActivity> builder) {
            this.builder = builder;
        }
    }

    private final Map<String, Record> records = new HashMap<>();

    public ComponentActivity launch(String id, Supplier<ComponentActivity> builder) {
        Record record = new Record(builder);
        record.current = builder.get();
        record.current.attach(null);
        records.put(id, record);
        return record.current;
    }

    // A rotation. The old screen is destroyed and a brand new one is built, and
    // the only thing that crosses the gap is the retained slot.
    public ComponentActivity configurationChange(String id) {
        Record record = recordFor(id);
        ComponentActivity.NonConfigurationInstances retained =
                record.current.retainNonConfigurationInstances();
        record.current.performDestroy(true);
        record.current = record.builder.get();
        record.current.attach(retained);
        return record.current;
    }

    // A real finish, a back press or a call to finish. Nothing is retained, so
    // the destroy clears the store and every model gets onCleared.
    public void finish(String id) {
        recordFor(id).current.performDestroy(false);
        records.remove(id);
    }

    // Process death, modelled honestly. No destroy runs, no store is cleared,
    // and every model in the process is simply gone.
    public void killProcess() {
        records.clear();
    }

    private Record recordFor(String id) {
        Record record = records.get(id);
        if (record == null) {
            throw new IllegalArgumentException("no activity with id " + id);
        }
        return record;
    }
}

com.androidinterview.viewmodel.host.ComponentActivity.java

package com.androidinterview.viewmodel.host;

import com.androidinterview.viewmodel.ViewModelStore;
import com.androidinterview.viewmodel.ViewModelStoreOwner;

// The screen. It owns a store, hands that store to the framework on the way
// out, and clears it only when it is going away for good. It never keeps a
// static reference to anything and it never asks whether the device rotated. It
// answers one question, am I finishing or am I being rebuilt.
public class ComponentActivity implements ViewModelStoreOwner {

    // The retained slot. In real Android this also carries retained fragments
    // and a custom object, and the only field that matters here is the store.
    public record NonConfigurationInstances(ViewModelStore viewModelStore) { }

    private final Lifecycle lifecycle = new Lifecycle();
    private ViewModelStore store;
    private boolean changingConfigurations;

    public ComponentActivity() {
        // The one line the whole question is about. A destroy that is part of a
        // rebuild leaves the store alone, because the next instance is about to
        // pick it up. Any other destroy ends every model in it.
        lifecycle.addObserver(state -> {
            if (state == Lifecycle.State.DESTROYED && !changingConfigurations && store != null) {
                store.clear();
            }
        });
    }

    public final Lifecycle lifecycle() {
        return lifecycle;
    }

    // Created on first use, so a screen that never asks for a ViewModel never
    // allocates a store and never has one retained for it.
    @Override
    public final ViewModelStore viewModelStore() {
        if (store == null) {
            store = new ViewModelStore();
        }
        return store;
    }

    public final boolean isChangingConfigurations() {
        return changingConfigurations;
    }

    protected void onCreate() { }

    protected void onDestroy() { }

    // The three calls below belong to the framework side and nothing else.
    final void attach(NonConfigurationInstances last) {
        if (last != null) {
            store = last.viewModelStore();
        }
        onCreate();
        lifecycle.moveTo(Lifecycle.State.RESUMED);
    }

    final NonConfigurationInstances retainNonConfigurationInstances() {
        return store == null ? null : new NonConfigurationInstances(store);
    }

    final void performDestroy(boolean changingConfigurations) {
        this.changingConfigurations = changingConfigurations;
        onDestroy();
        lifecycle.moveTo(Lifecycle.State.DESTROYED);
    }
}

com.androidinterview.viewmodel.host.Fragment.java

package com.androidinterview.viewmodel.host;

import com.androidinterview.viewmodel.ViewModelProvider;
import com.androidinterview.viewmodel.ViewModelStore;
import com.androidinterview.viewmodel.ViewModelStoreOwner;

// A fragment is a second owner living inside the first, and that sentence is
// the whole of ViewModel sharing. Ask this fragment for a store and you get one
// private to it. Ask the activity and every fragment gets the same one, so a
// shared model is a choice of which owner you ask, not a feature.
public class Fragment implements ViewModelStoreOwner {

    private final String tag;
    private ComponentActivity host;

    public Fragment(String tag) {
        this.tag = tag;
    }

    public void attachTo(ComponentActivity activity) {
        this.host = activity;
    }

    public ComponentActivity requireActivity() {
        if (host == null) {
            throw new IllegalStateException("fragment " + tag + " is not attached");
        }
        return host;
    }

    @Override
    public ViewModelStore viewModelStore() {
        return holder().storeFor(tag);
    }

    // The shared owner. Two fragments that go through this get the same
    // instance, because there is one activity and one activity store.
    public ViewModelStoreOwner activityOwner() {
        return requireActivity();
    }

    // Popped off the back stack rather than rotated. The fragment's own store
    // is cleared here, which is why a fragment scoped model dies with the
    // fragment while an activity scoped one carries on.
    public void removeForGood() {
        holder().remove(tag);
        host = null;
    }

    private FragmentStoreHolder holder() {
        return new ViewModelProvider(requireActivity()).get(FragmentStoreHolder.class);
    }
}

com.androidinterview.viewmodel.host.FragmentStoreHolder.java

package com.androidinterview.viewmodel.host;

import java.util.HashMap;
import java.util.Map;

import com.androidinterview.viewmodel.ViewModel;
import com.androidinterview.viewmodel.ViewModelStore;

// A ViewModel whose only job is to hold other stores, one per fragment tag.
// Fragments have no retained slot of their own, so this is the trick that gives
// them one. Their stores live inside a model in the activity's store, so they
// ride the same rotation and are cleared by the same call.
public final class FragmentStoreHolder extends ViewModel {

    private final Map<String, ViewModelStore> stores = new HashMap<>();

    public ViewModelStore storeFor(String tag) {
        return stores.computeIfAbsent(tag, key -> new ViewModelStore());
    }

    public void remove(String tag) {
        ViewModelStore removed = stores.remove(tag);
        if (removed != null) {
            removed.clear();
        }
    }

    @Override
    protected void onCleared() {
        for (ViewModelStore store : stores.values()) {
            store.clear();
        }
        stores.clear();
    }
}

com.androidinterview.viewmodel.host.Lifecycle.java

package com.androidinterview.viewmodel.host;

import java.util.ArrayList;
import java.util.List;

// The minimum lifecycle the design needs. Real Android has more states and a
// bigger observer registry, and none of that changes the answer. What matters
// is that something can say the owner has reached its end, once, and that the
// store can listen for it.
public final class Lifecycle {

    public enum State { INITIALIZED, CREATED, STARTED, RESUMED, DESTROYED }

    public interface Observer {
        void onStateChanged(State state);
    }

    private final List<Observer> observers = new ArrayList<>();
    private State state = State.INITIALIZED;

    public State currentState() {
        return state;
    }

    public void addObserver(Observer observer) {
        observers.add(observer);
        observer.onStateChanged(state);
    }

    // Package private, so only the owner in this package drives the states.
    void moveTo(State next) {
        if (next == state) {
            return;
        }
        state = next;
        new ArrayList<>(observers).forEach(observer -> observer.onStateChanged(next));
    }
}

com.androidinterview.viewmodel.scope.CloseableTaskScope.java

package com.androidinterview.viewmodel.scope;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.Executor;
import java.util.function.Consumer;

// A scope of background work that can be shut down in one call. No coroutines
// on purpose, because the shape is the point and the shape is the same. Work is
// launched into it, closing it cancels everything, and it is an AutoCloseable so
// a ViewModel can hold it without knowing what is inside.
public final class CloseableTaskScope implements AutoCloseable {

    // A handle on one piece of work. Cancellation is cooperative, as it is for a
    // coroutine. Work that has not started never runs, and work that has sees
    // this go false at its next check.
    public static final class Job {

        private volatile boolean active = true;

        public boolean isActive() {
            return active;
        }

        public void cancel() {
            active = false;
        }
    }

    private final Executor executor;
    private final List<Job> jobs = new ArrayList<>();
    private boolean closed;

    public CloseableTaskScope() {
        this(Runnable::run);
    }

    public CloseableTaskScope(Executor executor) {
        this.executor = executor;
    }

    public Job launch(Consumer<Job> work) {
        Job job = new Job();
        synchronized (this) {
            if (closed) {
                job.cancel();
            } else {
                jobs.add(job);
            }
        }
        if (job.isActive()) {
            executor.execute(() -> {
                if (job.isActive()) {
                    work.accept(job);
                }
            });
        }
        return job;
    }

    @Override
    public void close() {
        List<Job> running;
        synchronized (this) {
            if (closed) {
                return;
            }
            closed = true;
            running = new ArrayList<>(jobs);
            jobs.clear();
        }
        for (Job job : running) {
            job.cancel();
        }
    }
}

com.androidinterview.viewmodel.scope.ViewModelScope.java

package com.androidinterview.viewmodel.scope;

import com.androidinterview.viewmodel.ViewModel;

// The wiring behind viewModelScope, in one method. The scope is stored on the
// model under a known key the first time anybody asks, and it is closed by the
// model's own clear, so nobody has to remember to cancel anything. The real one
// is a property holding a coroutine scope. The mechanism is identical.
public final class ViewModelScope {

    private static final String KEY = "com.androidinterview.viewmodel.scope.ViewModelScope.JOB_KEY";

    private ViewModelScope() { }

    public static CloseableTaskScope of(ViewModel model) {
        AutoCloseable existing = model.getCloseable(KEY);
        if (existing instanceof CloseableTaskScope scope) {
            return scope;
        }
        // addCloseable hands back whichever scope actually landed in the map,
        // so two threads racing here still end up sharing one.
        return model.addCloseable(KEY, new CloseableTaskScope());
    }
}

Kotlin

com.androidinterview.viewmodel.ViewModel.kt

package com.androidinterview.viewmodel

// The base class. Two things make it a ViewModel rather than a plain object.
// It has one hook that runs exactly once, at the end of its life, and a bag of
// AutoCloseable resources closed at that same moment. Every scope and every
// subscription a real ViewModel holds is one of those closeables, so there is a
// single place that has to be right.
abstract class ViewModel {

    private val closeables = linkedMapOf<String, AutoCloseable>()
    private var anonymous = 0
    private var cleared = false

    // Called once, when the store that owns this model lets it go. Never on a
    // rotation. Override it to drop anything that is not already a closeable.
    protected open fun onCleared() = Unit

    fun addCloseable(closeable: AutoCloseable) {
        addCloseable("anonymous:${anonymous++}", closeable)
    }

    // The keyed form, and the one viewModelScope uses. First caller wins, so a
    // scope is created once even if two threads ask at the same moment.
    @Suppress("UNCHECKED_CAST")
    fun <T : AutoCloseable> addCloseable(key: String, closeable: T): T = synchronized(this) {
        // Registering on an already cleared model is a race, not a bug. Close
        // the resource straight away rather than leak it.
        if (cleared) closeable.also { it.closeQuietly() }
        else closeables.getOrPut(key) { closeable } as T
    }

    fun getCloseable(key: String): AutoCloseable? = synchronized(this) { closeables[key] }

    // internal on purpose. Only the store holding this model may end its life,
    // so no screen can clear a model another screen is still using.
    internal fun clear() {
        val toClose = synchronized(this) {
            if (cleared) return
            cleared = true
            closeables.values.toList().also { closeables.clear() }
        }
        toClose.forEach { it.closeQuietly() }
        onCleared()
    }
}

// One resource failing to close must not strand the rest.
private fun AutoCloseable.closeQuietly() {
    runCatching { close() }
}

com.androidinterview.viewmodel.ViewModelProvider.kt

package com.androidinterview.viewmodel

// Look up by key, and only build when the lookup misses. That is the entire
// class. It owns no models and holds no cache of its own, so constructing a
// second provider over the same owner changes nothing.
class ViewModelProvider(
    owner: ViewModelStoreOwner,
    private val factory: Factory = NewInstanceFactory,
) {

    private val store = owner.viewModelStore

    // Creation is a separate job from lookup, because a model with constructor
    // arguments still has to be found in the store on the way back. In Kotlin
    // this is usually a lambda, never a class.
    interface Factory {
        fun <T : ViewModel> create(type: Class<T>): T
    }

    // The fallback for a model that needs nothing.
    object NewInstanceFactory : Factory {
        override fun <T : ViewModel> create(type: Class<T>): T =
            runCatching { type.getDeclaredConstructor().newInstance() }
                .getOrElse { throw IllegalArgumentException("${type.name} needs a no argument constructor", it) }
    }

    // The key defaults to one derived from the class, so a screen gets one model
    // of each type without anybody naming anything. Pass a key and two models of
    // the same class live side by side.
    fun <T : ViewModel> get(type: Class<T>, key: String = defaultKey(type)): T {
        store[key]?.let { existing ->
            require(type.isInstance(existing)) { "key $key already holds a ${existing.javaClass.name}" }
            return type.cast(existing)
        }
        return factory.create(type).also { store[key] = it }
    }

    companion object {
        private const val PREFIX = "com.androidinterview.viewmodel.ViewModelProvider.DefaultKey"

        fun defaultKey(type: Class<*>): String = "$PREFIX:${type.name}"
    }
}

// Reified generics remove the Class argument the Java side has to carry.
inline fun <reified T : ViewModel> ViewModelProvider.get(key: String? = null): T =
    get(T::class.java, key ?: ViewModelProvider.defaultKey(T::class.java))

// The delegate, and the reason Kotlin reads better here. Written as
// `private val model: CartViewModel by viewModels()`, the lookup happens on
// first touch and the owner is whatever the screen is.
inline fun <reified T : ViewModel> ViewModelStoreOwner.viewModels(
    key: String? = null,
    crossinline factory: () -> ViewModelProvider.Factory = { ViewModelProvider.NewInstanceFactory },
): Lazy<T> = lazy { ViewModelProvider(this, factory()).get<T>(key) }

com.androidinterview.viewmodel.ViewModelStore.kt

package com.androidinterview.viewmodel

// One method, and it is the seam the whole design turns on. Anything that can
// hand back a store can own ViewModels, so the provider never has to know
// whether it is talking to an activity, a fragment or a test double.
interface ViewModelStoreOwner {

    val viewModelStore: ViewModelStore
}

// A map from key to instance, and nothing else. This is the whole of the
// mechanism people think is magic. It is deliberately not a singleton and not
// an object, because there is one per screen and its life is the screen's life.
class ViewModelStore {

    private val models = mutableMapOf<String, ViewModel>()

    operator fun get(key: String): ViewModel? = models[key]

    // Replacing under a live key would silently leak the old model's resources,
    // so the one being displaced is ended here.
    operator fun set(key: String, model: ViewModel) {
        models.put(key, model)?.clear()
    }

    // The only way a model ever dies. Called when the owner is finishing for
    // real, never when it is being recreated.
    fun clear() {
        models.values.forEach(ViewModel::clear)
        models.clear()
    }
}

com.androidinterview.viewmodel.host.ActivityThread.kt

package com.androidinterview.viewmodel.host

// The framework side, and the reason any of this survives a rotation. It lives
// in the process rather than in a screen, so it can hold the retained slot in
// the gap where the old activity object is gone and the new one does not exist
// yet. It is also the whole answer to process death. Kill the process and this
// map goes with it, no callback runs, and nothing was ever written down.
class ActivityThread {

    private class Record(val builder: () -> ComponentActivity) {
        lateinit var current: ComponentActivity
    }

    private val records = mutableMapOf<String, Record>()

    fun launch(id: String, builder: () -> ComponentActivity): ComponentActivity =
        Record(builder).also { record ->
            record.current = builder().apply { attach(null) }
            records[id] = record
        }.current

    // A rotation. The old screen is destroyed and a brand new one is built, and
    // the only thing that crosses the gap is the retained slot.
    fun configurationChange(id: String): ComponentActivity {
        val record = recordFor(id)
        val retained = record.current.retainNonConfigurationInstances()
        record.current.performDestroy(changingConfigurations = true)
        record.current = record.builder().apply { attach(retained) }
        return record.current
    }

    // A real finish, a back press or a call to finish. Nothing is retained, so
    // the destroy clears the store and every model gets onCleared.
    fun finish(id: String) {
        recordFor(id).current.performDestroy(changingConfigurations = false)
        records -= id
    }

    // Process death, modelled honestly. No destroy runs, no store is cleared,
    // and every model in the process is simply gone.
    fun killProcess() = records.clear()

    private fun recordFor(id: String) = records[id] ?: error("no activity with id $id")
}

com.androidinterview.viewmodel.host.ComponentActivity.kt

package com.androidinterview.viewmodel.host

import com.androidinterview.viewmodel.ViewModelStore
import com.androidinterview.viewmodel.ViewModelStoreOwner

// The minimum lifecycle the design needs. What matters is that something can
// say the owner has reached its end, once, and that the store can hear it.
class Lifecycle {

    enum class State { INITIALIZED, CREATED, STARTED, RESUMED, DESTROYED }

    fun interface Observer {
        fun onStateChanged(state: State)
    }

    private val observers = mutableListOf<Observer>()

    var currentState: State = State.INITIALIZED
        private set

    fun addObserver(observer: Observer) {
        observers += observer
        observer.onStateChanged(currentState)
    }

    internal fun moveTo(next: State) {
        if (next == currentState) return
        currentState = next
        observers.toList().forEach { it.onStateChanged(next) }
    }
}

// The screen. It owns a store, hands that store to the framework on the way
// out, and clears it only when it is going away for good. It never keeps a
// static reference to anything and it never asks whether the device rotated. It
// answers one question, am I finishing or am I being rebuilt.
open class ComponentActivity : ViewModelStoreOwner {

    // The retained slot. In real Android this also carries retained fragments
    // and a custom object, and the only field that matters here is the store.
    data class NonConfigurationInstances(val viewModelStore: ViewModelStore)

    val lifecycle = Lifecycle()

    private var store: ViewModelStore? = null

    var isChangingConfigurations: Boolean = false
        private set

    init {
        // The one line the whole question is about. A destroy that is part of a
        // rebuild leaves the store alone, because the next instance is about to
        // pick it up. Any other destroy ends every model in it.
        lifecycle.addObserver { state ->
            if (state == Lifecycle.State.DESTROYED && !isChangingConfigurations) store?.clear()
        }
    }

    // Created on first use, so a screen that never asks for a ViewModel never
    // allocates a store and never has one retained for it.
    override val viewModelStore: ViewModelStore
        get() = store ?: ViewModelStore().also { store = it }

    protected open fun onCreate() = Unit

    protected open fun onDestroy() = Unit

    // The three calls below belong to the framework side and nothing else.
    internal fun attach(last: NonConfigurationInstances?) {
        last?.let { store = it.viewModelStore }
        onCreate()
        lifecycle.moveTo(Lifecycle.State.RESUMED)
    }

    internal fun retainNonConfigurationInstances(): NonConfigurationInstances? =
        store?.let(::NonConfigurationInstances)

    internal fun performDestroy(changingConfigurations: Boolean) {
        isChangingConfigurations = changingConfigurations
        onDestroy()
        lifecycle.moveTo(Lifecycle.State.DESTROYED)
    }
}

com.androidinterview.viewmodel.host.Fragment.kt

package com.androidinterview.viewmodel.host

import com.androidinterview.viewmodel.ViewModel
import com.androidinterview.viewmodel.ViewModelProvider
import com.androidinterview.viewmodel.ViewModelStore
import com.androidinterview.viewmodel.ViewModelStoreOwner
import com.androidinterview.viewmodel.get

// A ViewModel whose only job is to hold other stores, one per fragment tag.
// Fragments have no retained slot of their own, so this is the trick that gives
// them one. Their stores live inside a model in the activity's store, so they
// ride the same rotation and are cleared by the same call.
class FragmentStoreHolder : ViewModel() {

    private val stores = mutableMapOf<String, ViewModelStore>()

    fun storeFor(tag: String): ViewModelStore = stores.getOrPut(tag) { ViewModelStore() }

    fun remove(tag: String) {
        stores.remove(tag)?.clear()
    }

    override fun onCleared() {
        stores.values.forEach(ViewModelStore::clear)
        stores.clear()
    }
}

// A fragment is a second owner living inside the first, and that sentence is
// the whole of ViewModel sharing. Ask this fragment for a store and you get one
// private to it. Ask the activity and every fragment gets the same one, so a
// shared model is a choice of which owner you ask, not a feature.
class Fragment(private val tag: String) : ViewModelStoreOwner {

    private var host: ComponentActivity? = null

    fun attachTo(activity: ComponentActivity) {
        host = activity
    }

    fun requireActivity(): ComponentActivity = checkNotNull(host) { "fragment $tag is not attached" }

    override val viewModelStore: ViewModelStore
        get() = holder().storeFor(tag)

    // Popped off the back stack rather than rotated. The fragment's own store is
    // cleared here, which is why a fragment scoped model dies with the fragment
    // while an activity scoped one carries on.
    fun removeForGood() {
        holder().remove(tag)
        host = null
    }

    private fun holder() = ViewModelProvider(requireActivity()).get<FragmentStoreHolder>()
}

// The shared lookup, written the way androidx writes it. Two fragments that go
// through this get the same instance, because there is one activity store.
inline fun <reified T : ViewModel> Fragment.activityViewModels(key: String? = null): Lazy<T> =
    lazy { ViewModelProvider(requireActivity()).get<T>(key) }

com.androidinterview.viewmodel.scope.CloseableTaskScope.kt

package com.androidinterview.viewmodel.scope

import java.util.concurrent.Executor

import com.androidinterview.viewmodel.ViewModel

// A scope of background work that can be shut down in one call. No coroutines
// on purpose, because the shape is the point and the shape is the same. Work is
// launched into it, closing it cancels everything, and it is an AutoCloseable so
// a ViewModel can hold it without knowing what is inside.
class CloseableTaskScope(private val executor: Executor = Executor(Runnable::run)) : AutoCloseable {

    // A handle on one piece of work. Cancellation is cooperative, as it is for a
    // coroutine. Work that has not started never runs, and work that has sees
    // this go false at its next check.
    class Job {

        @Volatile
        var isActive: Boolean = true
            private set

        fun cancel() {
            isActive = false
        }
    }

    private val jobs = mutableListOf<Job>()
    private var closed = false

    fun launch(work: (Job) -> Unit): Job {
        val job = Job()
        synchronized(this) { if (closed) job.cancel() else jobs += job }
        if (job.isActive) executor.execute { if (job.isActive) work(job) }
        return job
    }

    override fun close() {
        val running = synchronized(this) {
            if (closed) return
            closed = true
            jobs.toList().also { jobs.clear() }
        }
        running.forEach(Job::cancel)
    }
}

private const val SCOPE_KEY = "com.androidinterview.viewmodel.scope.CloseableTaskScope.JOB_KEY"

// The wiring behind viewModelScope, in one property. The scope is stored on the
// model under a known key the first time anybody asks, and it is closed by the
// model's own clear, so nobody has to remember to cancel anything. The real one
// holds a coroutine scope. The mechanism is identical.
val ViewModel.viewModelScope: CloseableTaskScope
    get() = getCloseable(SCOPE_KEY) as CloseableTaskScope?
        ?: addCloseable(SCOPE_KEY, CloseableTaskScope())

Concurrency and edge cases

The provider is not thread safe, and that is a decision rather than an oversight. Get is a check and then an insert, so two threads racing on the same key can both miss and both create. Real ViewModels are obtained on the main thread, so the honest answer is to say that the store is confined to one thread and to name the confinement rather than to lock a map nobody contends. If you were pushed, one lock around the whole of get is enough, and it must cover the factory call as well as the lookup, otherwise you have moved the race rather than fixed it.

Closeables are thread safe, because those genuinely are not confined. Background work can register a resource at the same moment the screen is finishing. Two cases have to be right. Registering after clear closes the resource immediately instead of leaking it, and registering the same key twice hands back the one that is already there instead of quietly replacing it.

The rest of the cases, each in a sentence.

  • Clearing twice. Guarded by a flag on the model, so onCleared runs once even if a store is cleared by two paths.
  • Putting a second model under a live key. The one being displaced is cleared, because leaving it in the graph with an open scope is a leak that will not show up until much later.
  • Asking for the wrong type under an existing key. An error, not a replacement. Silently swapping types here would be a bug that surfaces three screens away.
  • A model that throws in onCleared. One resource failing to close must not strand the rest, so closes are individually guarded and the hook runs after them.
  • A screen that never asks for a model. No store is ever allocated and nothing is retained, which is why this design costs nothing on screens that do not use it.
  • A fragment popped off the back stack. Its own store is cleared, while the Activity scoped model it was sharing carries on, which is exactly the difference people reach for shared models to get.
  • Rotating during in flight work. The work keeps running, because its scope belongs to the model and the model did not move. This is the actual reason ViewModels exist, and it is worth ending on.

Watch