diff --git a/engine/components-api/src/main/java/com/cloud/template/TemplateManager.java b/engine/components-api/src/main/java/com/cloud/template/TemplateManager.java index f1891c774edd..7b2cdf0f6169 100644 --- a/engine/components-api/src/main/java/com/cloud/template/TemplateManager.java +++ b/engine/components-api/src/main/java/com/cloud/template/TemplateManager.java @@ -45,6 +45,8 @@ public interface TemplateManager { static final String AllowPublicUserTemplatesCK = "allow.public.user.templates"; static final String TemplatePreloaderPoolSizeCK = "template.preloader.pool.size"; + static final String PublicTemplateSecStorageCopyCK = "secstorage.public.template.copy.max"; + static final String PrivateTemplateSecStorageCopyCK = "secstorage.private.template.copy.max"; static final ConfigKey AllowPublicUserTemplates = new ConfigKey("Advanced", Boolean.class, AllowPublicUserTemplatesCK, "true", "If false, users will not be able to create public Templates.", true, ConfigKey.Scope.Account); @@ -64,6 +66,18 @@ public interface TemplateManager { true, ConfigKey.Scope.Global); + ConfigKey PublicTemplateSecStorageCopy = new ConfigKey("Advanced", Integer.class, + PublicTemplateSecStorageCopyCK, "0", + "Maximum number of secondary storage pools to which a public template is copied. " + + "0 means copy to all secondary storage pools (default behavior).", + true, ConfigKey.Scope.Zone); + + ConfigKey PrivateTemplateSecStorageCopy = new ConfigKey("Advanced", Integer.class, + PrivateTemplateSecStorageCopyCK, "1", + "Maximum number of secondary storage pools to which a private template is copied. " + + "Default is 1 to preserve existing behavior.", + true, ConfigKey.Scope.Zone); + static final String VMWARE_TOOLS_ISO = "vmware-tools.iso"; static final String XS_TOOLS_ISO = "xs-tools.iso"; diff --git a/engine/storage/image/src/main/java/org/apache/cloudstack/storage/image/TemplateServiceImpl.java b/engine/storage/image/src/main/java/org/apache/cloudstack/storage/image/TemplateServiceImpl.java index e29e89cf431c..b63a3f4dc48a 100644 --- a/engine/storage/image/src/main/java/org/apache/cloudstack/storage/image/TemplateServiceImpl.java +++ b/engine/storage/image/src/main/java/org/apache/cloudstack/storage/image/TemplateServiceImpl.java @@ -295,6 +295,40 @@ public void handleSysTemplateDownload(HypervisorType hostHyper, Long dcId) { } } + private boolean hasReachedSecStorageCopyLimit(VMTemplateVO template, long zoneId) { + boolean isPrivate = !template.isPublicTemplate() && !template.isFeatured() + && !TemplateType.SYSTEM.equals(template.getTemplateType()); + int copyLimit = isPrivate + ? TemplateManager.PrivateTemplateSecStorageCopy.valueIn(zoneId) + : TemplateManager.PublicTemplateSecStorageCopy.valueIn(zoneId); + if (copyLimit <= 0) { + return false; + } + List stores = _storeMgr.getImageStoresByScope(new ZoneScope(zoneId)); + if (stores == null || stores.isEmpty()) { + return false; + } + int count = 0; + for (DataStore ds : stores) { + List rows = _vmTemplateStoreDao.listByTemplateStore(template.getId(), ds.getId()); + if (rows == null) { + continue; + } + for (TemplateDataStoreVO row : rows) { + State st = row.getState(); + Status ds_state = row.getDownloadState(); + if (st != State.Failed && st != State.Destroyed + && ds_state != Status.ABANDONED && ds_state != Status.DOWNLOAD_ERROR) { + count++; + break; + } + } + } + logger.debug("Template [{}] secstorage copy check in zone [{}]: count={}, limit={}", + template.getUniqueName(), zoneId, count, copyLimit); + return count >= copyLimit; + } + protected boolean shouldDownloadTemplateToStore(VMTemplateVO template, DataStore store) { Long zoneId = store.getScope().getScopeId(); DataStore directedStore = _tmpltMgr.verifyHeuristicRulesForZone(template, zoneId); @@ -304,6 +338,12 @@ protected boolean shouldDownloadTemplateToStore(VMTemplateVO template, DataStore return false; } + if (zoneId != null && hasReachedSecStorageCopyLimit(template, zoneId)) { + logger.info("Skipping sync of template [{}] to image store [{}]: zone [{}] has reached the configured copy limit.", + template.getUniqueName(), store.getName(), zoneId); + return false; + } + if (template.isPublicTemplate()) { logger.debug("Download of template [{}] to image store [{}] cannot be skipped, as it is public.", template.getUniqueName(), store.getName()); @@ -531,10 +571,13 @@ public void handleTemplateSync(DataStore store) { && tmpltStore.getState() == State.Ready && tmpltStore.getInstallPath() == null) { logger.info("Keep fake entry in template store table for migration of previous NFS to object store"); - } else { + } else if (tmpltStore.getDownloadState() == VMTemplateStorageResourceAssoc.Status.DOWNLOADED + || tmpltStore.getState() == State.Ready) { logger.info("Removing leftover template {} entry from template store table", tmplt); - // remove those leftover entries _vmTemplateStoreDao.remove(tmpltStore.getId()); + } else { + logger.debug("Template {} entry on store {} is in pre-download state ({}/{}); not treating as leftover.", + tmplt, store, tmpltStore.getState(), tmpltStore.getDownloadState()); } } } diff --git a/server/src/main/java/com/cloud/template/HypervisorTemplateAdapter.java b/server/src/main/java/com/cloud/template/HypervisorTemplateAdapter.java index a7a6cea7d907..8df7b9b1704c 100644 --- a/server/src/main/java/com/cloud/template/HypervisorTemplateAdapter.java +++ b/server/src/main/java/com/cloud/template/HypervisorTemplateAdapter.java @@ -19,10 +19,10 @@ import java.util.ArrayList; import java.util.Collections; import java.util.Date; -import java.util.HashSet; +import java.util.HashMap; import java.util.LinkedList; import java.util.List; -import java.util.Set; +import java.util.Map; import java.util.concurrent.ExecutionException; import java.util.stream.Collectors; @@ -292,9 +292,10 @@ protected void createTemplateWithinZones(TemplateProfile profile, VMTemplateVO t if (imageStore == null) { List imageStores = getImageStoresThrowsExceptionIfNotFound(zoneId, profile); - standardImageStoreAllocation(imageStores, template); + standardImageStoreAllocation(imageStores, template, zoneId); } else { - validateSecondaryStorageAndCreateTemplate(List.of(imageStore), template, null); + int copyLimit = getSecStorageCopyLimit(template, zoneId); + validateSecondaryStorageAndCreateTemplate(List.of(imageStore), template, new HashMap<>(), copyLimit); } } } @@ -307,17 +308,17 @@ protected List getImageStoresThrowsExceptionIfNotFound(long zoneId, T return imageStores; } - protected void standardImageStoreAllocation(List imageStores, VMTemplateVO template) { - Set zoneSet = new HashSet(); + protected void standardImageStoreAllocation(List imageStores, VMTemplateVO template, long zoneId) { + int copyLimit = getSecStorageCopyLimit(template, zoneId); Collections.shuffle(imageStores); - validateSecondaryStorageAndCreateTemplate(imageStores, template, zoneSet); + validateSecondaryStorageAndCreateTemplate(imageStores, template, new HashMap<>(), copyLimit); } - protected void validateSecondaryStorageAndCreateTemplate(List imageStores, VMTemplateVO template, Set zoneSet) { + protected void validateSecondaryStorageAndCreateTemplate(List imageStores, VMTemplateVO template, Map zoneCopyCount, int copyLimit) { for (DataStore imageStore : imageStores) { Long zoneId = imageStore.getScope().getScopeId(); - if (!isZoneAndImageStoreAvailable(imageStore, zoneId, zoneSet, isPrivateTemplate(template))) { + if (!isZoneAndImageStoreAvailable(imageStore, zoneId, zoneCopyCount, copyLimit)) { continue; } diff --git a/server/src/main/java/com/cloud/template/TemplateAdapterBase.java b/server/src/main/java/com/cloud/template/TemplateAdapterBase.java index 97c6e5903592..4ff2cec0e9d2 100644 --- a/server/src/main/java/com/cloud/template/TemplateAdapterBase.java +++ b/server/src/main/java/com/cloud/template/TemplateAdapterBase.java @@ -20,11 +20,9 @@ import java.util.Arrays; import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; -import java.util.Set; import javax.inject.Inject; @@ -169,7 +167,13 @@ protected DataStore verifyHeuristicRulesForZone(VMTemplateVO template, Long zone return heuristicRuleHelper.getImageStoreIfThereIsHeuristicRule(zoneId, heuristicType, template); } - protected boolean isZoneAndImageStoreAvailable(DataStore imageStore, Long zoneId, Set zoneSet, boolean isTemplatePrivate) { + protected int getSecStorageCopyLimit(VMTemplateVO template, long zoneId) { + return isPrivateTemplate(template) + ? TemplateManager.PrivateTemplateSecStorageCopy.valueIn(zoneId) + : TemplateManager.PublicTemplateSecStorageCopy.valueIn(zoneId); + } + + protected boolean isZoneAndImageStoreAvailable(DataStore imageStore, Long zoneId, Map zoneCopyCount, int copyLimit) { if (zoneId == null) { logger.warn(String.format("Zone ID is null, cannot allocate ISO/template in image store [%s].", imageStore)); return false; @@ -191,19 +195,13 @@ protected boolean isZoneAndImageStoreAvailable(DataStore imageStore, Long zoneId return false; } - if (zoneSet == null) { - logger.info(String.format("Zone set is null; therefore, the ISO/template should be allocated in every secondary storage of zone [%s].", zone)); - return true; - } - - if (isTemplatePrivate && zoneSet.contains(zoneId)) { - logger.info(String.format("The template is private and it is already allocated in a secondary storage in zone [%s]; therefore, image store [%s] will be skipped.", - zone, imageStore)); + int currentCount = zoneCopyCount.getOrDefault(zoneId, 0); + if (copyLimit > 0 && currentCount >= copyLimit) { + logger.info("Copy limit of {} reached for zone [{}]; skipping image store [{}].", copyLimit, zone, imageStore); return false; } - logger.info(String.format("Private template will be allocated in image store [%s] in zone [%s].", imageStore, zone)); - zoneSet.add(zoneId); + zoneCopyCount.put(zoneId, currentCount + 1); return true; } @@ -212,12 +210,13 @@ protected boolean isZoneAndImageStoreAvailable(DataStore imageStore, Long zoneId * {@link TemplateProfile#getZoneIdList()}. */ protected void postUploadAllocation(List imageStores, VMTemplateVO template, List payloads) { - Set zoneSet = new HashSet<>(); + Map zoneCopyCount = new HashMap<>(); Collections.shuffle(imageStores); for (DataStore imageStore : imageStores) { Long zoneId_is = imageStore.getScope().getScopeId(); + int copyLimit = zoneId_is == null ? 0 : getSecStorageCopyLimit(template, zoneId_is); - if (!isZoneAndImageStoreAvailable(imageStore, zoneId_is, zoneSet, isPrivateTemplate(template))) { + if (!isZoneAndImageStoreAvailable(imageStore, zoneId_is, zoneCopyCount, copyLimit)) { continue; } diff --git a/server/src/main/java/com/cloud/template/TemplateManagerImpl.java b/server/src/main/java/com/cloud/template/TemplateManagerImpl.java index 67f7128e864d..3ad26e09f117 100755 --- a/server/src/main/java/com/cloud/template/TemplateManagerImpl.java +++ b/server/src/main/java/com/cloud/template/TemplateManagerImpl.java @@ -2493,7 +2493,9 @@ public ConfigKey[] getConfigKeys() { return new ConfigKey[] {AllowPublicUserTemplates, TemplatePreloaderPoolSize, ValidateUrlIsResolvableBeforeRegisteringTemplate, - TemplateDeleteFromPrimaryStorage}; + TemplateDeleteFromPrimaryStorage, + PublicTemplateSecStorageCopy, + PrivateTemplateSecStorageCopy}; } public List getTemplateAdapters() { diff --git a/server/src/test/java/com/cloud/template/HypervisorTemplateAdapterTest.java b/server/src/test/java/com/cloud/template/HypervisorTemplateAdapterTest.java index e2a97be469ff..4cd48e686b01 100644 --- a/server/src/test/java/com/cloud/template/HypervisorTemplateAdapterTest.java +++ b/server/src/test/java/com/cloud/template/HypervisorTemplateAdapterTest.java @@ -32,10 +32,8 @@ import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; -import java.util.Set; import java.util.concurrent.ExecutionException; import org.apache.cloudstack.engine.subsystem.api.storage.DataStore; @@ -339,7 +337,7 @@ public void createTemplateWithinZonesTestZoneIdsNotNullShouldNotCallListAllZones Mockito.when(templateProfileMock.getZoneIdList()).thenReturn(zoneIds); Mockito.doReturn(null).when(_adapter).getImageStoresThrowsExceptionIfNotFound(Mockito.any(Long.class), Mockito.any(TemplateProfile.class)); Mockito.doReturn(null).when(_templateMgr).verifyHeuristicRulesForZone(Mockito.any(VMTemplateVO.class), Mockito.anyLong()); - Mockito.doNothing().when(_adapter).standardImageStoreAllocation(Mockito.isNull(), Mockito.any(VMTemplateVO.class)); + Mockito.doNothing().when(_adapter).standardImageStoreAllocation(Mockito.isNull(), Mockito.any(VMTemplateVO.class), Mockito.anyLong()); _adapter.createTemplateWithinZones(templateProfileMock, vmTemplateVOMock); @@ -355,11 +353,11 @@ public void createTemplateWithinZonesTestZoneDoesNotHaveActiveHeuristicRulesShou Mockito.when(templateProfileMock.getZoneIdList()).thenReturn(zoneIds); Mockito.doReturn(null).when(_adapter).getImageStoresThrowsExceptionIfNotFound(Mockito.any(Long.class), Mockito.any(TemplateProfile.class)); Mockito.doReturn(null).when(_templateMgr).verifyHeuristicRulesForZone(Mockito.any(VMTemplateVO.class), Mockito.anyLong()); - Mockito.doNothing().when(_adapter).standardImageStoreAllocation(Mockito.isNull(), Mockito.any(VMTemplateVO.class)); + Mockito.doNothing().when(_adapter).standardImageStoreAllocation(Mockito.isNull(), Mockito.any(VMTemplateVO.class), Mockito.anyLong()); _adapter.createTemplateWithinZones(templateProfileMock, vmTemplateVOMock); - Mockito.verify(_adapter, Mockito.times(1)).standardImageStoreAllocation(Mockito.isNull(), Mockito.any(VMTemplateVO.class)); + Mockito.verify(_adapter, Mockito.times(1)).standardImageStoreAllocation(Mockito.isNull(), Mockito.any(VMTemplateVO.class), Mockito.anyLong()); } @Test @@ -371,11 +369,11 @@ public void createTemplateWithinZonesTestZoneWithHeuristicRuleShouldCallValidate Mockito.when(templateProfileMock.getZoneIdList()).thenReturn(zoneIds); Mockito.doReturn(dataStoreMock).when(_templateMgr).verifyHeuristicRulesForZone(Mockito.any(VMTemplateVO.class), Mockito.anyLong()); - Mockito.doNothing().when(_adapter).validateSecondaryStorageAndCreateTemplate(Mockito.any(List.class), Mockito.any(VMTemplateVO.class), Mockito.isNull()); + Mockito.doNothing().when(_adapter).validateSecondaryStorageAndCreateTemplate(Mockito.any(List.class), Mockito.any(VMTemplateVO.class), Mockito.any(Map.class), Mockito.anyInt()); _adapter.createTemplateWithinZones(templateProfileMock, vmTemplateVOMock); - Mockito.verify(_adapter, Mockito.times(1)).validateSecondaryStorageAndCreateTemplate(Mockito.any(List.class), Mockito.any(VMTemplateVO.class), Mockito.isNull()); + Mockito.verify(_adapter, Mockito.times(1)).validateSecondaryStorageAndCreateTemplate(Mockito.any(List.class), Mockito.any(VMTemplateVO.class), Mockito.any(Map.class), Mockito.anyInt()); } @Test(expected = CloudRuntimeException.class) @@ -411,11 +409,8 @@ public void getImageStoresThrowsExceptionIfNotFoundTestNonEmptyImageStoreShouldN @Test public void isZoneAndImageStoreAvailableTestZoneIdIsNullShouldReturnFalse() { DataStore dataStoreMock = Mockito.mock(DataStore.class); - Long zoneId = null; - Set zoneSet = null; - boolean isTemplatePrivate = false; - boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneSet, isTemplatePrivate); + boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, null, new HashMap<>(), 0); Mockito.verify(loggerMock, Mockito.times(1)).warn(String.format("Zone ID is null, cannot allocate ISO/template in image store [%s].", dataStoreMock)); Assert.assertFalse(result); @@ -425,13 +420,10 @@ public void isZoneAndImageStoreAvailableTestZoneIdIsNullShouldReturnFalse() { public void isZoneAndImageStoreAvailableTestZoneIsNullShouldReturnFalse() { DataStore dataStoreMock = Mockito.mock(DataStore.class); Long zoneId = 1L; - Set zoneSet = null; - boolean isTemplatePrivate = false; - DataCenterVO dataCenterVOMock = null; - Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(dataCenterVOMock); + Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(null); - boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneSet, isTemplatePrivate); + boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, new HashMap<>(), 0); Mockito.verify(loggerMock, Mockito.times(1)).warn("Unable to find zone by id [{}], so skip downloading template to its image store [{}].", zoneId, dataStoreMock); @@ -442,14 +434,12 @@ public void isZoneAndImageStoreAvailableTestZoneIsNullShouldReturnFalse() { public void isZoneAndImageStoreAvailableTestZoneIsDisabledShouldReturnFalse() { DataStore dataStoreMock = Mockito.mock(DataStore.class); Long zoneId = 1L; - Set zoneSet = null; - boolean isTemplatePrivate = false; DataCenterVO dataCenterVOMock = Mockito.mock(DataCenterVO.class); Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(dataCenterVOMock); Mockito.when(dataCenterVOMock.getAllocationState()).thenReturn(Grouping.AllocationState.Disabled); - boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneSet, isTemplatePrivate); + boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, new HashMap<>(), 0); Mockito.verify(loggerMock, Mockito.times(1)).info("Zone [{}] is disabled. Skip downloading template to its image store [{}].", dataCenterVOMock, dataStoreMock); Assert.assertFalse(result); @@ -459,15 +449,13 @@ public void isZoneAndImageStoreAvailableTestZoneIsDisabledShouldReturnFalse() { public void isZoneAndImageStoreAvailableTestImageStoreDoesNotHaveEnoughCapacityShouldReturnFalse() { DataStore dataStoreMock = Mockito.mock(DataStore.class); Long zoneId = 1L; - Set zoneSet = null; - boolean isTemplatePrivate = false; DataCenterVO dataCenterVOMock = Mockito.mock(DataCenterVO.class); Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(dataCenterVOMock); Mockito.when(dataCenterVOMock.getAllocationState()).thenReturn(Grouping.AllocationState.Enabled); Mockito.when(statsCollectorMock.imageStoreHasEnoughCapacity(any(DataStore.class))).thenReturn(false); - boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneSet, isTemplatePrivate); + boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, new HashMap<>(), 0); Mockito.verify(loggerMock, times(1)).info("Image store doesn't have enough capacity. Skip downloading template to this image store [{}].", dataStoreMock); @@ -475,60 +463,72 @@ public void isZoneAndImageStoreAvailableTestImageStoreDoesNotHaveEnoughCapacityS } @Test - public void isZoneAndImageStoreAvailableTestImageStoreHasEnoughCapacityAndZoneSetIsNullShouldReturnTrue() { + public void isZoneAndImageStoreAvailableTestReplicaLimitZeroShouldCopyToAllStores() { DataStore dataStoreMock = Mockito.mock(DataStore.class); Long zoneId = 1L; - Set zoneSet = null; - boolean isTemplatePrivate = false; DataCenterVO dataCenterVOMock = Mockito.mock(DataCenterVO.class); + Map zoneCopyCount = new HashMap<>(); + zoneCopyCount.put(zoneId, 999); Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(dataCenterVOMock); Mockito.when(dataCenterVOMock.getAllocationState()).thenReturn(Grouping.AllocationState.Enabled); Mockito.when(statsCollectorMock.imageStoreHasEnoughCapacity(any(DataStore.class))).thenReturn(true); - boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneSet, isTemplatePrivate); + boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneCopyCount, 0); - Mockito.verify(loggerMock, times(1)).info(String.format("Zone set is null; therefore, the ISO/template should be allocated in every secondary storage " + - "of zone [%s].", dataCenterVOMock)); Assert.assertTrue(result); + Assert.assertEquals(1000, (int) zoneCopyCount.get(zoneId)); } @Test - public void isZoneAndImageStoreAvailableTestTemplateIsPrivateAndItIsAlreadyAllocatedToTheSameZoneShouldReturnFalse() { + public void isZoneAndImageStoreAvailableTestReplicaLimitReachedShouldReturnFalse() { DataStore dataStoreMock = Mockito.mock(DataStore.class); Long zoneId = 1L; - Set zoneSet = Set.of(1L); - boolean isTemplatePrivate = true; DataCenterVO dataCenterVOMock = Mockito.mock(DataCenterVO.class); + Map zoneCopyCount = new HashMap<>(); + zoneCopyCount.put(zoneId, 1); Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(dataCenterVOMock); Mockito.when(dataCenterVOMock.getAllocationState()).thenReturn(Grouping.AllocationState.Enabled); Mockito.when(statsCollectorMock.imageStoreHasEnoughCapacity(any(DataStore.class))).thenReturn(true); - boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneSet, isTemplatePrivate); + boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneCopyCount, 1); - Mockito.verify(loggerMock, times(1)).info(String.format("The template is private and it is already allocated in a secondary storage in zone [%s]; " + - "therefore, image store [%s] will be skipped.", dataCenterVOMock, dataStoreMock)); + Mockito.verify(loggerMock, times(1)).info("Copy limit of {} reached for zone [{}]; skipping image store [{}].", 1, dataCenterVOMock, dataStoreMock); Assert.assertFalse(result); } @Test - public void isZoneAndImageStoreAvailableTestTemplateIsPrivateAndItIsNotAlreadyAllocatedToTheSameZoneShouldReturnTrue() { + public void isZoneAndImageStoreAvailableTestReplicaLimitNotYetReachedShouldReturnTrueAndIncrementCount() { DataStore dataStoreMock = Mockito.mock(DataStore.class); Long zoneId = 1L; - Set zoneSet = new HashSet<>(); - boolean isTemplatePrivate = true; DataCenterVO dataCenterVOMock = Mockito.mock(DataCenterVO.class); + Map zoneCopyCount = new HashMap<>(); Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(dataCenterVOMock); Mockito.when(dataCenterVOMock.getAllocationState()).thenReturn(Grouping.AllocationState.Enabled); Mockito.when(statsCollectorMock.imageStoreHasEnoughCapacity(any(DataStore.class))).thenReturn(true); - boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneSet, isTemplatePrivate); + boolean result = _adapter.isZoneAndImageStoreAvailable(dataStoreMock, zoneId, zoneCopyCount, 2); - Mockito.verify(loggerMock, times(1)).info(String.format("Private template will be allocated in image store [%s] in zone [%s].", - dataStoreMock, dataCenterVOMock)); Assert.assertTrue(result); + Assert.assertEquals(1, (int) zoneCopyCount.get(zoneId)); + } + + @Test + public void isZoneAndImageStoreAvailableTestReplicaLimitOfTwoShouldCopyToExactlyTwoStores() { + Long zoneId = 1L; + DataCenterVO dataCenterVOMock = Mockito.mock(DataCenterVO.class); + Map zoneCopyCount = new HashMap<>(); + + Mockito.when(_dcDao.findById(Mockito.anyLong())).thenReturn(dataCenterVOMock); + Mockito.when(dataCenterVOMock.getAllocationState()).thenReturn(Grouping.AllocationState.Enabled); + Mockito.when(statsCollectorMock.imageStoreHasEnoughCapacity(any(DataStore.class))).thenReturn(true); + + Assert.assertTrue(_adapter.isZoneAndImageStoreAvailable(Mockito.mock(DataStore.class), zoneId, zoneCopyCount, 2)); + Assert.assertTrue(_adapter.isZoneAndImageStoreAvailable(Mockito.mock(DataStore.class), zoneId, zoneCopyCount, 2)); + Assert.assertFalse(_adapter.isZoneAndImageStoreAvailable(Mockito.mock(DataStore.class), zoneId, zoneCopyCount, 2)); + Assert.assertEquals(2, (int) zoneCopyCount.get(zoneId)); } @Test