TimerScope.kt

/*
 * Copyright 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 androidx.glance.session

import java.util.concurrent.atomic.AtomicReference
import kotlin.time.Duration
import kotlin.time.Duration.Companion.milliseconds
import kotlinx.coroutines.CancellationException
import kotlinx.coroutines.CoroutineScope
import kotlinx.coroutines.Job
import kotlinx.coroutines.cancel
import kotlinx.coroutines.coroutineScope
import kotlinx.coroutines.delay
import kotlinx.coroutines.launch

internal class TimeoutCancellationException(
    override val message: String,
    internal val block: Int,
) : CancellationException(message) {
    override fun toString() = "TimeoutCancellationException($message, $block)"
    override fun fillInStackTrace() = this
}

/**
 * This interface is similar to [kotlin.time.TimeSource], which is still marked experimental.
 */
internal fun interface TimeSource {
    /**
     * Current time in milliseconds.
     */
    fun markNow(): Long

    companion object {
        val Monotonic = TimeSource { System.currentTimeMillis() }
    }
}

/**
 * TimerScope is a CoroutineScope that allows setting an adjustable timeout for all of the
 * coroutines in the scope.
 */
internal interface TimerScope : CoroutineScope {
    /**
     * Amount of time left before this timer cancels the scope. This is not valid before
     * [startTimer] is called.
     */
    val timeLeft: Duration

    /**
     * Start the timer with an [initialTimeout].
     *
     * Once the [initialTimeout] has passed, the scope is cancelled. If [startTimer] is called again
     * while the timer is running, it will reset the timer if [initialTimeout] is less than
     * [timeLeft]. If [initialTimeout] is larger than [timeLeft], the current timer is kept.
     *
     * In order to extend the deadline, call [addTime].
     */
    fun startTimer(initialTimeout: Duration)

    /**
     *  Shift the deadline for this timer forward by [time].
     */
    fun addTime(time: Duration)
}

internal suspend fun <T> withTimer(
    timeSource: TimeSource = TimeSource.Monotonic,
    block: suspend TimerScope.() -> T,
): T = coroutineScope {
    val timerScope = this
    val timerJob: AtomicReference<Job?> = AtomicReference(null)
    coroutineScope {
        val blockScope = object : TimerScope, CoroutineScope by this {
            override val timeLeft: Duration
                get() = (deadline.get()?.minus(timeSource.markNow()))?.milliseconds
                    ?: Duration.INFINITE
            private val deadline: AtomicReference<Long?> = AtomicReference(null)

            override fun addTime(time: Duration) {
                deadline.update {
                    checkNotNull(it) { "Start the timer with startTimer before calling addTime" }
                    require(time.isPositive()) { "Cannot call addTime with a negative duration" }
                    it + time.inWholeMilliseconds
                }
            }

            override fun startTimer(initialTimeout: Duration) {
                if (initialTimeout.inWholeMilliseconds <= 0) {
                    timerScope.cancel(
                        TimeoutCancellationException("Timed out immediately", block.hashCode())
                    )
                    return
                }
                if (timeLeft < initialTimeout) return

                deadline.set(timeSource.markNow() + initialTimeout.inWholeMilliseconds)
                // Loop until the deadline is reached.
                timerJob.getAndSet(
                    timerScope.launch {
                        while (deadline.get()!! > timeSource.markNow()) {
                            delay(timeLeft)
                        }
                        timerScope.cancel(
                            TimeoutCancellationException(
                                "Timed out of executing block.",
                                block.hashCode()
                            )
                        )
                    }
                )?.cancel()
            }
        }
        blockScope.block()
    }.also {
        timerJob.get()?.cancel()
    }
}

internal suspend fun <T> withTimerOrNull(
    timeSource: TimeSource = TimeSource.Monotonic,
    block: suspend TimerScope.() -> T,
): T? = try {
    withTimer(timeSource, block)
} catch (e: TimeoutCancellationException) {
    // Return null if it's our exception, else propagate it upstream in case there are nested
    // withTimers
    if (e.block == block.hashCode()) null else throw e
}

// Update the value of the AtomicReference using the given updater function. Will throw an error
// if unable to successfully set the value.
private fun <T> AtomicReference<T>.update(updater: (T) -> T) {
    while (true) {
        get().let {
            if (compareAndSet(it, updater(it))) return
        }
    }
}