util.Backoff.kt Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of experiment-jvm-server Show documentation
Amplitude Experiment server-side SDK for JVM (Java, Kotlin)
The newest version!
package com.amplitude.experiment.util

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

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

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

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,
        retry: ((Throwable) -> Boolean)?
    ): CompletableFuture {
        backoff(0, config.min, function, retry)
        return completableFuture
    }

    private fun backoff(
        attempt: Int,
        delay: Long,
        function: () -> CompletableFuture,
        retry: ((Throwable) -> Boolean)? = null
    ) {
        Experiment.scheduler.schedule({
            if (completableFuture.isCancelled) {
                return@schedule
            }
            function.invoke().whenComplete { result, t ->
                if (t != null || result == null) {
                    val unwrapped = when (t) {
                        is CompletionException -> t.cause ?: t
                        else -> t
                    }
                    val shouldRetry = retry?.invoke(unwrapped) ?: true
                    val nextAttempt = attempt + 1
                    if (shouldRetry && nextAttempt < config.attempts) {
                        val nextDelay = min(delay * config.scalar, config.max.toDouble()).toLong()
                        val jitter = Random.nextLong(
                            (nextDelay - (nextDelay * 0.1).toLong()) * -1,
                            nextDelay + (nextDelay + 0.1).toLong()
                        )
                        backoff(nextAttempt, nextDelay + jitter, function, retry)
                    } else {
                        completableFuture.completeExceptionally(unwrapped)
                    }
                } else {
                    completableFuture.complete(result)
                }
            }
        }, delay, TimeUnit.MILLISECONDS)
    }
}