Datalayer for Desktop windowing education

As the data we need is required to persist through device reboots, we can not rely on in-memory storage. Hence we are using Android ProtoDataStore.
WindowingEducationDatastoreRepository- Repository that accesses datastore by reading values and writing into it.
More details can be read at go/dw-education-design-doc.

Bug: 352552060, 352552062
Test: AppHandleEducationDatastoreRepositoryTest
Flag: com.android.window.flags.enable_desktop_windowing_app_handle_education
Change-Id: I4cd7f7a45b9ec89b56907d65acd2e6253081520f
diff --git a/libs/WindowManager/Shell/Android.bp b/libs/WindowManager/Shell/Android.bp
index 1e68241..0bf46d0 100644
--- a/libs/WindowManager/Shell/Android.bp
+++ b/libs/WindowManager/Shell/Android.bp
@@ -166,6 +166,16 @@
     },
 }
 
+java_library {
+    name: "WindowManager-Shell-lite-proto",
+
+    srcs: ["src/com/android/wm/shell/desktopmode/education/data/proto/**/*.proto"],
+
+    proto: {
+        type: "lite",
+    },
+}
+
 filegroup {
     name: "wm_shell-shared-aidls",
 
@@ -206,6 +216,7 @@
         "androidx.core_core-animation",
         "androidx.core_core-ktx",
         "androidx.arch.core_core-runtime",
+        "androidx.datastore_datastore",
         "androidx.compose.material3_material3",
         "androidx-constraintlayout_constraintlayout",
         "androidx.dynamicanimation_dynamicanimation",
@@ -217,6 +228,7 @@
         "com_android_wm_shell_flags_lib",
         "com.android.window.flags.window-aconfig-java",
         "WindowManager-Shell-proto",
+        "WindowManager-Shell-lite-proto",
         "WindowManager-Shell-shared",
         "perfetto_trace_java_protos",
         "dagger2",
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
index aa499d9..b277af4 100644
--- a/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/dagger/WMShellModule.java
@@ -66,6 +66,7 @@
 import com.android.wm.shell.desktopmode.EnterDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.ExitDesktopTaskTransitionHandler;
 import com.android.wm.shell.desktopmode.ToggleResizeDesktopTaskTransitionHandler;
+import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository;
 import com.android.wm.shell.draganddrop.DragAndDropController;
 import com.android.wm.shell.draganddrop.GlobalDragListener;
 import com.android.wm.shell.freeform.FreeformComponents;
@@ -629,6 +630,13 @@
         return new DesktopModeEventLogger();
     }
 
