From 562f7516513946f60f33e28479c0729b24d72faa Mon Sep 17 00:00:00 2001 From: Pradeep Makineni <7807964+mpc00620@users.noreply.github.com> Date: Tue, 20 Oct 2026 22:04:57 +0000 Subject: [PATCH] Add isReady() and document lazy model loading in the Java library (#17) TaskInitializer.init() loads the model and experiment configurations and schedules the feature registry setup and the periodic model-result and bloom-filter loading. The first model load stays on the periodic schedule, after a random delay of up to one refresh period (300 s by default), as before: blocking init() on it would add up to a minute of model load time to every deployment batch on large fleets, stop bidding for that long on an all-at-once restart, and send the fleet's first S3 reads in one burst. Callers had no way to tell when the model data was in place, and the README suggested sleeping for 10 seconds. Add BidRequestEvaluatorFactory.isReady(): - true once init() has completed, the extractor and transformer registries hold every ServiceLoader provider, and every model-result and bloom-filter loading task of that TaskInitializer has finished its first periodic run, whatever the outcome (a missing file, an S3 error or a loader exception still counts, so model data never blocks readiness forever); once true it stays true - the base class method throws UnsupportedOperationException, so existing subclasses keep compiling; DefaultBidRequestEvaluatorFactory overrides it Supporting changes, all additive: - PeriodicTaskWithRandomizedStart.hasFinishedFirstRun(), a volatile flag set in a finally block of the scheduled run - a TaskInitializer constructor taking a callback that init() runs once both stages have completed; init() now also throws when a stage ends with tasks still pending, instead of returning normally - getTask() on InitializerTaskOnPeriodicTask and ModelResultPeriodicLoadingInitializerTask Fix the init() and loading task javadoc to say what happens, and replace the README's Thread.sleep(10000) advice with isReady(), stating that this can take up to one refresh period plus the load time and that until a model loads it returns its no-match value: LowValue models forward (fail open), HighValue models score the request low value. Fixes #17 --- README.md | 8 +- .../factory/BidRequestEvaluatorFactory.java | 41 ++ .../DefaultBidRequestEvaluatorFactory.java | 70 +++- .../DefaultTaskInitializerFactory.java | 4 +- .../task/InitializerTaskOnPeriodicTask.java | 7 + .../task/PeriodicTaskWithRandomizedStart.java | 21 + .../task/TaskInitializer.java | 29 ++ .../BloomFilterPeriodicLoadingTask.java | 3 +- ...lResultPeriodicLoadingInitializerTask.java | 7 + .../ModelResultPeriodicLoadingTask.java | 3 +- ...dRequestEvaluatorFactoryReadinessTest.java | 364 ++++++++++++++++++ .../PeriodicTaskWithRandomizedStartTest.java | 122 ++++++ .../task/TaskInitializerTest.java | 42 ++ 13 files changed, 705 insertions(+), 16 deletions(-) create mode 100644 java/src/test/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactoryReadinessTest.java create mode 100644 java/src/test/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStartTest.java diff --git a/README.md b/README.md index 3b2ab8e..b6d2e12 100644 --- a/README.md +++ b/README.md @@ -657,16 +657,18 @@ BidRequestEvaluatorFactory bidRequestEvaluatorFactory = bucket, executor); -// initialize background tasks to periodically fetch new signal data to evaluate. +// load the configurations and schedule background tasks to periodically fetch new signal data to evaluate. bidRequestEvaluatorFactory.getTaskInitializer().init(); // create an evaluator instance to receive recommendations. BidRequestEvaluator bidRequestEvaluator = bidRequestEvaluatorFactory.getEvaluator(); -// sleep to ensure initialization is complete before evaluating requests. -Thread.sleep(10000); +// report readiness (e.g. from a health check) instead of sleeping for a fixed time. +boolean ready = bidRequestEvaluatorFactory.isReady(); ``` +`init()` loads the model and experiment configurations and schedules the periodic refresh; it throws if the configurations cannot be loaded within `task.overall.execution-timeout`. It does not wait for the model results: each model loading task runs first after a random delay of up to one refresh period (5 minutes by default), so that hosts that start together do not all read S3 at once. Until a model's results are loaded, that model returns its no-match value: 1.0 (forward) for a `LowValue` model and 0.0 (low value) for a `HighValue` model, combined with the other models by the aggregation as usual. With only `LowValue` models, evaluations therefore fail open (traffic is forwarded) until the models load. `isReady()` turns true once `init()` has completed, the feature registries are populated and every model loading task has finished its first run, which can take up to one refresh period (the longest configured model-result period) plus the model load time. A first run that fails, for example because a model file is missing or S3 returns an error, still counts as finished: that model keeps returning its no-match value until a later refresh loads it. Once true, `isReady()` stays true. Poll it from a readiness probe, or before taking traffic, if you prefer not to evaluate before the models are loaded. + In addition, update the logging configurations under `src/main/resources/log4j2.xml` as needed. For the Golang library, the initialization method is slightly different. diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/BidRequestEvaluatorFactory.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/BidRequestEvaluatorFactory.java index fc225bf..143c32f 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/BidRequestEvaluatorFactory.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/BidRequestEvaluatorFactory.java @@ -32,4 +32,45 @@ public static BidRequestEvaluatorFactory create(String supplierName, AwsCredenti public abstract BidRequestEvaluator getEvaluator(); + /** + * Tells whether the evaluator has had its first chance to load the data it needs to make recommendations. + *

+ * True once all of the following hold: + *

+ * + *

+ * The first periodic run starts after a random delay of up to one refresh period (300 s by default; the + * longest of the configured model-result periods), so this can take up to one refresh period plus the model + * load time after {@code init()} returns, longer if the executor has no free thread. Once true it stays true; + * later refresh failures do not reset it. It stays false if {@code init()} throws, and while a feature + * extractor or transformer provider has not been registered, for example because it failed to load. + *

+ *

+ * A first run counts as finished whatever its outcome: a model file that is missing, an S3 error or a loader + * exception does not block readiness. + *

+ *

+ * Until a model's results are loaded, evaluations still run: that model returns its no-match value, which + * forwards for a {@code LowValue} model (1.0) and marks the request low value for a {@code HighValue} model + * (0.0), and the aggregation combines it with the other models as usual. With only {@code LowValue} models, + * evaluations therefore fail open (forward) until the models load. + *

+ *

+ * This default implementation throws; {@link DefaultBidRequestEvaluatorFactory} overrides it. + *

+ * + * @return true when the conditions above hold. + * @throws UnsupportedOperationException if the factory does not track readiness. + */ + public boolean isReady() { + throw new UnsupportedOperationException("isReady is not implemented by " + getClass().getName()); + } + } diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactory.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactory.java index e0c38dc..aad5e45 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactory.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactory.java @@ -14,7 +14,9 @@ import com.amazon.demanddriventrafficevaluator.evaluation.experiment.ExperimentManager; import com.amazon.demanddriventrafficevaluator.modelfeature.Extraction; import com.amazon.demanddriventrafficevaluator.modelfeature.Transformation; +import com.amazon.demanddriventrafficevaluator.modelfeature.extractor.Extractor; import com.amazon.demanddriventrafficevaluator.modelfeature.extractor.ExtractorRegistry; +import com.amazon.demanddriventrafficevaluator.modelfeature.transformer.Transformer; import com.amazon.demanddriventrafficevaluator.modelfeature.transformer.TransformerRegistry; import com.amazon.demanddriventrafficevaluator.repository.dao.BloomFilterDao; import com.amazon.demanddriventrafficevaluator.repository.dao.Dao; @@ -33,9 +35,11 @@ import com.amazon.demanddriventrafficevaluator.repository.provider.model.RuleBasedModelResultProvider; import com.amazon.demanddriventrafficevaluator.task.InitializerTask; import com.amazon.demanddriventrafficevaluator.task.InitializerTaskOnPeriodicTask; +import com.amazon.demanddriventrafficevaluator.task.PeriodicTaskWithRandomizedStart; import com.amazon.demanddriventrafficevaluator.task.TaskConfiguration; import com.amazon.demanddriventrafficevaluator.task.TaskInitializer; import com.amazon.demanddriventrafficevaluator.task.dataloading.BloomFilterPeriodicLoadingTask; +import com.amazon.demanddriventrafficevaluator.task.dataloading.ModelResultPeriodicLoadingInitializerTask; import com.amazon.demanddriventrafficevaluator.util.PropertiesUtil; import lombok.extern.log4j.Log4j2; import org.apache.commons.configuration2.Configuration; @@ -44,7 +48,11 @@ import java.io.InputStream; import java.util.List; import java.util.Map; +import java.util.ServiceConfigurationError; +import java.util.ServiceLoader; +import java.util.Set; import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.stream.Collectors; /** * A factory for creating BidRequestEvaluator instances that support both rule-based @@ -75,6 +83,12 @@ public class DefaultBidRequestEvaluatorFactory extends BidRequestEvaluatorFactor private final String bucket; private final ScheduledThreadPoolExecutor executor; private final Dao fileDao; + // The model loading tasks of the last TaskInitializer whose init() completed; null until one has. + private volatile List modelLoadingTasks; + private volatile boolean ready; + // The names the registry setup tasks register; listed once, on the first isReady() call that needs them. + private volatile Set extractorNames; + private volatile Set transformerNames; public DefaultBidRequestEvaluatorFactory( String supplierName, @@ -110,9 +124,53 @@ public TaskInitializer getTaskInitializer() { } List stageOneTasks = getStageOneTasks(); - List stageTwoTasks = getStageTwoTasks(); + ModelResultPeriodicLoadingInitializerTask ruleBasedTask = + defaultTaskInitializerFactory.getInitializerTaskForPeriodicLoadingRuleBasedModelResult(); + InitializerTaskOnPeriodicTask bloomFilterTask = getInitializerTaskForPeriodicLoadingBloomFilterModelResult(); + List stageTwoTasks = List.of(ruleBasedTask, bloomFilterTask); + List tasks = List.of(ruleBasedTask.getTask(), bloomFilterTask.getTask()); - return new TaskInitializer(stageOneTasks, stageTwoTasks, overallTimeoutMs); + return new TaskInitializer(stageOneTasks, stageTwoTasks, overallTimeoutMs, () -> modelLoadingTasks = tasks); + } + + /** + * {@inheritDoc} + *

+ * This factory waits for the rule-based model-result and the bloom-filter loading tasks of the + * {@link TaskInitializer} whose {@code init()} completed last. + *

+ */ + @Override + public boolean isReady() { + if (ready) { + return true; + } + List tasks = modelLoadingTasks; + if (tasks == null || !tasks.stream().allMatch(PeriodicTaskWithRandomizedStart::hasFinishedFirstRun)) { + return false; + } + try { + if (extractorNames == null || transformerNames == null) { + extractorNames = getProviderNames(Extractor.class); + transformerNames = getProviderNames(Transformer.class); + } + } catch (ServiceConfigurationError e) { + log.error("isReady: cannot list the feature extractor and transformer providers", e); + return false; + } + if (ExtractorRegistryFactory.getInstance().getSingleton().getRecords().keySet().containsAll(extractorNames) + && TransformerRegistryFactory.getInstance().getSingleton().getRecords().keySet() + .containsAll(transformerNames)) { + ready = true; + } + return ready; + } + + // The names the registry setup tasks register: the simple class names of the ServiceLoader providers. + private static Set getProviderNames(Class type) { + return ServiceLoader.load(type).stream() + .map(provider -> provider.type().getSimpleName()) + .collect(Collectors.toUnmodifiableSet()); } @Override @@ -167,13 +225,7 @@ private List getStageOneTasks() { return defaultTaskInitializerFactory.getStageOneTasks(); } - private List getStageTwoTasks() { - InitializerTask ruleBasedTask = defaultTaskInitializerFactory.getInitializerTaskForPeriodicLoadingRuleBasedModelResult(); - InitializerTask bloomFilterTask = getInitializerTaskForPeriodicLoadingBloomFilterModelResult(); - return List.of(ruleBasedTask, bloomFilterTask); - } - - private InitializerTask getInitializerTaskForPeriodicLoadingBloomFilterModelResult() { + private InitializerTaskOnPeriodicTask getInitializerTaskForPeriodicLoadingBloomFilterModelResult() { LocalCacheRegistry localCacheRegistry = DefaultLocalCacheRegistryFactory.getInstance().getDefaultLocalCacheRegistrySingleton(); LocalCacheDao modelConfigurationCacheDao = new LocalCacheDao<>(localCacheRegistry); ConfigurationProvider modelConfigurationProvider = new ModelConfigurationProvider(modelConfigurationCacheDao); diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultTaskInitializerFactory.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultTaskInitializerFactory.java index 3fc9e2f..2ba843f 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultTaskInitializerFactory.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/factory/DefaultTaskInitializerFactory.java @@ -165,7 +165,7 @@ private InitializerTaskOnPeriodicTask getInitializerTaskForPeriodicLoadingExperi ); } - InitializerTask getInitializerTaskForPeriodicLoadingRuleBasedModelResult() { + ModelResultPeriodicLoadingInitializerTask getInitializerTaskForPeriodicLoadingRuleBasedModelResult() { LocalCacheDao modelConfigurationCacheDao = new LocalCacheDao<>(localCacheRegistry); ConfigurationProvider modelConfigurationProvider = new ModelConfigurationProvider(modelConfigurationCacheDao); Dao modelResultsCacheDao = new LocalCacheDao<>(localCacheRegistry); @@ -216,7 +216,7 @@ private InitializerTaskOnPeriodicTask getInitializerTaskForPeriodicLoadingConfig ); } - private InitializerTask getInitializerTaskForPeriodicLoadingModelResult( + private ModelResultPeriodicLoadingInitializerTask getInitializerTaskForPeriodicLoadingModelResult( String taskName, long periodMs, int maximumAttempts, diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/InitializerTaskOnPeriodicTask.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/InitializerTaskOnPeriodicTask.java index 94630e3..cb6fa0f 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/InitializerTaskOnPeriodicTask.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/InitializerTaskOnPeriodicTask.java @@ -18,6 +18,13 @@ public InitializerTaskOnPeriodicTask( this.task = task; } + /** + * @return the periodic task that {@link #run()} initializes. + */ + public PeriodicTaskWithRandomizedStart getTask() { + return task; + } + /** * When an object implementing interface {@code Runnable} is used * to create a thread, starting the thread causes the object's diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStart.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStart.java index 396f2c2..72dbb1a 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStart.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStart.java @@ -3,6 +3,7 @@ package com.amazon.demanddriventrafficevaluator.task; +import lombok.AccessLevel; import lombok.Getter; import java.util.Random; @@ -30,6 +31,10 @@ public abstract class PeriodicTaskWithRandomizedStart { private final Random random = new Random(System.currentTimeMillis()); + // Written on the scheduler thread, read by callers of hasFinishedFirstRun(). + @Getter(AccessLevel.NONE) + private volatile boolean firstRunFinished; + protected PeriodicTaskWithRandomizedStart(String sspIdentifier, String taskName, long periodMs, ScheduledThreadPoolExecutor executor) { this.sspIdentifier = sspIdentifier; this.taskName = taskName; @@ -49,6 +54,20 @@ protected PeriodicTaskWithRandomizedStart(String sspIdentifier, String taskName, */ abstract public void initialize(); + /** + * Tells whether the first scheduled execution of this task has finished. + *

