CoroutineScope + ViewModels

Test: ViewModelTest#testVmScope
Change-Id: I728866d62dc87a066449c8d7c4f2a085e8b41d0f
diff --git a/buildSrc/src/main/kotlin/androidx/build/dependencies/Dependencies.kt b/buildSrc/src/main/kotlin/androidx/build/dependencies/Dependencies.kt
index be2a257..a7f3958 100644
--- a/buildSrc/src/main/kotlin/androidx/build/dependencies/Dependencies.kt
+++ b/buildSrc/src/main/kotlin/androidx/build/dependencies/Dependencies.kt
@@ -40,6 +40,7 @@
 const val JUNIT = "junit:junit:4.12"
 const val KOTLIN_STDLIB = "org.jetbrains.kotlin:kotlin-stdlib:1.2.70"
 const val KOTLIN_METADATA = "me.eugeniomarletti.kotlin.metadata:kotlin-metadata:1.4.0"
+const val KOTLIN_COROUTINES = "org.jetbrains.kotlinx:kotlinx-coroutines-android:0.30.2"
 const val MOCKITO_CORE = "org.mockito:mockito-core:2.19.0"
 const val MULTIDEX = "androidx.multidex:multidex:2.0.0"
 const val NULLAWAY = "com.uber.nullaway:nullaway:0.3.7"
diff --git a/lifecycle/viewmodel-savedstate/src/main/java/androidx/lifecycle/SavedStateVMFactory.java b/lifecycle/viewmodel-savedstate/src/main/java/androidx/lifecycle/SavedStateVMFactory.java
index 257079a..a5226d8 100644
--- a/lifecycle/viewmodel-savedstate/src/main/java/androidx/lifecycle/SavedStateVMFactory.java
+++ b/lifecycle/viewmodel-savedstate/src/main/java/androidx/lifecycle/SavedStateVMFactory.java
@@ -41,7 +41,7 @@
         SavedStateHandle handle = new SavedStateHandle(mInitialArgs, savedState);
         mSavedStateStore.registerSaveStateCallback(key, handle.savedStateComponent());
         T viewmodel = mWrappedFactory.create(key, modelClass, handle);
-        viewmodel.setTag(TAG_SAVED_STATE_HANDLE, handle);
+        viewmodel.setTagIfAbsent(TAG_SAVED_STATE_HANDLE, handle);
         return viewmodel;
     }
 }
