fix review4: modelcache thrcache as bean

Author: EvgeniiMunin

extra/modules/greenbids-real-time-data/src/main/java/org/prebid/server/hooks/modules/greenbids/real/time/data/config/GreenbidsRealTimeDataConfiguration.java

@@ -12,8 +12,10 @@
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.config.GreenbidsRealTimeDataProperties;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.data.GreenbidsInferenceDataService;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.FilterService;
+import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.ModelCache;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.OnnxModelRunner;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.OnnxModelRunnerWithThresholds;
+import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.ThresholdCache;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.v1.GreenbidsRealTimeDataModule;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.v1.GreenbidsRealTimeDataProcessedAuctionRequestHook;
 import org.prebid.server.json.ObjectMapperProvider;
@@ -33,8 +35,10 @@
 public class GreenbidsRealTimeDataConfiguration {
 
     @Bean
-    public GreenbidsInferenceDataService greenbidsInferenceDataService(GreenbidsRealTimeDataProperties properties) {
-        final ObjectMapper mapper = ObjectMapperProvider.mapper();
+    public GreenbidsInferenceDataService greenbidsInferenceDataService(
+            GreenbidsRealTimeDataProperties properties,
+            ObjectMapper mapper) {
+
         final File database = new File(properties.getGeoLiteCountryPath());
         DatabaseReader dbReader;
         try {
@@ -49,8 +53,9 @@ public GreenbidsInferenceDataService greenbidsInferenceDataService(GreenbidsReal
     GreenbidsRealTimeDataModule greenbidsRealTimeDataModule(
             FilterService filterService,
             OnnxModelRunnerWithThresholds onnxModelRunnerWithThresholds,
-            GreenbidsInferenceDataService greenbidsInferenceDataService) {
-        final ObjectMapper mapper = ObjectMapperProvider.mapper();
+            GreenbidsInferenceDataService greenbidsInferenceDataService,
+            ObjectMapper mapper) {
+
         return new GreenbidsRealTimeDataModule(List.of(
                 new GreenbidsRealTimeDataProcessedAuctionRequestHook(
                         mapper,
@@ -65,27 +70,51 @@ public FilterService filterService() {
     }
 
     @Bean
-    public OnnxModelRunnerWithThresholds onnxModelRunnerWithThresholds(
-            GreenbidsRealTimeDataProperties properties, Vertx vertx) {
-
-        final Storage storage = StorageOptions.newBuilder()
+    public Storage storage(GreenbidsRealTimeDataProperties properties) {
+        return StorageOptions.newBuilder()
                 .setProjectId(properties.getGoogleCloudGreenbidsProject()).build().getService();
+    }
 
+    @Bean
+    public ModelCache modelCache(GreenbidsRealTimeDataProperties properties, Vertx vertx, Storage storage) {
         final Cache<String, OnnxModelRunner> modelCacheWithExpiration = Caffeine.newBuilder()
                 .expireAfterWrite(properties.getCacheExpirationMinutes(), TimeUnit.MINUTES)
                 .build();
 
+        return new ModelCache(
+                storage,
+                properties.getGcsBucketName(),
+                modelCacheWithExpiration,
+                properties.getOnnxModelCacheKeyPrefix(),
+                vertx
+        );
+    }
+
+    @Bean
+    public ThresholdCache thresholdCache(
+            GreenbidsRealTimeDataProperties properties,
+            Vertx vertx,
+            Storage storage,
+            ObjectMapper mapper) {
+
         final Cache<String, ThrottlingThresholds> thresholdsCacheWithExpiration = Caffeine.newBuilder()
                 .expireAfterWrite(properties.getCacheExpirationMinutes(), TimeUnit.MINUTES)
                 .build();
 
-        return new OnnxModelRunnerWithThresholds(
-                modelCacheWithExpiration,
-                thresholdsCacheWithExpiration,
+        return new ThresholdCache(
                 storage,
                 properties.getGcsBucketName(),
-                properties.getOnnxModelCacheKeyPrefix(),
+                mapper,
+                thresholdsCacheWithExpiration,
                 properties.getThresholdsCacheKeyPrefix(),
                 vertx);
     }
+
+    @Bean
+    public OnnxModelRunnerWithThresholds onnxModelRunnerWithThresholds(
+            ModelCache modelCache,
+            ThresholdCache thresholdCache) {
+
+        return new OnnxModelRunnerWithThresholds(modelCache, thresholdCache);
+    }
 }

extra/modules/greenbids-real-time-data/src/main/java/org/prebid/server/hooks/modules/greenbids/real/time/data/model/predictor/ModelCache.java

@@ -20,8 +20,6 @@ public class ModelCache {
 
     String gcsBucketName;
 
-    String modelPath;
-
     Cache<String, OnnxModelRunner> cache;
 
     Storage storage;
@@ -33,22 +31,20 @@ public class ModelCache {
     Vertx vertx;
 
     public ModelCache(
-            String modelPath,
             Storage storage,
             String gcsBucketName,
             Cache<String, OnnxModelRunner> cache,
             String onnxModelCacheKeyPrefix,
             Vertx vertx) {
         this.gcsBucketName = gcsBucketName;
-        this.modelPath = modelPath;
         this.cache = cache;
         this.storage = storage;
         this.onnxModelCacheKeyPrefix = onnxModelCacheKeyPrefix;
         this.isFetching = new AtomicBoolean(false);
         this.vertx = vertx;
     }
 
-    public Future<OnnxModelRunner> getModelRunner(String pbuid) {
+    public Future<OnnxModelRunner> getModelRunner(String onnxModelPath, String pbuid) {
         final String cacheKey = onnxModelCacheKeyPrefix + pbuid;
         final OnnxModelRunner cachedOnnxModelRunner = cache.getIfPresent(cacheKey);
 
@@ -58,7 +54,7 @@ public Future<OnnxModelRunner> getModelRunner(String pbuid) {
 
         if (isFetching.compareAndSet(false, true)) {
             try {
-                fetchAndCacheModelRunner(cacheKey);
+                fetchAndCacheModelRunner(onnxModelPath, cacheKey);
             } finally {
                 isFetching.set(false);
             }
@@ -67,17 +63,20 @@ public Future<OnnxModelRunner> getModelRunner(String pbuid) {
         return Future.failedFuture("ModelRunner fetching in progress");
     }
 
-    private void fetchAndCacheModelRunner(String cacheKey) {
-        vertx.executeBlocking(this::getBlob)
-                .map(this::loadModelRunner)
+    private void fetchAndCacheModelRunner(String onnxModelPath, String cacheKey) {
+        vertx.executeBlocking(promise -> {
+                    Blob blob = getBlob(onnxModelPath);
+                    promise.complete(blob);
+                })
+                .map(blob -> loadModelRunner((Blob) blob))
                 .onSuccess(onnxModelRunner -> cache.put(cacheKey, onnxModelRunner))
                 .onFailure(error -> logger.error("Failed to fetch ONNX model"));
     }
 
-    private Blob getBlob() {
+    private Blob getBlob(String onnxModelPath) {
         try {
             return Optional.ofNullable(storage.get(gcsBucketName))
-                    .map(bucket -> bucket.get(modelPath))
+                    .map(bucket -> bucket.get(onnxModelPath))
                     .orElseThrow(() -> new PreBidException("Bucket not found: " + gcsBucketName));
         } catch (StorageException e) {
             throw new PreBidException("Error accessing GCS artefact for model: ", e);

extra/modules/greenbids-real-time-data/src/main/java/org/prebid/server/hooks/modules/greenbids/real/time/data/model/predictor/OnnxModelRunnerWithThresholds.java

@@ -10,54 +10,24 @@
 
 public class OnnxModelRunnerWithThresholds {
 
-    private final Cache<String, OnnxModelRunner> modelCacheWithExpiration;
-    private final Cache<String, ThrottlingThresholds> thresholdsCacheWithExpiration;
-    private final String gcsBucketName;
-    private final String onnxModelCacheKeyPrefix;
-    private final String thresholdsCacheKeyPrefix;
-    private final Storage storage;
-    private final Vertx vertx;
+    private final ModelCache modelCache;
+    private final ThresholdCache thresholdCache;
 
     public OnnxModelRunnerWithThresholds(
-            Cache<String, OnnxModelRunner> modelCacheWithExpiration,
-            Cache<String, ThrottlingThresholds> thresholdsCacheWithExpiration,
-            Storage storage,
-            String gcsBucketName,
-            String onnxModelCacheKeyPrefix,
-            String thresholdsCacheKeyPrefix,
-            Vertx vertx) {
-        this.modelCacheWithExpiration = modelCacheWithExpiration;
-        this.thresholdsCacheWithExpiration = thresholdsCacheWithExpiration;
-        this.gcsBucketName = gcsBucketName;
-        this.onnxModelCacheKeyPrefix = onnxModelCacheKeyPrefix;
-        this.thresholdsCacheKeyPrefix = thresholdsCacheKeyPrefix;
-        this.storage = storage;
-        this.vertx = vertx;
+            ModelCache modelCache,
+            ThresholdCache thresholdCache) {
+        this.modelCache = modelCache;
+        this.thresholdCache = thresholdCache;
     }
 
     public Future<OnnxModelRunner> retrieveOnnxModelRunner(Partner partner) {
         final String onnxModelPath = "models_pbuid=" + partner.getPbuid() + ".onnx";
-        final ModelCache modelCache = new ModelCache(
-                onnxModelPath,
-                storage,
-                gcsBucketName,
-                modelCacheWithExpiration,
-                onnxModelCacheKeyPrefix,
-                vertx);
-        return modelCache.getModelRunner(partner.getPbuid());
+        return modelCache.getModelRunner(onnxModelPath, partner.getPbuid());
     }
 
-    public Future<Double> retrieveThreshold(Partner partner, ObjectMapper mapper) {
+    public Future<Double> retrieveThreshold(Partner partner) {
         final String thresholdJsonPath = "thresholds_pbuid=" + partner.getPbuid() + ".json";
-        final ThresholdCache thresholdCache = new ThresholdCache(
-                thresholdJsonPath,
-                storage,
-                gcsBucketName,
-                mapper,
-                thresholdsCacheWithExpiration,
-                thresholdsCacheKeyPrefix,
-                vertx);
-        return thresholdCache.getThrottlingThresholds(partner.getPbuid())
+        return thresholdCache.getThrottlingThresholds(thresholdJsonPath, partner.getPbuid())
                 .map(partner::getThreshold);
     }
 }

extra/modules/greenbids-real-time-data/src/main/java/org/prebid/server/hooks/modules/greenbids/real/time/data/model/predictor/ThresholdCache.java

@@ -23,8 +23,6 @@ public class ThresholdCache {
 
     String gcsBucketName;
 
-    String thresholdPath;
-
     Cache<String, ThrottlingThresholds> cache;
 
     Storage storage;
@@ -38,15 +36,13 @@ public class ThresholdCache {
     Vertx vertx;
 
     public ThresholdCache(
-            String thresholdPath,
             Storage storage,
             String gcsBucketName,
             ObjectMapper mapper,
             Cache<String, ThrottlingThresholds> cache,
             String thresholdsCacheKeyPrefix,
             Vertx vertx) {
         this.gcsBucketName = gcsBucketName;
-        this.thresholdPath = thresholdPath;
         this.cache = cache;
         this.storage = storage;
         this.mapper = mapper;
@@ -55,7 +51,7 @@ public ThresholdCache(
         this.vertx = vertx;
     }
 
-    public Future<ThrottlingThresholds> getThrottlingThresholds(String pbuid) {
+    public Future<ThrottlingThresholds> getThrottlingThresholds(String thresholdJsonPath, String pbuid) {
         final String cacheKey = thresholdsCacheKeyPrefix + pbuid;
         final ThrottlingThresholds cachedThrottlingThresholds = cache.getIfPresent(cacheKey);
 
@@ -65,7 +61,7 @@ public Future<ThrottlingThresholds> getThrottlingThresholds(String pbuid) {
 
         if (isFetching.compareAndSet(false, true)) {
             try {
-                fetchAndCacheThrottlingThresholds(cacheKey);
+                fetchAndCacheThrottlingThresholds(thresholdJsonPath, cacheKey);
             } finally {
                 isFetching.set(false);
             }
@@ -74,17 +70,20 @@ public Future<ThrottlingThresholds> getThrottlingThresholds(String pbuid) {
         return Future.failedFuture("ThrottlingThresholds fetching in progress");
     }
 
-    private void fetchAndCacheThrottlingThresholds(String cacheKey) {
-        vertx.executeBlocking(this::getBlob)
-                .map(this::loadThrottlingThresholds)
+    private void fetchAndCacheThrottlingThresholds(String thresholdJsonPath, String cacheKey) {
+        vertx.executeBlocking(promise -> {
+                    Blob blob = getBlob(thresholdJsonPath);
+                    promise.complete(blob);
+                })
+                .map(blob -> loadThrottlingThresholds((Blob) blob))
                 .onSuccess(thresholds -> cache.put(cacheKey, thresholds))
-                .onFailure(error -> logger.error("Failed to fetch thresholds"));
+                .onFailure(error -> logger.error("Failed to fetch thresholds"));;
     }
 
-    private Blob getBlob() {
+    private Blob getBlob(String thresholdJsonPath) {
         try {
             return Optional.ofNullable(storage.get(gcsBucketName))
-                    .map(bucket -> bucket.get(thresholdPath))
+                    .map(bucket -> bucket.get(thresholdJsonPath))
                     .orElseThrow(() -> new PreBidException("Bucket not found: " + gcsBucketName));
         } catch (StorageException e) {
             throw new PreBidException("Error accessing GCS artefact for threshold: ", e);

extra/modules/greenbids-real-time-data/src/main/java/org/prebid/server/hooks/modules/greenbids/real/time/data/v1/GreenbidsRealTimeDataProcessedAuctionRequestHook.java

@@ -81,7 +81,7 @@ public Future<InvocationResult<AuctionRequestPayload>> call(
 
         return Future.join(
                         onnxModelRunnerWithThresholds.retrieveOnnxModelRunner(partner),
-                        onnxModelRunnerWithThresholds.retrieveThreshold(partner, mapper))
+                        onnxModelRunnerWithThresholds.retrieveThreshold(partner))
                 .compose(compositeFuture -> toInvocationResult(
                         bidRequest,
                         partner,

extra/modules/greenbids-real-time-data/src/test/java/org/prebid/server/hooks/modules/greenbids/real/time/data/v1/GreenbidsRealTimeDataProcessedAuctionRequestHookTest.java

@@ -24,8 +24,10 @@
 import org.prebid.server.hooks.modules.greenbids.real.time.data.core.ThrottlingThresholds;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.data.GreenbidsInferenceDataService;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.FilterService;
+import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.ModelCache;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.OnnxModelRunner;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.OnnxModelRunnerWithThresholds;
+import org.prebid.server.hooks.modules.greenbids.real.time.data.model.predictor.ThresholdCache;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.result.AnalyticsResult;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.result.ExplorationResult;
 import org.prebid.server.hooks.modules.greenbids.real.time.data.model.result.Ortb2ImpExtResult;
@@ -84,14 +86,22 @@ public void setUp() throws IOException {
         final File database = new File("src/test/resources/GeoLite2-Country.mmdb");
         final DatabaseReader dbReader = new DatabaseReader.Builder(database).build();
         final FilterService filterService = new FilterService();
-        final OnnxModelRunnerWithThresholds onnxModelRunnerWithThresholds = new OnnxModelRunnerWithThresholds(
-                modelCacheWithExpiration,
-                thresholdsCacheWithExpiration,
+        final ModelCache modelCache = new ModelCache(
                 storage,
                 "test_bucket",
+                modelCacheWithExpiration,
                 "onnxModelRunner_",
+                Vertx.vertx());
+        final ThresholdCache thresholdCache = new ThresholdCache(
+                storage,
+                "test_bucket",
+                mapper,
+                thresholdsCacheWithExpiration,
                 "throttlingThresholds_",
                 Vertx.vertx());
+        final OnnxModelRunnerWithThresholds onnxModelRunnerWithThresholds = new OnnxModelRunnerWithThresholds(
+                modelCache,
+                thresholdCache);
         final GreenbidsInferenceDataService greenbidsInferenceDataService = new GreenbidsInferenceDataService(
                 dbReader,
                 mapper);