Merge "Improve unfold related logging in perfetto traces" into main

12 files changed, 476 insertions, 10 deletions

diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/DeviceStateRepository.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/DeviceStateRepository.kt
new file mode 100644
index 000000000000..f219cece3ff8
--- /dev/null
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/DeviceStateRepository.kt
@@ -0,0 +1,50 @@
+/*
+ * Copyright (C) 2023 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.systemui.unfold.system
+
+import com.android.systemui.unfold.dagger.UnfoldMain
+import com.android.systemui.unfold.updates.FoldProvider
+import com.android.systemui.unfold.updates.FoldProvider.FoldCallback
+import java.util.concurrent.Executor
+import javax.inject.Inject
+import javax.inject.Singleton
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.buffer
+import kotlinx.coroutines.flow.callbackFlow
+
+/** Provides whether the device is folded. */
+interface DeviceStateRepository {
+    val isFolded: Flow<Boolean>
+}
+
+@Singleton
+class DeviceStateRepositoryImpl
+@Inject
+constructor(
+    private val foldProvider: FoldProvider,
+    @UnfoldMain private val executor: Executor,
+) : DeviceStateRepository {
+
+    override val isFolded: Flow<Boolean>
+        get() =
+            callbackFlow {
+                    val callback = FoldCallback { isFolded -> trySend(isFolded) }
+                    foldProvider.registerCallback(callback, executor)
+                    awaitClose { foldProvider.unregisterCallback(callback) }
+                }
+                .buffer(capacity = Channel.CONFLATED)
+}
diff --git a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt
index fe607e16661c..7b67e3f3c920 100644
--- a/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt
+++ b/packages/SystemUI/shared/src/com/android/systemui/unfold/system/SystemUnfoldSharedModule.kt
@@ -48,6 +48,9 @@ abstract class SystemUnfoldSharedModule {
     abstract fun foldState(provider: DeviceStateManagerFoldProvider): FoldProvider
 
     @Binds
+    abstract fun deviceStateRepository(provider: DeviceStateRepositoryImpl): DeviceStateRepository
+
+    @Binds
     @UnfoldMain
     abstract fun mainExecutor(@Main executor: Executor): Executor
 
diff --git a/packages/SystemUI/shared/src/com/android/systemui/util/TraceStateLogger.kt b/packages/SystemUI/shared/src/com/android/systemui/util/TraceStateLogger.kt
new file mode 100644
index 000000000000..63ea1165ee04
--- /dev/null
+++ b/packages/SystemUI/shared/src/com/android/systemui/util/TraceStateLogger.kt
@@ -0,0 +1,54 @@
+/*
+ * Copyright (C) 2023 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.systemui.util
+
+import android.os.Trace
+
+/**
+ * Utility class used to log state changes easily in a track with a custom name.
+ *
+ * Example of usage:
+ * ```kotlin
+ * class MyClass {
+ *    val screenStateLogger = TraceStateLogger("Screen state")
+ *
+ *    fun onTurnedOn() { screenStateLogger.log("on") }
+ *    fun onTurnedOff() { screenStateLogger.log("off") }
+ * }
+ * ```
+ *
+ * This creates a new slice in a perfetto trace only if the state is different than the previous
+ * one.
+ */
+class TraceStateLogger(
+    private val trackName: String,
+    private val logOnlyIfDifferent: Boolean = true,
+    private val instantEvent: Boolean = true
+) {
+
+    private var previousValue: String? = null
+
+    /** If needed, logs the value to a track with name [trackName]. */
+    fun log(newValue: String) {
+        if (instantEvent) {
+            Trace.instantForTrack(Trace.TRACE_TAG_APP, trackName, newValue)
+        }
+        if (logOnlyIfDifferent && previousValue == newValue) return
+        Trace.asyncTraceForTrackEnd(Trace.TRACE_TAG_APP, trackName, 0)
+        Trace.asyncTraceForTrackBegin(Trace.TRACE_TAG_APP, trackName, newValue, 0)
+        previousValue = newValue
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLatencyTracker.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLatencyTracker.kt
index 9269df31e37e..8c66c2f0fab3 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLatencyTracker.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldLatencyTracker.kt
@@ -19,6 +19,7 @@ package com.android.systemui.unfold
 import android.content.ContentResolver
 import android.content.Context
 import android.hardware.devicestate.DeviceStateManager
+import android.os.Trace
 import android.util.Log
 import com.android.internal.util.LatencyTracker
 import com.android.systemui.dagger.SysUISingleton
@@ -57,6 +58,7 @@ constructor(
     private var folded: Boolean? = null
     private var isTransitionEnabled: Boolean? = null
     private val foldStateListener = FoldStateListener(context)
+    private var unfoldInProgress = false
     private val isFoldable: Boolean
         get() =
             context.resources
@@ -95,7 +97,7 @@ constructor(
         // the unfold animation (e.g. it could be disabled because of battery saver).
         // When animation is enabled finishing of the tracking will be done in onTransitionStarted.
         if (folded == false && isTransitionEnabled == false) {
-            latencyTracker.onActionEnd(LatencyTracker.ACTION_SWITCH_DISPLAY_UNFOLD)
+            onUnfoldEnded()
 
             if (DEBUG) {
                 Log.d(TAG, "onScreenTurnedOn: ending ACTION_SWITCH_DISPLAY_UNFOLD")
@@ -116,7 +118,7 @@ constructor(
         }
 
         if (folded == false && isTransitionEnabled == true) {
-            latencyTracker.onActionEnd(LatencyTracker.ACTION_SWITCH_DISPLAY_UNFOLD)
+            onUnfoldEnded()
 
             if (DEBUG) {
                 Log.d(TAG, "onTransitionStarted: ending ACTION_SWITCH_DISPLAY_UNFOLD")
@@ -124,6 +126,22 @@ constructor(
         }
     }
 
+    private fun onUnfoldStarted() {
+        if (unfoldInProgress) return
+        unfoldInProgress = true
+        // As LatencyTracker might be disabled, let's also log a parallel slice to the trace to be
+        // able to debug all cases.
+        latencyTracker.onActionStart(LatencyTracker.ACTION_SWITCH_DISPLAY_UNFOLD)
+        Trace.asyncTraceBegin(Trace.TRACE_TAG_APP, UNFOLD_IN_PROGRESS_TRACE_NAME, /* cookie= */ 0)
+    }
+
+    private fun onUnfoldEnded() {
+        if (!unfoldInProgress) return
+        unfoldInProgress = false
+        latencyTracker.onActionEnd(LatencyTracker.ACTION_SWITCH_DISPLAY_UNFOLD)
+        Trace.endAsyncSection(UNFOLD_IN_PROGRESS_TRACE_NAME, 0)
+    }
+
     private fun onFoldEvent(folded: Boolean) {
         val oldFolded = this.folded
 
@@ -139,7 +157,7 @@ constructor(
             // unfolding the device.
             if (oldFolded != null && !folded) {
                 // Unfolding started
-                latencyTracker.onActionStart(LatencyTracker.ACTION_SWITCH_DISPLAY_UNFOLD)
+                onUnfoldStarted()
                 isTransitionEnabled =
                     transitionProgressProvider.isPresent && contentResolver.areAnimationsEnabled()
 
@@ -159,4 +177,5 @@ constructor(
 }
 
 private const val TAG = "UnfoldLatencyTracker"
+private const val UNFOLD_IN_PROGRESS_TRACE_NAME = "Switch displays during unfold"
 private val DEBUG = Compile.IS_DEBUG && Log.isLoggable(TAG, Log.VERBOSE)
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTraceLogger.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTraceLogger.kt
new file mode 100644
index 000000000000..ed960f31228a
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTraceLogger.kt
@@ -0,0 +1,75 @@
+/*
+ * Copyright (C) 2023 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.systemui.unfold
+
+import android.content.Context
+import android.os.Trace
+import com.android.systemui.CoreStartable
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.dagger.qualifiers.Application
+import com.android.systemui.unfold.system.DeviceStateRepository
+import com.android.systemui.unfold.updates.FoldStateRepository
+import com.android.systemui.util.TraceStateLogger
+import javax.inject.Inject
+import kotlinx.coroutines.CoroutineScope
+import kotlinx.coroutines.launch
+
+/**
+ * Logs several unfold related details in a trace. Mainly used for debugging and investigate
+ * droidfooders traces.
+ */
+@SysUISingleton
+class UnfoldTraceLogger
+@Inject
+constructor(
+    private val context: Context,
+    private val foldStateRepository: FoldStateRepository,
+    @Application private val applicationScope: CoroutineScope,
+    private val deviceStateRepository: DeviceStateRepository
+) : CoreStartable {
+    private val isFoldable: Boolean
+        get() =
+            context.resources
+                .getIntArray(com.android.internal.R.array.config_foldedDeviceStates)
+                .isNotEmpty()
+
+    override fun start() {
+        if (!isFoldable) return
+
+        applicationScope.launch {
+            val foldUpdateLogger = TraceStateLogger("FoldUpdate")
+            foldStateRepository.foldUpdate.collect { foldUpdateLogger.log(it.name) }
+        }
+
+        applicationScope.launch {
+            foldStateRepository.hingeAngle.collect {
+                Trace.traceCounter(Trace.TRACE_TAG_APP, "hingeAngle", it.toInt())
+            }
+        }
+        applicationScope.launch {
+            val foldedStateLogger = TraceStateLogger("FoldedState")
+            deviceStateRepository.isFolded.collect { isFolded ->
+                foldedStateLogger.log(
+                    if (isFolded) {
+                        "folded"
+                    } else {
+                        "unfolded"
+                    }
+                )
+            }
+        }
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
index ed3eacd27c0a..71314f1f1775 100644
--- a/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/unfold/UnfoldTransitionModule.kt
@@ -19,6 +19,7 @@ package com.android.systemui.unfold
 import android.content.Context
 import android.hardware.devicestate.DeviceStateManager
 import android.os.SystemProperties
+import com.android.systemui.CoreStartable
 import com.android.systemui.dagger.qualifiers.Application
 import com.android.systemui.dagger.qualifiers.Main
 import com.android.systemui.keyguard.LifecycleScreenStatusProvider
@@ -34,16 +35,26 @@ import com.android.systemui.unfold.util.UnfoldOnlyProgressProvider
 import com.android.systemui.unfold.util.UnfoldTransitionATracePrefix
 import com.android.systemui.util.time.SystemClockImpl
 import com.android.wm.shell.unfold.ShellUnfoldProgressProvider
+import dagger.Binds
 import dagger.Lazy
 import dagger.Module
 import dagger.Provides
+import dagger.multibindings.ClassKey
+import dagger.multibindings.IntoMap
 import java.util.Optional
 import java.util.concurrent.Executor
 import javax.inject.Named
 import javax.inject.Provider
 import javax.inject.Singleton
 
-@Module(includes = [UnfoldSharedModule::class, SystemUnfoldSharedModule::class])
+@Module(
+    includes =
+        [
+            UnfoldSharedModule::class,
+            SystemUnfoldSharedModule::class,
+            UnfoldTransitionModule.Bindings::class
+        ]
+)
 class UnfoldTransitionModule {
 
     @Provides @UnfoldTransitionATracePrefix fun tracingTagPrefix() = "systemui"
@@ -136,13 +147,22 @@ class UnfoldTransitionModule {
                 null
             }
 
-        return resultingProvider?.get()?.orElse(null)?.let {
-            unfoldProgressProvider -> UnfoldProgressProvider(unfoldProgressProvider, foldProvider)
-        } ?: ShellUnfoldProgressProvider.NO_PROVIDER
+        return resultingProvider?.get()?.orElse(null)?.let { unfoldProgressProvider ->
+            UnfoldProgressProvider(unfoldProgressProvider, foldProvider)
+        }
+            ?: ShellUnfoldProgressProvider.NO_PROVIDER
     }
 
     @Provides
     fun screenStatusProvider(impl: LifecycleScreenStatusProvider): ScreenStatusProvider = impl
+
+    @Module
+    interface Bindings {
+        @Binds
+        @IntoMap
+        @ClassKey(UnfoldTraceLogger::class)
+        fun bindUnfoldTraceLogger(impl: UnfoldTraceLogger): CoreStartable
+    }
 }
 
 const val UNFOLD_STATUS_BAR = "unfold_status_bar"
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceStateRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceStateRepositoryTest.kt
new file mode 100644
index 000000000000..4eb159103b49
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/DeviceStateRepositoryTest.kt
@@ -0,0 +1,68 @@
+/*
+ * Copyright (C) 2023 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.systemui.unfold.updates
+
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.unfold.system.DeviceStateRepositoryImpl
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.util.mockito.capture
+import com.android.systemui.util.mockito.mock
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.verify
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class DeviceStateRepositoryTest : SysuiTestCase() {
+
+    private val foldProvider = mock<FoldProvider>()
+    private val testScope = TestScope(UnconfinedTestDispatcher())
+
+    private val foldStateRepository = DeviceStateRepositoryImpl(foldProvider) { r -> r.run() }
+
+    @Test
+    fun onHingeAngleUpdate_received() =
+        testScope.runTest {
+            val flowValue = collectLastValue(foldStateRepository.isFolded)
+            val foldCallback = argumentCaptor<FoldProvider.FoldCallback>()
+
+            verify(foldProvider).registerCallback(capture(foldCallback), any())
+
+            foldCallback.value.onFoldUpdated(true)
+            assertThat(flowValue()).isEqualTo(true)
+
+            foldCallback.value.onFoldUpdated(false)
+            assertThat(flowValue()).isEqualTo(false)
+        }
+
+    @Test
+    fun onHingeAngleUpdate_unregisters() {
+        testScope.runTest {
+            val flowValue = collectLastValue(foldStateRepository.isFolded)
+
+            verify(foldProvider).registerCallback(any(), any())
+        }
+        verify(foldProvider).unregisterCallback(any())
+    }
+}
diff --git a/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/FoldStateRepositoryTest.kt b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/FoldStateRepositoryTest.kt
new file mode 100644
index 000000000000..065132300564
--- /dev/null
+++ b/packages/SystemUI/tests/src/com/android/systemui/unfold/updates/FoldStateRepositoryTest.kt
@@ -0,0 +1,78 @@
+/*
+ * Copyright (C) 2023 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.systemui.unfold.updates
+
+import android.testing.AndroidTestingRunner
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.coroutines.collectLastValue
+import com.android.systemui.unfold.updates.FoldStateRepository.FoldUpdate
+import com.android.systemui.util.mockito.argumentCaptor
+import com.android.systemui.util.mockito.capture
+import com.android.systemui.util.mockito.mock
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.TestScope
+import kotlinx.coroutines.test.UnconfinedTestDispatcher
+import kotlinx.coroutines.test.runTest
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mockito.verify
+
+@RunWith(AndroidTestingRunner::class)
+@SmallTest
+class FoldStateRepositoryTest : SysuiTestCase() {
+
+    private val foldStateProvider = mock<FoldStateProvider>()
+    private val foldUpdatesListener = argumentCaptor<FoldStateProvider.FoldUpdatesListener>()
+    private val testScope = TestScope(UnconfinedTestDispatcher())
+
+    private val foldStateRepository = FoldStateRepositoryImpl(foldStateProvider)
+    @Test
+    fun onHingeAngleUpdate_received() =
+        testScope.runTest {
+            val flowValue = collectLastValue(foldStateRepository.hingeAngle)
+
+            verify(foldStateProvider).addCallback(capture(foldUpdatesListener))
+            foldUpdatesListener.value.onHingeAngleUpdate(42f)
+
+            assertThat(flowValue()).isEqualTo(42f)
+        }
+
+    @Test
+    fun onFoldUpdate_received() =
+        testScope.runTest {
+            val flowValue = collectLastValue(foldStateRepository.foldUpdate)
+
+            verify(foldStateProvider).addCallback(capture(foldUpdatesListener))
+            foldUpdatesListener.value.onFoldUpdate(FOLD_UPDATE_START_OPENING)
+
+            assertThat(flowValue()).isEqualTo(FoldUpdate.START_OPENING)
+        }
+
+    @Test
+    fun foldUpdates_mappedCorrectly() {
+        mapOf(
+                FOLD_UPDATE_START_OPENING to FoldUpdate.START_OPENING,
+                FOLD_UPDATE_START_CLOSING to FoldUpdate.START_CLOSING,
+                FOLD_UPDATE_FINISH_HALF_OPEN to FoldUpdate.FINISH_HALF_OPEN,
+                FOLD_UPDATE_FINISH_FULL_OPEN to FoldUpdate.FINISH_FULL_OPEN,
+                FOLD_UPDATE_FINISH_CLOSED to FoldUpdate.FINISH_CLOSED
+            )
+            .forEach { (id, expected) ->
+                assertThat(FoldUpdate.fromFoldUpdateId(id)).isEqualTo(expected)
+            }
+    }
+}
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
index 5ffc094b88b3..7473ca6a6486 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/UnfoldSharedModule.kt
@@ -22,6 +22,8 @@ import com.android.systemui.unfold.progress.PhysicsBasedUnfoldTransitionProgress
 import com.android.systemui.unfold.progress.UnfoldTransitionProgressForwarder
 import com.android.systemui.unfold.updates.DeviceFoldStateProvider
 import com.android.systemui.unfold.updates.FoldStateProvider
+import com.android.systemui.unfold.updates.FoldStateRepository
+import com.android.systemui.unfold.updates.FoldStateRepositoryImpl
 import com.android.systemui.unfold.updates.hinge.EmptyHingeAngleProvider
 import com.android.systemui.unfold.updates.hinge.HingeAngleProvider
 import com.android.systemui.unfold.updates.hinge.HingeSensorAngleProvider
@@ -55,6 +57,12 @@ class UnfoldSharedModule {
     fun unfoldKeyguardVisibilityManager(
         impl: UnfoldKeyguardVisibilityManagerImpl
     ): UnfoldKeyguardVisibilityManager = impl
+
+    @Provides
+    @Singleton
+    fun foldStateRepository(
+            impl: FoldStateRepositoryImpl
+    ): FoldStateRepository = impl
 }
 
 /**
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
index 6743515c2ec7..003013e18583 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/DeviceFoldStateProvider.kt
@@ -17,7 +17,6 @@ package com.android.systemui.unfold.updates
 
 import android.content.Context
 import android.os.Handler
-import android.os.Trace
 import android.util.Log
 import androidx.annotation.FloatRange
 import androidx.annotation.VisibleForTesting
@@ -130,7 +129,6 @@ constructor(
                     "lastHingeAngleBeforeTransition: $lastHingeAngleBeforeTransition"
             )
         }
-        Trace.setCounter("DeviceFoldStateProvider#onHingeAngle", angle.toLong())
 
         val currentDirection =
                 if (angle < lastHingeAngle) FOLD_UPDATE_START_CLOSING else FOLD_UPDATE_START_OPENING
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/FoldProvider.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/FoldProvider.kt
index 6e87beeb295f..ea6786e6c4bc 100644
--- a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/FoldProvider.kt
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/FoldProvider.kt
@@ -20,7 +20,7 @@ interface FoldProvider {
     fun registerCallback(callback: FoldCallback, executor: Executor)
     fun unregisterCallback(callback: FoldCallback)
 
-    interface FoldCallback {
+    fun interface FoldCallback {
         fun onFoldUpdated(isFolded: Boolean)
     }
 }
diff --git a/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/FoldStateRepository.kt b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/FoldStateRepository.kt
new file mode 100644
index 000000000000..61b0b40a55bf
--- /dev/null
+++ b/packages/SystemUI/unfold/src/com/android/systemui/unfold/updates/FoldStateRepository.kt
@@ -0,0 +1,93 @@
+/*
+ * Copyright (C) 2023 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.systemui.unfold.updates
+
+import com.android.systemui.unfold.updates.FoldStateRepository.FoldUpdate
+import javax.inject.Inject
+import kotlinx.coroutines.channels.Channel
+import kotlinx.coroutines.channels.awaitClose
+import kotlinx.coroutines.flow.Flow
+import kotlinx.coroutines.flow.buffer
+import kotlinx.coroutines.flow.callbackFlow
+
+/**
+ * Allows to subscribe to main events related to fold/unfold process such as hinge angle update,
+ * start folding/unfolding, screen availability
+ */
+interface FoldStateRepository {
+    /** Latest fold update, as described by [FoldStateProvider.FoldUpdate]. */
+    val foldUpdate: Flow<FoldUpdate>
+
+    /** Provides the hinge angle while the fold/unfold is in progress. */
+    val hingeAngle: Flow<Float>
+
+    enum class FoldUpdate {
+        START_OPENING,
+        START_CLOSING,
+        FINISH_HALF_OPEN,
+        FINISH_FULL_OPEN,
+        FINISH_CLOSED;
+
+        companion object {
+            /** Maps the old [FoldStateProvider.FoldUpdate] to [FoldStateRepository.FoldUpdate]. */
+            fun fromFoldUpdateId(@FoldStateProvider.FoldUpdate oldId: Int): FoldUpdate {
+                return when (oldId) {
+                    FOLD_UPDATE_START_OPENING -> START_OPENING
+                    FOLD_UPDATE_START_CLOSING -> START_CLOSING
+                    FOLD_UPDATE_FINISH_HALF_OPEN -> FINISH_HALF_OPEN
+                    FOLD_UPDATE_FINISH_FULL_OPEN -> FINISH_FULL_OPEN
+                    FOLD_UPDATE_FINISH_CLOSED -> FINISH_CLOSED
+                    else -> error("FoldUpdateNotFound")
+                }
+            }
+        }
+    }
+}
+
+class FoldStateRepositoryImpl
+@Inject
+constructor(
+    private val foldStateProvider: FoldStateProvider,
+) : FoldStateRepository {
+
+    override val hingeAngle: Flow<Float>
+        get() =
+            callbackFlow {
+                    val callback =
+                        object : FoldStateProvider.FoldUpdatesListener {
+                            override fun onHingeAngleUpdate(angle: Float) {
+                                trySend(angle)
+                            }
+                        }
+                    foldStateProvider.addCallback(callback)
+                    awaitClose { foldStateProvider.removeCallback(callback) }
+                }
+                .buffer(capacity = Channel.CONFLATED)
+
+    override val foldUpdate: Flow<FoldUpdate>
+        get() =
+            callbackFlow {
+                    val callback =
+                        object : FoldStateProvider.FoldUpdatesListener {
+                            override fun onFoldUpdate(update: Int) {
+                                trySend(FoldUpdate.fromFoldUpdateId(update))
+                            }
+                        }
+                    foldStateProvider.addCallback(callback)
+                    awaitClose { foldStateProvider.removeCallback(callback) }
+                }
+                .buffer(capacity = Channel.CONFLATED)
+}