diff --git a/lifecycle/viewmodel/ktx/api/current.txt b/lifecycle/viewmodel/ktx/api/current.txt
index 183eaa0..1b04cb9 100644
--- a/lifecycle/viewmodel/ktx/api/current.txt
+++ b/lifecycle/viewmodel/ktx/api/current.txt
@@ -1,6 +1,11 @@
 // Signature format: 2.0
 package androidx.lifecycle {
 
+  public final class ViewModelKt {
+    ctor public ViewModelKt();
+    method public static kotlinx.coroutines.experimental.CoroutineScope getViewModelScope(androidx.lifecycle.ViewModel);
+  }
+
   public final class ViewModelProviderKt {
     ctor public ViewModelProviderKt();
   }
diff --git a/lifecycle/viewmodel/ktx/build.gradle b/lifecycle/viewmodel/ktx/build.gradle
index 96727de..e887a64 100644
--- a/lifecycle/viewmodel/ktx/build.gradle
+++ b/lifecycle/viewmodel/ktx/build.gradle
@@ -34,9 +34,13 @@
 dependencies {
     api(project(":lifecycle:lifecycle-viewmodel"))
     api(KOTLIN_STDLIB)
+    api(KOTLIN_COROUTINES)
 
     testImplementation(JUNIT)
     testImplementation(TEST_RUNNER)
+
+    androidTestImplementation(TRUTH)
+    androidTestImplementation(TEST_RUNNER)
 }
 
 supportLibrary {
@@ -48,3 +52,9 @@
     description = "Kotlin extensions for 'viewmodel' artifact"
     useMetalava = true
 }
+
+kotlin {
+    experimental {
+        coroutines "enable"
+    }
+}
diff --git a/lifecycle/viewmodel/ktx/src/androidTest/java/androidx/lifecycle/ViewModelTest.kt b/lifecycle/viewmodel/ktx/src/androidTest/java/androidx/lifecycle/ViewModelTest.kt
new file mode 100644
index 0000000..2f90453
--- /dev/null
+++ b/lifecycle/viewmodel/ktx/src/androidTest/java/androidx/lifecycle/ViewModelTest.kt
@@ -0,0 +1,56 @@
+/*
+ * Copyright 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.lifecycle
+
+import androidx.test.filters.SmallTest
+import com.google.common.truth.Truth
+import kotlinx.coroutines.experimental.delay
+import kotlinx.coroutines.experimental.launch
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.junit.runners.JUnit4
+
+@RunWith(JUnit4::class)
+@SmallTest
+class ViewModelTest {
+
+    @Test fun testVmScope() {
+        val vm = object : ViewModel() {}
+        val job1 = vm.viewModelScope.launch { delay(1000) }
+        val job2 = vm.viewModelScope.launch { delay(1000) }
+        vm.clear()
+        Truth.assertThat(job1.isCancelled).isTrue()
+        Truth.assertThat(job2.isCancelled).isTrue()
+    }
+
+    @Test fun testStartJobInClearedVM() {
+        val vm = object : ViewModel() {}
+        vm.clear()
+        val job1 = vm.viewModelScope.launch { delay(1000) }
+        Truth.assertThat(job1.isCancelled).isTrue()
+    }
+
+    @Test fun testSameScope() {
+        val vm = object : ViewModel() {}
+        val scope1 = vm.viewModelScope
+        val scope2 = vm.viewModelScope
+        Truth.assertThat(scope1).isSameAs(scope2)
+        vm.clear()
+        val scope3 = vm.viewModelScope
+        Truth.assertThat(scope3).isSameAs(scope2)
+    }
+}
\ No newline at end of file
diff --git a/lifecycle/viewmodel/ktx/src/main/java/androidx/lifecycle/ViewModel.kt b/lifecycle/viewmodel/ktx/src/main/java/androidx/lifecycle/ViewModel.kt
new file mode 100644
index 0000000..1d121c4
--- /dev/null
+++ b/lifecycle/viewmodel/ktx/src/main/java/androidx/lifecycle/ViewModel.kt
@@ -0,0 +1,50 @@
+/*
+ * Copyright 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.lifecycle
+
+import kotlinx.coroutines.experimental.CoroutineScope
+import kotlinx.coroutines.experimental.Dispatchers
+import kotlinx.coroutines.experimental.Job
+import kotlinx.coroutines.experimental.cancel
+import java.io.Closeable
+import kotlin.coroutines.experimental.CoroutineContext
+
+private const val JOB_KEY = "androidx.lifecycle.ViewModelCoroutineScope.JOB_KEY"
+
+/**
+ * [CoroutineScope] tied to this [ViewModel].
+ * This scope will be canceled when ViewModel will be cleared, i.e [ViewModel.onCleared] is called
+ *
+ * This scope is bound to [Dispatchers.Main]
+ */
+val ViewModel.viewModelScope: CoroutineScope
+        get() {
+            val scope: CoroutineScope? = this.getTag(JOB_KEY)
+            if (scope != null) {
+                return scope
+            }
+            return setTagIfAbsent(JOB_KEY,
+                CloseableCoroutineScope(Job() + Dispatchers.Main))
+        }
+
+internal class CloseableCoroutineScope(context: CoroutineContext) : Closeable, CoroutineScope {
+    override val coroutineContext: CoroutineContext = context
+
+    override fun close() {
+        coroutineContext.cancel()
+    }
+}
diff --git a/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModel.java b/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModel.java
index fc75941..639aa7e 100644
--- a/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModel.java
+++ b/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModel.java
@@ -19,8 +19,9 @@
 import androidx.annotation.MainThread;
 import androidx.annotation.Nullable;
 
-import java.util.HashMap;
-import java.util.Map;
+import java.io.Closeable;
+import java.io.IOException;
+import java.util.concurrent.ConcurrentHashMap;
 
 /**
  * ViewModel is a class that is responsible for preparing and managing the data for
@@ -105,7 +106,8 @@
  */
 public abstract class ViewModel {
     @Nullable
-    private Map<String, Object> mBagOfTags;
+    private final ConcurrentHashMap<String, Object> mBagOfTags = new ConcurrentHashMap<>();
+    private volatile boolean mCleared = false;
 
     /**
      * This method will be called when this ViewModel is no longer used and will be destroyed.
@@ -117,24 +119,56 @@
     protected void onCleared() {
     }
 
+    @MainThread
+    final void clear() {
+        mCleared = true;
+        for (Object value: mBagOfTags.values()) {
+            // see comment for the similar call in setTagIfAbsent
+            closeWithRuntimeException(value);
+        }
+        onCleared();
+    }
+
     /**
      * Sets a tag associated with this viewmodel and a key.
+     * If the given {@code newValue} is {@link Closeable},
+     * it will be closed once {@link #clear()}.
+     * <p>
+     * If a value was already set for the given key, this calls do nothing and
+     * returns currently associated value, the given {@code newValue} would be ignored
+     * <p>
+     * If the ViewModel was already cleared then close() would be called on the returned object if
+     * it implements {@link Closeable}. The same object may receive multiple close calls, so method
+     * should be idempotent.
      */
-    @MainThread
-    void setTag(String key, Object obj) {
-        if (mBagOfTags == null) {
-            mBagOfTags = new HashMap<>();
+    <T> T setTagIfAbsent(String key, T newValue) {
+        @SuppressWarnings("unchecked") T previous = (T) mBagOfTags.putIfAbsent(key, newValue);
+        T result = previous == null ? newValue : previous;
+        if (mCleared) {
+            // It is possible that we'll call close() multiple times on the same object, but
+            // Closeable interface requires close method to be idempotent:
+            // "if the stream is already closed then invoking this method has no effect." (c)
+            closeWithRuntimeException(result);
         }
-        mBagOfTags.put(key, obj);
+        return result;
     }
 
     /**
      * Returns the tag associated with this viewmodel and the specified key.
      */
     @SuppressWarnings("TypeParameterUnusedInFormals")
-    @MainThread
     <T> T getTag(String key) {
         //noinspection unchecked
-        return mBagOfTags != null ? (T) mBagOfTags.get(key) : null;
+        return (T) mBagOfTags.get(key);
+    }
+
+    private static void closeWithRuntimeException(Object obj) {
+        if (obj instanceof Closeable) {
+            try {
+                ((Closeable) obj).close();
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
     }
 }
diff --git a/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModelStore.java b/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModelStore.java
index 15ec1fb..e5c69f3 100644
--- a/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModelStore.java
+++ b/lifecycle/viewmodel/src/main/java/androidx/lifecycle/ViewModelStore.java
@@ -59,7 +59,7 @@
      */
     public final void clear() {
         for (ViewModel vm : mMap.values()) {
-            vm.onCleared();
+            vm.clear();
         }
         mMap.clear();
     }
diff --git a/lifecycle/viewmodel/src/test/java/androidx/lifecycle/ViewModelTest.java b/lifecycle/viewmodel/src/test/java/androidx/lifecycle/ViewModelTest.java
new file mode 100644
index 0000000..0a456b7
--- /dev/null
+++ b/lifecycle/viewmodel/src/test/java/androidx/lifecycle/ViewModelTest.java
@@ -0,0 +1,69 @@
+/*
+ * Copyright 2018 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.lifecycle;
+
+import static org.hamcrest.CoreMatchers.is;
+import static org.hamcrest.MatcherAssert.assertThat;
+import static org.junit.Assert.assertTrue;
+
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.JUnit4;
+
+import java.io.Closeable;
+
+@RunWith(JUnit4.class)
+public class ViewModelTest {
+
+    static class CloseableImpl implements Closeable {
+        boolean mWasClosed;
+        @Override
+        public void close() {
+            mWasClosed = true;
+        }
+    }
+
+    class ViewModel extends androidx.lifecycle.ViewModel {
+    }
+
+    @Test
+    public void testCloseableTag() {
+        ViewModel vm = new ViewModel();
+        CloseableImpl impl = new CloseableImpl();
+        vm.setTagIfAbsent("totally_not_coroutine_context", impl);
+        vm.clear();
+        assertTrue(impl.mWasClosed);
+    }
+
+    @Test
+    public void testCloseableTagAlreadyClearedVM() {
+        ViewModel vm = new ViewModel();
+        vm.clear();
+        CloseableImpl impl = new CloseableImpl();
+        vm.setTagIfAbsent("key", impl);
+        assertTrue(impl.mWasClosed);
+
+    }
+
+    @Test
+    public void testAlreadyAssociatedKey() {
+        ViewModel vm = new ViewModel();
+        assertThat(vm.setTagIfAbsent("key", "first"), is("first"));
+        assertThat(vm.setTagIfAbsent("key", "second"), is("first"));
+    }
+
+}