1 /*
2  * Copyright 2023 The Android Open Source Project
3  *
4  * Licensed under the Apache License, Version 2.0 (the "License");
5  * you may not use this file except in compliance with the License.
6  * You may obtain a copy of the License at
7  *
8  *      http://www.apache.org/licenses/LICENSE-2.0
9  *
10  * Unless required by applicable law or agreed to in writing, software
11  * distributed under the License is distributed on an "AS IS" BASIS,
12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13  * See the License for the specific language governing permissions and
14  * limitations under the License.
15  */
16 
17 package androidx.glance.session
18 
19 import androidx.annotation.RestrictTo
20 import androidx.annotation.RestrictTo.Scope.LIBRARY_GROUP
21 import java.util.concurrent.atomic.AtomicReference
22 import kotlin.time.Duration
23 import kotlin.time.Duration.Companion.milliseconds
24 import kotlinx.coroutines.CancellationException
25 import kotlinx.coroutines.CoroutineScope
26 import kotlinx.coroutines.Job
27 import kotlinx.coroutines.cancel
28 import kotlinx.coroutines.coroutineScope
29 import kotlinx.coroutines.delay
30 import kotlinx.coroutines.launch
31 
32 internal class TimeoutCancellationException(
33     override val message: String,
34     internal val block: Int,
35 ) : CancellationException(message) {
toStringnull36     override fun toString() = "TimeoutCancellationException($message, $block)"
37 
38     override fun fillInStackTrace() = this
39 }
40 
41 /** This interface is similar to [kotlin.time.TimeSource], which is still marked experimental. */
42 @RestrictTo(LIBRARY_GROUP)
43 fun interface TimeSource {
44     /** Current time in milliseconds. */
45     fun markNow(): Long
46 
47     companion object {
48         val Monotonic = TimeSource { System.currentTimeMillis() }
49     }
50 }
51 
52 /**
53  * TimerScope is a CoroutineScope that allows setting an adjustable timeout for all of the
54  * coroutines in the scope.
55  */
56 internal interface TimerScope : CoroutineScope {
57     /**
58      * Amount of time left before this timer cancels the scope. This is not valid before
59      * [startTimer] is called.
60      */
61     val timeLeft: Duration
62 
63     /**
64      * Start the timer with an [initialTimeout].
65      *
66      * Once the [initialTimeout] has passed, the scope is cancelled. If [startTimer] is called again
67      * while the timer is running, it will reset the timer if [initialTimeout] is less than
68      * [timeLeft]. If [initialTimeout] is larger than [timeLeft], the current timer is kept.
69      *
70      * In order to extend the deadline, call [addTime].
71      */
startTimernull72     fun startTimer(initialTimeout: Duration)
73 
74     /** Shift the deadline for this timer forward by [time]. */
75     fun addTime(time: Duration)
76 }
77 
78 internal suspend fun <T> withTimer(
79     timeSource: TimeSource = TimeSource.Monotonic,
80     block: suspend TimerScope.() -> T,
81 ): T = coroutineScope {
82     val timerScope = this
83     val timerJob: AtomicReference<Job?> = AtomicReference(null)
84     coroutineScope {
85             val blockScope =
86                 object : TimerScope, CoroutineScope by this {
87                     override val timeLeft: Duration
88                         get() =
89                             (deadline.get()?.minus(timeSource.markNow()))?.milliseconds
90                                 ?: Duration.INFINITE
91 
92                     private val deadline: AtomicReference<Long?> = AtomicReference(null)
93 
94                     override fun addTime(time: Duration) {
95                         deadline.update {
96                             checkNotNull(it) {
97                                 "Start the timer with startTimer before calling addTime"
98                             }
99                             require(time.isPositive()) {
100                                 "Cannot call addTime with a negative duration"
101                             }
102                             it + time.inWholeMilliseconds
103                         }
104                     }
105 
106                     override fun startTimer(initialTimeout: Duration) {
107                         if (initialTimeout.inWholeMilliseconds <= 0) {
108                             timerScope.cancel(
109                                 TimeoutCancellationException(
110                                     "Timed out immediately",
111                                     block.hashCode()
112                                 )
113                             )
114                             return
115                         }
116                         if (timeLeft < initialTimeout) return
117 
118                         deadline.set(timeSource.markNow() + initialTimeout.inWholeMilliseconds)
119                         // Loop until the deadline is reached.
120                         timerJob
121                             .getAndSet(
122                                 timerScope.launch {
123                                     while (deadline.get()!! > timeSource.markNow()) {
124                                         delay(timeLeft)
125                                     }
126                                     timerScope.cancel(
127                                         TimeoutCancellationException(
128                                             "Timed out of executing block.",
129                                             block.hashCode()
130                                         )
131                                     )
132                                 }
133                             )
134                             ?.cancel()
135                     }
136                 }
137             blockScope.block()
138         }
139         .also { timerJob.get()?.cancel() }
140 }
141 
withTimerOrNullnull142 internal suspend fun <T> withTimerOrNull(
143     timeSource: TimeSource = TimeSource.Monotonic,
144     block: suspend TimerScope.() -> T,
145 ): T? =
146     try {
147         withTimer(timeSource, block)
148     } catch (e: TimeoutCancellationException) {
149         // Return null if it's our exception, else propagate it upstream in case there are nested
150         // withTimers
151         if (e.block == block.hashCode()) null else throw e
152     }
153 
154 // Update the value of the AtomicReference using the given updater function. Will throw an error
155 // if unable to successfully set the value.
updatenull156 private fun <T> AtomicReference<T>.update(updater: (T) -> T) {
157     while (true) {
158         get().let { if (compareAndSet(it, updater(it))) return }
159     }
160 }
161 
noopTimernull162 internal suspend fun <T> noopTimer(
163     block: suspend TimerScope.() -> T,
164 ): T = coroutineScope {
165     val timerScope =
166         object : TimerScope, CoroutineScope by this {
167             override val timeLeft = Duration.INFINITE
168 
169             override fun startTimer(initialTimeout: Duration) {}
170 
171             override fun addTime(time: Duration) {}
172         }
173     timerScope.block()
174 }
175