+    @WMSingleton
+    @Provides
+    static AppHandleEducationDatastoreRepository provideAppHandleEducationDatastoreRepository(
+            Context context) {
+        return new AppHandleEducationDatastoreRepository(context);
+    }
+
     //
     // Drag and drop
     //
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
new file mode 100644
index 0000000..bf4a2ab
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/AppHandleEducationDatastoreRepository.kt
@@ -0,0 +1,81 @@
+/*
+ * Copyright (C) 2024 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 com.android.wm.shell.desktopmode.education.data
+
+import android.content.Context
+import android.util.Log
+import androidx.datastore.core.CorruptionException
+import androidx.datastore.core.DataStore
+import androidx.datastore.core.DataStoreFactory
+import androidx.datastore.core.Serializer
+import androidx.datastore.dataStore
+import androidx.datastore.dataStoreFile
+import com.android.framework.protobuf.InvalidProtocolBufferException
+import com.android.internal.annotations.VisibleForTesting
+import java.io.InputStream
+import java.io.OutputStream
+import kotlinx.coroutines.flow.first
+
+/**
+ * Manages interactions with the App Handle education datastore.
+ *
+ * This class provides a layer of abstraction between the UI/business logic and the underlying
+ * DataStore.
+ */
+class AppHandleEducationDatastoreRepository
+@VisibleForTesting
+constructor(private val dataStore: DataStore<WindowingEducationProto>) {
+  constructor(
+      context: Context
+  ) : this(
+      DataStoreFactory.create(
+          serializer = WindowingEducationProtoSerializer,
+          produceFile = { context.dataStoreFile(APP_HANDLE_EDUCATION_DATASTORE_FILEPATH) }))
+
+  /**
+   * Reads and returns the [WindowingEducationProto] Proto object from the DataStore. If the
+   * DataStore is empty or there's an error reading, it returns the default value of Proto.
+   */
+  suspend fun windowingEducationProto(): WindowingEducationProto =
+      try {
+        dataStore.data.first()
+      } catch (e: Exception) {
+        Log.e(TAG, "Unable to read from datastore")
+        WindowingEducationProto.getDefaultInstance()
+      }
+
+  companion object {
+    private const val TAG = "AppHandleEducationDatastoreRepository"
+    private const val APP_HANDLE_EDUCATION_DATASTORE_FILEPATH = "app_handle_education.pb"
+
+    object WindowingEducationProtoSerializer : Serializer<WindowingEducationProto> {
+
+      override val defaultValue: WindowingEducationProto =
+          WindowingEducationProto.getDefaultInstance()
+
+      override suspend fun readFrom(input: InputStream): WindowingEducationProto =
+          try {
+            WindowingEducationProto.parseFrom(input)
+          } catch (exception: InvalidProtocolBufferException) {
+            throw CorruptionException("Cannot read proto.", exception)
+          }
+
+      override suspend fun writeTo(windowingProto: WindowingEducationProto, output: OutputStream) =
+          windowingProto.writeTo(output)
+    }
+  }
+}
diff --git a/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/proto/windowing_education.proto b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/proto/windowing_education.proto
new file mode 100644
index 0000000..d29ec53
--- /dev/null
+++ b/libs/WindowManager/Shell/src/com/android/wm/shell/desktopmode/education/data/proto/windowing_education.proto
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2024 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.
+ */
+
+syntax = "proto2";
+
+option java_package = "com.android.wm.shell.desktopmode.education.data";
+option java_multiple_files = true;
+
+// Desktop Windowing education data
+message WindowingEducationProto {
+  // Timestamp in milliseconds of when the education was last viewed.
+  optional int64 education_viewed_timestamp_millis = 1;
+  // Timestamp in milliseconds of when the feature was last used.
+  optional int64 feature_used_timestamp_millis = 2;
+  oneof education_data {
+    // Fields specific to app handle education
+    AppHandleEducation app_handle_education = 3;
+  }
+
+  message AppHandleEducation {
+    // Map that stores app launch count for corresponding package
+    map<string, int32> app_usage_stats = 1;
+    // Timestamp of when app_usage_stats was last cached
+    optional int64 app_usage_stats_last_update_timestamp_millis = 2;
+  }
+}
\ No newline at end of file
diff --git a/libs/WindowManager/Shell/tests/unittest/Android.bp b/libs/WindowManager/Shell/tests/unittest/Android.bp
index 92be4f9..a461d53 100644
--- a/libs/WindowManager/Shell/tests/unittest/Android.bp
+++ b/libs/WindowManager/Shell/tests/unittest/Android.bp
@@ -43,6 +43,8 @@
         "androidx.test.runner",
         "androidx.test.rules",
         "androidx.test.ext.junit",
+        "androidx.datastore_datastore",
+        "kotlinx_coroutines_test",
         "androidx.dynamicanimation_dynamicanimation",
         "dagger2",
         "frameworks-base-testutils",
