Repository: ktorio/ktor
File: ktor-client/ktor-client-plugins/ktor-client-auth/common/src/io/ktor/client/plugins/auth/providers/AuthTokenHolder.kt (Kotlin)

Code before:
/*
 * Copyright 2014-2021 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
 */

package io.ktor.client.plugins.auth.providers

import io.ktor.util.collections.*
import io.ktor.utils.io.*
import io.ktor.utils.io.locks.*
import kotlinx.atomicfu.*
import kotlinx.atomicfu.locks.synchronized
import kotlinx.coroutines.*
import kotlinx.coroutines.CancellationException

/**
 * Thread-safe
 */
internal class AuthTokenHolder<T>(private val loadTokens: suspend () -> T?) {
    private val ref = atomic<T?>(null)

    @OptIn(InternalAPI::class) private val lock = SynchronizedObject()

    private val loadJobs = ConcurrentSet<Deferred<T?>>()
    private val setJobs = ConcurrentSet<Deferred<T?>>()

    private object JobCancelledException : CancellationException("Job cancelled")

    /**
     * Returns a cached reference if any. Otherwise, computes a value using [loadTokens] and caches it.
     * If there are loadToken coroutines in-progress, the first resumed one wins and the other ones are cancelled.
     */
    @OptIn(InternalAPI::class)
    internal suspend fun loadToken(): T? {
        if (ref.value != null) return ref.value

        return coroutineScope {
            val deferred = async {
                loadTokens()
            }

            loadJobs.add(deferred)


Review comment (suggestion):
If we're using a lock, why not just use it for synchronization in general?

Something like...

```
var job: Job? = null
var ref: T? = null

suspend fun loadToken() {
    job?.cancelAndJoin() // cancel previous
    
    // do work with lock
    synchronized(lock) {
         job = launch {
             ref = block()
         }
         job.join()
    }
}
```

You could have two coroutines waiting for the lock at the same time, but they'll still run in sequence.

Code after:
/*
 * Copyright 2014-2021 JetBrains s.r.o and contributors. Use of this source code is governed by the Apache 2.0 license.
 */

package io.ktor.client.plugins.auth.providers

import kotlinx.coroutines.DelicateCoroutinesApi
import kotlinx.coroutines.GlobalScope
import kotlinx.coroutines.launch
import kotlinx.coroutines.sync.Mutex
import kotlinx.coroutines.sync.withLock
import kotlinx.coroutines.withContext
import kotlin.concurrent.Volatile
import kotlin.coroutines.CoroutineContext
import kotlin.coroutines.coroutineContext

internal class AuthTokenHolder<T>(private val loadTokens: suspend () -> T?) {

    @Volatile private var value: T? = null

    @Volatile private var isLoadRequest = false

    private val mutex = Mutex()

    /**
     * Exist only for testing
     */
    internal fun get(): T? = value

    /**
     * Returns a cached value if any. Otherwise, computes a value using [loadTokens] and caches it.
     * Only one [loadToken] call can be executed at a time. The other calls are suspended and have no effect on the cached value.
     */
    internal suspend fun loadToken(): T? {
        if (value != null) return value // Hot path
        val prevValue = value

        return if (coroutineContext[SetTokenContext] != null) { // Already locked by setToken
            value = loadTokens()
            value
        } else {
            mutex.withLock {