+ * True once the first execution scheduled by {@link #schedulePeriodically(long)} has returned or thrown, + * whatever its outcome. An execution run directly through {@link #executeTask()}, for example by + * {@link #initialize()}, does not count. Once true it stays true. + *

+ * + * @return true when the first scheduled execution has finished. + */ + public boolean hasFinishedFirstRun() { + return firstRunFinished; + } + /** * Schedules the task to run periodically with a specified initial delay. * @@ -61,6 +80,8 @@ protected void schedulePeriodically(long initialDelay) { } catch (Exception e) { log.error("{}: task failed", taskName, e); // Optionally rethrow or handle the exception + } finally { + firstRunFinished = true; } }, initialDelay, periodMs, TimeUnit.MILLISECONDS); } diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializer.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializer.java index 168f1e3..f8e8aba 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializer.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializer.java @@ -4,6 +4,7 @@ package com.amazon.demanddriventrafficevaluator.task; import java.util.List; +import java.util.Objects; import java.util.concurrent.CompletionService; import java.util.concurrent.ExecutionException; import java.util.concurrent.ExecutorCompletionService; @@ -25,13 +26,24 @@ public class TaskInitializer { private final List stageOneTasks; private final List stageTwoTasks; private final long overallTimeoutMs; + private final Runnable onInitialized; private ExecutorService executorPool; private long taskElapsedTime = 0L; public TaskInitializer(List stageOneTasks, List stageTwoTasks, long overallTimeoutMs) { + this(stageOneTasks, stageTwoTasks, overallTimeoutMs, () -> { }); + } + + /** + * @param onInitialized Run by {@link #init()} once the tasks of both stages have completed without an exception; + * not run when there are no stage one tasks. + */ + public TaskInitializer(List stageOneTasks, List stageTwoTasks, long overallTimeoutMs, + Runnable onInitialized) { this.stageOneTasks = stageOneTasks; this.stageTwoTasks = stageTwoTasks; this.overallTimeoutMs = overallTimeoutMs; + this.onInitialized = Objects.requireNonNull(onInitialized, "onInitialized"); } /** @@ -57,6 +69,18 @@ static long calculateDelay(int timesAlreadyAttempted, final long currentDelay, l * This method executes stage one tasks, followed by stage two tasks if present. * It manages the execution pool and handles overall timeout. *

+ *

+ * With the default factories, stage one loads the model and experiment configurations and + * schedules the feature extractor and transformer registration, which runs within about one + * second; stage two only schedules the periodic model-result and bloom-filter loading. The first + * model load runs on that schedule after a random delay of up to one refresh period (300 s by + * default), so this method does not wait for the model results. Until they are loaded, each + * model returns its no-match value: 1.0 (forward) for a {@code LowValue} model and 0.0 (low + * value) for a {@code HighValue} model. Callers that need to know when the first load has + * finished use {@code BidRequestEvaluatorFactory.isReady()}. + *

+ * + * @throws RuntimeException if a task fails after its retries or the overall timeout elapses. */ public void init() { if (stageOneTasks != null && !stageOneTasks.isEmpty()) { @@ -65,6 +89,7 @@ public void init() { if (stageTwoTasks != null && !stageTwoTasks.isEmpty()) { executeTasks(stageTwoTasks, "StageTwo", System.currentTimeMillis()); } + onInitialized.run(); } finally { executorPool.shutdown(); try { @@ -117,6 +142,10 @@ private void tasksResultIterator(CompletionService completionService, List } taskElapsedTime += elapsed; } + if (executedTaskCount < tasks.size()) { + throw new RuntimeException("Timed out after " + executedTaskCount + " of " + tasks.size() + " tasks and " + + (System.currentTimeMillis() - startTime) + " ms"); + } } /** diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/BloomFilterPeriodicLoadingTask.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/BloomFilterPeriodicLoadingTask.java index 64e3b5b..5dffe49 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/BloomFilterPeriodicLoadingTask.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/BloomFilterPeriodicLoadingTask.java @@ -95,7 +95,8 @@ public void executeTask() { *

* This method sets up the periodic execution schedule for the task. * Unlike some other tasks, this does not perform an immediate execution - * upon initialization. + * upon initialization: the first execution runs after a random delay of + * up to one period, and {@link #hasFinishedFirstRun()} tells when it has finished. *

*/ @Override diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingInitializerTask.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingInitializerTask.java index 4175c47..bfcbbb7 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingInitializerTask.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingInitializerTask.java @@ -14,6 +14,13 @@ public ModelResultPeriodicLoadingInitializerTask(String taskName, int maximumAtt this.task = task; } + /** + * @return the periodic task that {@link #run()} initializes. + */ + public ModelResultPeriodicLoadingTask getTask() { + return task; + } + /** * When an object implementing interface {@code Runnable} is used * to create a thread, starting the thread causes the object's diff --git a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingTask.java b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingTask.java index 8fe80d9..c6fa1cd 100644 --- a/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingTask.java +++ b/java/src/main/java/com/amazon/demanddriventrafficevaluator/task/dataloading/ModelResultPeriodicLoadingTask.java @@ -86,7 +86,8 @@ public void executeTask() { *

* This method sets up the periodic execution schedule for the task. * Unlike some other tasks, this does not perform an immediate execution - * upon initialization. + * upon initialization: the first execution runs after a random delay of + * up to one period, and {@link #hasFinishedFirstRun()} tells when it has finished. *

*/ @Override diff --git a/java/src/test/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactoryReadinessTest.java b/java/src/test/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactoryReadinessTest.java new file mode 100644 index 0000000..a5a8132 --- /dev/null +++ b/java/src/test/java/com/amazon/demanddriventrafficevaluator/factory/DefaultBidRequestEvaluatorFactoryReadinessTest.java @@ -0,0 +1,364 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package com.amazon.demanddriventrafficevaluator.factory; + +import com.amazon.demanddriventrafficevaluator.modelfeature.extractor.Extractor; +import com.amazon.demanddriventrafficevaluator.modelfeature.transformer.Transformer; +import com.amazon.demanddriventrafficevaluator.repository.dao.BloomFilterDao; +import com.amazon.demanddriventrafficevaluator.repository.dao.LocalCacheDao; +import com.amazon.demanddriventrafficevaluator.repository.localcache.LocalCacheRegistry; +import com.google.common.hash.BloomFilter; +import com.google.common.hash.Funnels; +import org.junit.jupiter.api.AfterEach; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentMatchers; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; +import software.amazon.awssdk.auth.credentials.AwsCredentialsProvider; +import software.amazon.awssdk.core.ResponseInputStream; +import software.amazon.awssdk.services.s3.S3Client; +import software.amazon.awssdk.services.s3.model.GetObjectRequest; +import software.amazon.awssdk.services.s3.model.GetObjectResponse; +import software.amazon.awssdk.services.s3.model.HeadObjectRequest; +import software.amazon.awssdk.services.s3.model.HeadObjectResponse; +import software.amazon.awssdk.services.s3.model.NoSuchKeyException; +import software.amazon.awssdk.services.s3.model.S3Exception; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.InputStream; +import java.lang.reflect.Field; +import java.nio.charset.StandardCharsets; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.List; +import java.util.Map; +import java.util.UUID; +import java.util.concurrent.ScheduledFuture; +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; +import java.util.function.Supplier; + +import static com.amazon.demanddriventrafficevaluator.repository.dao.LocalCacheDao.CACHE_KEY_EXPERIMENT_CONFIGURATION; +import static com.amazon.demanddriventrafficevaluator.repository.dao.LocalCacheDao.CACHE_KEY_MODEL_CONFIGURATION; +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.Mockito.doThrow; +import static org.mockito.Mockito.when; + +/** + * Verifies {@link DefaultBidRequestEvaluatorFactory#isReady()} while the first model load stays on the periodic + * schedule. + *

+ * S3 is faked with a mocked {@link S3Client} injected into {@link AWSServiceClientFactory}. The executor runs + * one-shot tasks (the registry setup) inline or defers them, and captures periodic tasks without running them, + * so each test decides when a periodic run happens. + *

+ */ +@ExtendWith(MockitoExtension.class) +class DefaultBidRequestEvaluatorFactoryReadinessTest { + + private static final String SSP_IDENTIFIER = "readiness-test-ssp"; + private static final String RULE_BASED_MODEL_IDENTIFIER = "adsp_low-value_v2"; + private static final String BLOOM_FILTER_MODEL_IDENTIFIER = "adsp_inv_v1"; + private static final String MODEL_RESULT_KEY = "site|video|5895-EB|USA|640x390|u|0"; + private static final long DEFAULT_PERIOD_MS = 300000L; + // stage one schedules the experiment and model configuration refreshes, stage two the two model loads + private static final int CONFIGURATION_TASK_COUNT = 2; + private static final int PERIODIC_TASK_COUNT = 4; + + private final LocalCacheRegistry localCacheRegistry = + DefaultLocalCacheRegistryFactory.getInstance().getDefaultLocalCacheRegistrySingleton(); + private final LocalCacheDao modelResultsCacheDao = new LocalCacheDao<>(localCacheRegistry); + private final ControllableExecutor executor = new ControllableExecutor(); + private final String eTag = UUID.randomUUID().toString(); + private final Map extractors = ExtractorRegistryFactory.getInstance().getSingleton().getRecords(); + private final Map transformers = TransformerRegistryFactory.getInstance().getSingleton().getRecords(); + private final Map extractorsBefore = new HashMap<>(extractors); + private final Map transformersBefore = new HashMap<>(transformers); + + @Mock + private AwsCredentialsProvider credentialsProvider; + + @Mock + private S3Client s3Client; + + private S3Client s3ClientBefore; + + @BeforeEach + void setUp() throws Exception { + s3ClientBefore = getS3Client(); + setS3Client(s3Client); + // the registries are JVM-wide singletons; start empty so readiness depends on this test's init() only + // (safe while test classes run sequentially, as configured today) + extractors.clear(); + transformers.clear(); + modelResultsCacheDao.clear(RULE_BASED_MODEL_IDENTIFIER); + when(s3Client.headObject(ArgumentMatchers.any())) + .thenReturn(HeadObjectResponse.builder().contentLength(1024L).build()); + } + + @AfterEach + void tearDown() throws Exception { + setS3Client(s3ClientBefore); + executor.shutdownNow(); + // the configurations loaded by init() live in JVM-wide caches; leave them as a fresh JVM would + modelResultsCacheDao.clear(CACHE_KEY_MODEL_CONFIGURATION); + modelResultsCacheDao.clear(CACHE_KEY_EXPERIMENT_CONFIGURATION); + modelResultsCacheDao.clear(RULE_BASED_MODEL_IDENTIFIER); + extractors.clear(); + extractors.putAll(extractorsBefore); + transformers.clear(); + transformers.putAll(transformersBefore); + } + + @Test + void testIsReady_falseAfterInitUntilEveryModelTaskHasRunOnce() throws Exception { + // Arrange + BloomFilter bloomFilter = BloomFilter.create(Funnels.stringFunnel(StandardCharsets.UTF_8), 100, 0.01); + bloomFilter.put("539014228|USA"); + stubS3("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/test/ModelConfigurationWithBloomFilter.json", () -> response(resource("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/test/ModelResult.csv")), + serialize(bloomFilter)); + DefaultBidRequestEvaluatorFactory factory = newFactory(); + assertFalse(factory.isReady(), "not ready before init()"); + + // Act + factory.getTaskInitializer().init(); + + // Assert: init() only scheduled the model loads, with a random first delay of up to one period + assertEquals(PERIODIC_TASK_COUNT, executor.periodicTasks.size()); + for (ScheduledRun run : executor.periodicTasks) { + assertTrue(run.initialDelayMs >= 0 && run.initialDelayMs < DEFAULT_PERIOD_MS); + } + assertFalse(modelResultsCacheDao.get(RULE_BASED_MODEL_IDENTIFIER, MODEL_RESULT_KEY).isPresent(), + "init() must not load the rule-based model results"); + assertFalse(getBloomFilterDao(factory).get(BLOOM_FILTER_MODEL_IDENTIFIER, null).isPresent(), + "init() must not load the bloom filter"); + assertFalse(factory.isReady(), "not ready before the first model runs"); + + executor.runPeriodicTasks(0, CONFIGURATION_TASK_COUNT); + assertFalse(factory.isReady(), "configuration refreshes do not make it ready"); + executor.runPeriodicTasks(CONFIGURATION_TASK_COUNT, CONFIGURATION_TASK_COUNT + 1); + assertFalse(factory.isReady(), "not ready while one model task has not run yet"); + executor.runPeriodicTasks(CONFIGURATION_TASK_COUNT + 1, PERIODIC_TASK_COUNT); + assertTrue(factory.isReady(), "ready once every model task has run once"); + assertTrue(modelResultsCacheDao.get(RULE_BASED_MODEL_IDENTIFIER, MODEL_RESULT_KEY).isPresent()); + assertTrue(getBloomFilterDao(factory).get(BLOOM_FILTER_MODEL_IDENTIFIER, null).isPresent()); + } + + @Test + void testIsReady_trueAfterFirstRunsWhenModelFilesAreMissing() { + // Arrange: configurations exist, but neither model file exists + stubS3("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/test/ModelConfigurationWithBloomFilter.json", null, null); + DefaultBidRequestEvaluatorFactory factory = newFactory(); + factory.getTaskInitializer().init(); + assertFalse(factory.isReady()); + + // Act + executor.runPeriodicTasks(0, PERIODIC_TASK_COUNT); + + // Assert + assertFalse(modelResultsCacheDao.get(RULE_BASED_MODEL_IDENTIFIER, MODEL_RESULT_KEY).isPresent()); + assertTrue(factory.isReady(), "a missing model file still finishes the first run"); + } + + @Test + void testIsReady_trueAfterFirstRunsWhenTheLoadersFail() { + // Arrange: the rule-based model file cannot be parsed, so its loader throws, and the bloom filter is missing + stubS3("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/test/ModelConfigurationWithBloomFilter.json", + () -> new ResponseInputStream<>(GetObjectResponse.builder().eTag(UUID.randomUUID().toString()).build(), + new InputStream() { + @Override + public int read() { + throw new IllegalArgumentException("malformed model data"); + } + }), + null); + DefaultBidRequestEvaluatorFactory factory = newFactory(); + factory.getTaskInitializer().init(); + assertFalse(factory.isReady()); + + // Act: the scheduled wrapper logs the exceptions + assertDoesNotThrow(() -> executor.runPeriodicTasks(0, PERIODIC_TASK_COUNT)); + + // Assert + assertFalse(modelResultsCacheDao.get(RULE_BASED_MODEL_IDENTIFIER, MODEL_RESULT_KEY).isPresent()); + assertTrue(factory.isReady(), "a failed first load still finishes the first run"); + } + + @Test + void testIsReady_bloomFilterTaskDoesNotBlockWhenNoBloomFilterModelIsConfigured() { + // Arrange: a configuration with a rule-based model only + stubS3("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/test/ModelConfiguration.json", () -> response(resource("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/test/ModelResult.csv")), null); + DefaultBidRequestEvaluatorFactory factory = newFactory(); + factory.getTaskInitializer().init(); + + // Act + executor.runPeriodicTasks(0, PERIODIC_TASK_COUNT); + + // Assert + assertTrue(factory.isReady()); + } + + @Test + void testIsReady_staysTrueAfterLaterRuns() { + // Arrange + stubS3("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/test/ModelConfigurationWithBloomFilter.json", null, null); + DefaultBidRequestEvaluatorFactory factory = newFactory(); + factory.getTaskInitializer().init(); + executor.runPeriodicTasks(0, PERIODIC_TASK_COUNT); + assertTrue(factory.isReady()); + + // Act: a later refresh fails and the registries are emptied + doThrow(S3Exception.builder().message("throttled").statusCode(503).build()) + .when(s3Client).getObject(ArgumentMatchers.any()); + executor.runPeriodicTasks(0, PERIODIC_TASK_COUNT); + extractors.clear(); + + // Assert + assertTrue(factory.isReady(), "once true it stays true"); + } + + @Test + void testIsReady_waitsForRegistrySetup() { + // Arrange: one-shot registry setup is deferred, as it is with a real executor + stubS3("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/test/ModelConfigurationWithBloomFilter.json", null, null); + executor.deferOneShotTasks = true; + DefaultBidRequestEvaluatorFactory factory = newFactory(); + factory.getTaskInitializer().init(); + executor.runPeriodicTasks(0, PERIODIC_TASK_COUNT); + + // Assert + assertFalse(factory.isReady(), "not ready while the registries are still empty"); + executor.runDeferredOneShotTasks(); + assertTrue(factory.isReady()); + } + + @Test + void testIsReady_falseWhenInitFails() { + // Arrange: no configuration file exists, so the configuration load fails on every attempt + when(s3Client.getObject(ArgumentMatchers.any())) + .thenThrow(NoSuchKeyException.builder().message("no such key").build()); + DefaultBidRequestEvaluatorFactory factory = newFactory(); + + // Act + assertThrows(RuntimeException.class, () -> factory.getTaskInitializer().init()); + executor.runPeriodicTasks(0, executor.periodicTasks.size()); + + // Assert + assertFalse(factory.isReady()); + } + + // --- helpers --- + + private DefaultBidRequestEvaluatorFactory newFactory() { + return new DefaultBidRequestEvaluatorFactory(SSP_IDENTIFIER, credentialsProvider, "us-east-1", "test-bucket", + executor); + } + + private void stubS3(String modelConfiguration, Supplier> modelResult, + byte[] bloomFilterBytes) { + when(s3Client.getObject(ArgumentMatchers.any())).thenAnswer(invocation -> { + String key = ((GetObjectRequest) invocation.getArgument(0)).key(); + if ((SSP_IDENTIFIER + "/configuration/model/config.json").equals(key)) { + return response(resource(modelConfiguration)); + } else if ((SSP_IDENTIFIER + "/configuration/experiment/config.json").equals(key)) { + return response(resource("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/test/ExperimentConfigurationWithAggregation.json")); + } else if (key.matches(SSP_IDENTIFIER + "/\\d{4}-\\d{2}-\\d{2}/\\d{2}/" + RULE_BASED_MODEL_IDENTIFIER + "\\.csv") + && modelResult != null) { + return modelResult.get(); + } else if ((SSP_IDENTIFIER + "/models/" + BLOOM_FILTER_MODEL_IDENTIFIER + ".bloom").equals(key) + && bloomFilterBytes != null) { + return response(new ByteArrayInputStream(bloomFilterBytes)); + } + throw NoSuchKeyException.builder().message("no such key: " + key).build(); + }); + } + + private ResponseInputStream response(InputStream content) { + return new ResponseInputStream<>(GetObjectResponse.builder().eTag(eTag).build(), content); + } + + private InputStream resource(String path) { + return getClass().getResourceAsStream(path); + } + + private static byte[] serialize(BloomFilter bloomFilter) throws Exception { + ByteArrayOutputStream out = new ByteArrayOutputStream(); + bloomFilter.writeTo(out); + return out.toByteArray(); + } + + private static S3Client getS3Client() throws Exception { + Field field = AWSServiceClientFactory.class.getDeclaredField("s3Client"); + field.setAccessible(true); + return (S3Client) field.get(AWSServiceClientFactory.getInstance()); + } + + private static void setS3Client(S3Client client) throws Exception { + Field field = AWSServiceClientFactory.class.getDeclaredField("s3Client"); + field.setAccessible(true); + field.set(AWSServiceClientFactory.getInstance(), client); + } + + private static BloomFilterDao getBloomFilterDao(DefaultBidRequestEvaluatorFactory factory) throws Exception { + Field field = DefaultBidRequestEvaluatorFactory.class.getDeclaredField("bloomFilterDao"); + field.setAccessible(true); + return (BloomFilterDao) field.get(factory); + } + + private static final class ScheduledRun { + private final Runnable command; + private final long initialDelayMs; + + ScheduledRun(Runnable command, long initialDelayMs) { + this.command = command; + this.initialDelayMs = initialDelayMs; + } + } + + /** + * Runs one-shot tasks inline or defers them, and captures periodic tasks in scheduling order without running + * them. Stage one completes before stage two starts, so the configuration refreshes come first. + */ + private static final class ControllableExecutor extends ScheduledThreadPoolExecutor { + private final List deferredOneShotTasks = new ArrayList<>(); + private final List periodicTasks = new ArrayList<>(); + private boolean deferOneShotTasks; + + ControllableExecutor() { + super(1); + } + + @Override + public synchronized ScheduledFuture schedule(Runnable command, long delay, TimeUnit unit) { + if (deferOneShotTasks) { + deferredOneShotTasks.add(command); + } else { + command.run(); + } + return null; + } + + @Override + public synchronized ScheduledFuture scheduleAtFixedRate(Runnable command, long initialDelay, long period, + TimeUnit unit) { + periodicTasks.add(new ScheduledRun(command, unit.toMillis(initialDelay))); + return null; + } + + synchronized void runPeriodicTasks(int fromIndex, int toIndex) { + periodicTasks.subList(fromIndex, toIndex).forEach(run -> run.command.run()); + } + + synchronized void runDeferredOneShotTasks() { + deferredOneShotTasks.forEach(Runnable::run); + } + } +} diff --git a/java/src/test/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStartTest.java b/java/src/test/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStartTest.java new file mode 100644 index 0000000..398fd4e --- /dev/null +++ b/java/src/test/java/com/amazon/demanddriventrafficevaluator/task/PeriodicTaskWithRandomizedStartTest.java @@ -0,0 +1,122 @@ +// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. +// SPDX-License-Identifier: Apache-2.0 + +package com.amazon.demanddriventrafficevaluator.task; + +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.extension.ExtendWith; +import org.mockito.ArgumentCaptor; +import org.mockito.Mock; +import org.mockito.junit.jupiter.MockitoExtension; + +import java.util.concurrent.ScheduledThreadPoolExecutor; +import java.util.concurrent.TimeUnit; + +import static org.junit.jupiter.api.Assertions.assertDoesNotThrow; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; +import static org.mockito.ArgumentMatchers.any; +import static org.mockito.ArgumentMatchers.anyLong; +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.verify; + +@ExtendWith(MockitoExtension.class) +class PeriodicTaskWithRandomizedStartTest { + + private static final long PERIOD_MS = 300000L; + + @Mock + private ScheduledThreadPoolExecutor mockExecutor; + + @Test + void testHasFinishedFirstRun_falseUntilTheFirstScheduledRunFinishes() { + // Arrange + TestTask task = new TestTask(null); + + // Act + task.schedulePeriodically(); + Runnable scheduledRun = captureScheduledRun(); + + // Assert + assertFalse(task.hasFinishedFirstRun(), "scheduling alone does not run the task"); + scheduledRun.run(); + assertTrue(task.hasFinishedFirstRun()); + assertTrue(task.executed); + } + + @Test + void testHasFinishedFirstRun_trueWhenTheFirstRunThrowsAnException() { + // Arrange + TestTask task = new TestTask(new IllegalStateException("load failed")); + task.schedulePeriodically(); + Runnable scheduledRun = captureScheduledRun(); + + // Act: the exception is logged by the scheduled wrapper, as before + assertDoesNotThrow(scheduledRun::run); + + // Assert + assertTrue(task.hasFinishedFirstRun()); + } + + @Test + void testHasFinishedFirstRun_trueWhenTheFirstRunThrowsAnError() { + // Arrange + TestTask task = new TestTask(new AssertionError("fatal")); + task.schedulePeriodically(); + Runnable scheduledRun = captureScheduledRun(); + + // Act: an Error is not caught and ends the schedule, but the run has still finished + assertThrows(AssertionError.class, scheduledRun::run); + + // Assert + assertTrue(task.hasFinishedFirstRun()); + } + + @Test + void testSchedulePeriodically_firstRunDelayIsWithinOnePeriod() { + // Arrange + TestTask task = new TestTask(null); + + // Act + task.schedulePeriodically(); + + // Assert + ArgumentCaptor initialDelay = ArgumentCaptor.forClass(Long.class); + verify(mockExecutor).scheduleAtFixedRate(any(Runnable.class), + initialDelay.capture(), eq(PERIOD_MS), eq(TimeUnit.MILLISECONDS)); + assertTrue(initialDelay.getValue() >= 0 && initialDelay.getValue() < PERIOD_MS); + } + + private Runnable captureScheduledRun() { + ArgumentCaptor runnable = ArgumentCaptor.forClass(Runnable.class); + verify(mockExecutor).scheduleAtFixedRate(runnable.capture(), anyLong(), eq(PERIOD_MS), + eq(TimeUnit.MILLISECONDS)); + return runnable.getValue(); + } + + private final class TestTask extends PeriodicTaskWithRandomizedStart { + private final Throwable failure; + private boolean executed; + + TestTask(Throwable failure) { + super("ssp", "TestTask", PERIOD_MS, mockExecutor); + this.failure = failure; + } + + @Override + public void executeTask() { + executed = true; + if (failure instanceof RuntimeException) { + throw (RuntimeException) failure; + } else if (failure instanceof Error) { + throw (Error) failure; + } + } + + @Override + public void initialize() { + schedulePeriodically(); + } + } +} diff --git a/java/src/test/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializerTest.java b/java/src/test/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializerTest.java index 46c1b8a..416284a 100644 --- a/java/src/test/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializerTest.java +++ b/java/src/test/java/com/amazon/demanddriventrafficevaluator/task/TaskInitializerTest.java @@ -10,9 +10,12 @@ import java.util.Arrays; import java.util.Collections; +import java.util.concurrent.atomic.AtomicBoolean; import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertThrows; +import static org.junit.jupiter.api.Assertions.assertTrue; import static org.mockito.Mockito.doAnswer; import static org.mockito.Mockito.doThrow; import static org.mockito.Mockito.times; @@ -90,6 +93,45 @@ void testInit_TaskExecutionFailure() { assertThrows(RuntimeException.class, () -> taskInitializer.init()); } + @Test + void testInit_RunsOnInitializedAfterBothStages() { + // Arrange + AtomicBoolean initialized = new AtomicBoolean(); + taskInitializer = new TaskInitializer(Collections.singletonList(mockTask1), + Collections.singletonList(mockTask2), 10000, () -> initialized.set(true)); + when(mockTask1.getName()).thenReturn("Task1"); + when(mockTask1.getMaximumAttempts()).thenReturn(1); + when(mockTask2.getName()).thenReturn("Task2"); + when(mockTask2.getMaximumAttempts()).thenReturn(1); + + // Act + taskInitializer.init(); + + // Assert + assertTrue(initialized.get()); + } + + @Test + void testInit_DoesNotRunOnInitializedWhenATaskFails() { + // Arrange + AtomicBoolean initialized = new AtomicBoolean(); + taskInitializer = new TaskInitializer(Collections.singletonList(mockTask1), null, 10000, + () -> initialized.set(true)); + when(mockTask1.getName()).thenReturn("Task1"); + when(mockTask1.getMaximumAttempts()).thenReturn(1); + doThrow(new RuntimeException("Task failed")).when(mockTask1).run(); + + // Act & Assert + assertThrows(RuntimeException.class, () -> taskInitializer.init()); + assertFalse(initialized.get()); + } + + @Test + void testConstructor_RejectsNullOnInitialized() { + assertThrows(NullPointerException.class, + () -> new TaskInitializer(Collections.singletonList(mockTask1), null, 10000, null)); + } + @Test void testInit_Timeout() { // Arrange