util.Backoff.kt Maven / Gradle / Ivy

Go to download
package com.amplitude.experiment.util

import com.amplitude.experiment.Experiment
import java.util.concurrent.CompletableFuture
import java.util.concurrent.TimeUnit
import kotlin.math.min

internal fun  backoff(
    config: BackoffConfig,
    function: () -> CompletableFuture,
): CompletableFuture {
    return Backoff(config).start(function)
}

internal data class BackoffConfig(
    val attempts: Int,
    val min: Long,
    val max: Long,
    val scalar: Double,
)

private class Backoff(
    private val config: BackoffConfig,
) {

    private val completableFuture = CompletableFuture()

    fun start(function: () -> CompletableFuture): CompletableFuture {
        backoff(0, config.min, function)
        return completableFuture
    }

    private fun backoff(
        attempt: Int,
        delay: Long,
        function: () -> CompletableFuture
    ) {
        Experiment.scheduler.schedule({
            if (completableFuture.isCancelled) {
                return@schedule
            }
            function.invoke().whenComplete { variants, t ->
                if (t != null || variants == null) {
                    // Retry the request function
                    val nextAttempt = attempt + 1
                    if (nextAttempt < config.attempts) {
                        val nextDelay = min(delay * config.scalar, config.max.toDouble()).toLong()
                        backoff(nextAttempt, nextDelay, function)
                    } else {
                        completableFuture.completeExceptionally(t)
                    }
                } else {
                    completableFuture.complete(variants)
                }
            }
        }, delay, TimeUnit.MILLISECONDS)
    }
}