diff --git a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java index 3cdaf661a6b..dd2a364e4db 100644 --- a/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java +++ b/clients/da-vinci-client/src/main/java/com/linkedin/davinci/DaVinciBackend.java @@ -401,7 +401,7 @@ private Function functionToCheckWhetherStorageEngineShouldBeKep private synchronized void bootstrap() { List storageEngines = - storageService.getStorageEngineRepository().getAllLocalStorageEngines(); + getStorageService().getStorageEngineRepository().getAllLocalStorageEngines(); LOGGER.info("Starting bootstrap, storageEngines: {}", storageEngines); Map storeNameToBootstrapVersionMap = new HashMap<>(); Map> storeNameToPartitionListMap = new HashMap<>(); @@ -440,7 +440,7 @@ private synchronized void bootstrap() { if (!(storeNameToBootstrapVersionMap.containsKey(storeName) && (storeNameToBootstrapVersionMap.get(storeName).getNumber() < versionNumber))) { storeNameToBootstrapVersionMap.put(storeName, version); - storeNameToPartitionListMap.put(storeName, storageService.getUserPartitions(kafkaTopicName)); + storeNameToPartitionListMap.put(storeName, getStorageService().getUserPartitions(kafkaTopicName)); } } @@ -455,12 +455,12 @@ private synchronized void bootstrap() { * In this case we will only need to close metadata partition, as it is supposed to be opened and managed by * forked ingestion process via following subscribe call. */ - for (AbstractStorageEngine storageEngine: storageService.getStorageEngineRepository() + for (AbstractStorageEngine storageEngine: getStorageService().getStorageEngineRepository() .getAllLocalStorageEngines()) { storageEngine.closeMetadataPartition(); } } else { - storageService.closeAllStorageEngines(); + getStorageService().closeAllStorageEngines(); } } @@ -470,13 +470,13 @@ private synchronized void bootstrap() { metricsRepository, storageMetadataService, ingestionService, - storageService, + getStorageService(), blobTransferManager, this::getVeniceCurrentVersionNumber) : new DefaultIngestionBackend( storageMetadataService, ingestionService, - storageService, + getStorageService(), blobTransferManager, configLoader.getVeniceServerConfig()); ingestionBackend.addIngestionNotifier(ingestionListener); @@ -487,7 +487,7 @@ private synchronized void bootstrap() { List partitions = storeNameToPartitionListMap.get(storeName); String versionTopic = version.kafkaTopicName(); LOGGER.info("Bootstrapping partitions {} for {}", partitions, versionTopic); - AbstractStorageEngine storageEngine = storageService.getStorageEngine(versionTopic); + AbstractStorageEngine storageEngine = getStorageService().getStorageEngine(versionTopic); aggVersionedStorageEngineStats.setStorageEngine(versionTopic, storageEngine); StoreBackend storeBackend = getStoreOrThrow(storeName); storeBackend.subscribe(ComplementSet.newSet(partitions), Optional.of(version)); diff --git a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/DaVinciBackendTest.java b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/DaVinciBackendTest.java index 375fce76d4f..07cac5df8b5 100644 --- a/clients/da-vinci-client/src/test/java/com/linkedin/davinci/DaVinciBackendTest.java +++ b/clients/da-vinci-client/src/test/java/com/linkedin/davinci/DaVinciBackendTest.java @@ -189,6 +189,7 @@ public void testBootstrappingSubscription() List localStorageEngines = new ArrayList<>(); localStorageEngines.add(abstractStorageEngine); + when(backend.getStorageService()).thenReturn(mockStorageService); when(mockStorageService.getStorageEngineRepository()).thenReturn(mockStorageEngineRepository); when(mockStorageService.getStorageEngine(resourceName)).thenReturn(abstractStorageEngine); when(mockStorageEngineRepository.getAllLocalStorageEngines()).thenReturn(localStorageEngines);