diff --git a/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationDatastoreRepositoryTest.kt b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationDatastoreRepositoryTest.kt
new file mode 100644
index 0000000..4d40738
--- /dev/null
+++ b/libs/WindowManager/Shell/tests/unittest/src/com/android/wm/shell/desktopmode/education/AppHandleEducationDatastoreRepositoryTest.kt
@@ -0,0 +1,125 @@
+/*
+ * Copyright (C) 2024 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 com.android.wm.shell.desktopmode.education
+
+import android.content.Context
+import android.testing.AndroidTestingRunner
+import androidx.datastore.core.DataStore
+import androidx.datastore.core.DataStoreFactory
+import androidx.datastore.dataStoreFile
+import androidx.test.core.app.ApplicationProvider
+import androidx.test.filters.SmallTest
+import androidx.test.platform.app.InstrumentationRegistry
+import com.android.wm.shell.desktopmode.education.data.AppHandleEducationDatastoreRepository
+import com.android.wm.shell.desktopmode.education.data.WindowingEducationProto
+import com.google.common.truth.Truth.assertThat
+import java.io.File
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.Dispatchers
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import kotlinx.coroutines.SupervisorJob
+import kotlinx.coroutines.cancel
+import kotlinx.coroutines.test.StandardTestDispatcher
+import kotlinx.coroutines.test.runTest
+import kotlinx.coroutines.test.setMain
+import org.junit.After
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+
+@SmallTest
+@RunWith(AndroidTestingRunner::class)
+@ExperimentalCoroutinesApi
+class AppHandleEducationDatastoreRepositoryTest {
+  private val testContext: Context = InstrumentationRegistry.getInstrumentation().targetContext
+  private lateinit var testDatastore: DataStore<WindowingEducationProto>
+  private lateinit var datastoreRepository: AppHandleEducationDatastoreRepository
+  private lateinit var datastoreScope: CoroutineScope
+
+  @Before
+  fun setUp() {
+    Dispatchers.setMain(StandardTestDispatcher())
+    datastoreScope = CoroutineScope(Dispatchers.Unconfined + SupervisorJob())
+    testDatastore =
+        DataStoreFactory.create(
+            serializer =
+                AppHandleEducationDatastoreRepository.Companion.WindowingEducationProtoSerializer,
+            scope = datastoreScope) {
+              testContext.dataStoreFile(APP_HANDLE_EDUCATION_DATASTORE_TEST_FILE)
+            }
+    datastoreRepository = AppHandleEducationDatastoreRepository(testDatastore)
+  }
+
+  @After
+  fun tearDown() {
+    File(ApplicationProvider.getApplicationContext<Context>().filesDir, "datastore")
+        .deleteRecursively()
+
+    datastoreScope.cancel()
+  }
+
+  @Test
+  fun getWindowingEducationProto_returnsCorrectProto() =
+      runTest(StandardTestDispatcher()) {
+        val windowingEducationProto =
+            createWindowingEducationProto(
+                educationViewedTimestampMillis = 123L,
+                featureUsedTimestampMillis = 124L,
+                appUsageStats = mapOf(GMAIL_PACKAGE_NAME to 2),
+                appUsageStatsLastUpdateTimestampMillis = 125L)
+        testDatastore.updateData { windowingEducationProto }
+
+        val resultProto = datastoreRepository.windowingEducationProto()
+
+        assertThat(resultProto).isEqualTo(windowingEducationProto)
+      }
+
+  private fun createWindowingEducationProto(
+      educationViewedTimestampMillis: Long? = null,
+      featureUsedTimestampMillis: Long? = null,
+      appUsageStats: Map<String, Int>? = null,
+      appUsageStatsLastUpdateTimestampMillis: Long? = null
+  ): WindowingEducationProto =
+      WindowingEducationProto.newBuilder()
+          .apply {
+            if (educationViewedTimestampMillis != null)
+                setEducationViewedTimestampMillis(educationViewedTimestampMillis)
+            if (featureUsedTimestampMillis != null)
+                setFeatureUsedTimestampMillis(featureUsedTimestampMillis)
+            setAppHandleEducation(
+                createAppHandleEducationProto(
+                    appUsageStats, appUsageStatsLastUpdateTimestampMillis))
+          }
+          .build()
+
+  private fun createAppHandleEducationProto(
+      appUsageStats: Map<String, Int>? = null,
+      appUsageStatsLastUpdateTimestampMillis: Long? = null
+  ): WindowingEducationProto.AppHandleEducation =
+      WindowingEducationProto.AppHandleEducation.newBuilder()
+          .apply {
+            if (appUsageStats != null) putAllAppUsageStats(appUsageStats)
+            if (appUsageStatsLastUpdateTimestampMillis != null)
+                setAppUsageStatsLastUpdateTimestampMillis(appUsageStatsLastUpdateTimestampMillis)
+          }
+          .build()
+
+  companion object {
+    private const val GMAIL_PACKAGE_NAME = "com.google.android.gm"
+    private const val APP_HANDLE_EDUCATION_DATASTORE_TEST_FILE = "app_handle_education_test.pb"
+  }
+}