use of com.github.ambry.account.InMemAccountService in project ambry by linkedin.
the class StorageQuotaEnforcerTest method testInitStorageUsage.
/**
* Test to initialize the storage usage with non-empty map.
*/
@Test
public void testInitStorageUsage() throws Exception {
Map<String, Map<String, Long>> containerUsage = TestUtils.makeStorageMap(10, 10, 1000, 100);
InMemAccountService accountService = new InMemAccountService(false, false);
Map<QuotaResource, Long> expectedStorageUsages = new HashMap<>();
// Account and container id's base is 1, not 0
for (int i = 1; i <= 10; i++) {
QuotaResourceType resourceType = i <= containerUsage.size() / 2 ? QuotaResourceType.CONTAINER : QuotaResourceType.ACCOUNT;
AccountBuilder accountBuilder = new AccountBuilder((short) i, String.valueOf(i), Account.AccountStatus.ACTIVE, resourceType);
for (int j = 1; j <= 10; j++) {
accountBuilder.addOrUpdateContainer(new ContainerBuilder((short) j, String.valueOf(j), Container.ContainerStatus.ACTIVE, "", (short) i).build());
}
accountService.updateAccounts(Collections.singleton(accountBuilder.build()));
if (resourceType == QuotaResourceType.ACCOUNT) {
expectedStorageUsages.put(QuotaResource.fromAccountId((short) i), containerUsage.get(String.valueOf(i)).values().stream().mapToLong(Long::longValue).sum());
} else {
for (Map.Entry<String, Long> containerEntry : containerUsage.get(String.valueOf(i)).entrySet()) {
expectedStorageUsages.put(QuotaResource.fromContainerId((short) i, Short.valueOf(containerEntry.getKey())), containerEntry.getValue());
}
}
}
StorageQuotaEnforcer enforcer = new StorageQuotaEnforcer(config, new JSONStringStorageQuotaSource(new HashMap<>(), accountService), (StorageUsageRefresher) null);
enforcer.initStorageUsage(containerUsage);
assertEquals(expectedStorageUsages, enforcer.getStorageUsages());
}
use of com.github.ambry.account.InMemAccountService in project ambry by linkedin.
the class StorageQuotaEnforcerTest method testGetQuotaAndUsageAndCharge.
/**
* Test {@link StorageQuotaEnforcer#getQuotaAndUsage} and {@link StorageQuotaEnforcer#charge} methods.
* @throws Exception
*/
@Test
public void testGetQuotaAndUsageAndCharge() throws Exception {
int initNumAccounts = 10;
Map<String, Map<String, Long>> containerUsage = TestUtils.makeStorageMap(initNumAccounts, 10, 10000, 1000);
InMemAccountService accountService = new InMemAccountService(false, false);
Map<String, Long> storageQuota = new HashMap<>();
// Account and container id's base is 1, not 0
for (int i = 1; i <= initNumAccounts; i++) {
QuotaResourceType resourceType = i <= containerUsage.size() / 2 ? QuotaResourceType.CONTAINER : QuotaResourceType.ACCOUNT;
AccountBuilder accountBuilder = new AccountBuilder((short) i, String.valueOf(i), Account.AccountStatus.ACTIVE, resourceType);
for (int j = 1; j <= 10; j++) {
accountBuilder.addOrUpdateContainer(new ContainerBuilder((short) j, String.valueOf(j), Container.ContainerStatus.ACTIVE, "", (short) i).build());
}
accountService.updateAccounts(Collections.singleton(accountBuilder.build()));
if (resourceType == QuotaResourceType.ACCOUNT) {
storageQuota.put(QuotaResource.fromAccount(accountService.getAccountById((short) i)).getResourceId(), containerUsage.get(String.valueOf(i)).values().stream().mapToLong(Long::longValue).sum());
} else {
accountService.getAccountById((short) i).getAllContainers().forEach(c -> storageQuota.put(QuotaResource.fromContainer(c).getResourceId(), containerUsage.get(String.valueOf(c.getParentAccountId())).get(String.valueOf(c.getId()))));
}
}
JSONStringStorageQuotaSource quotaSource = new JSONStringStorageQuotaSource(storageQuota, accountService);
StorageQuotaEnforcer enforcer = new StorageQuotaEnforcer(config, quotaSource, (StorageUsageRefresher) null);
enforcer.initStorageUsage(Collections.EMPTY_MAP);
for (Map.Entry<String, Map<String, Long>> accountEntry : containerUsage.entrySet()) {
short accountId = Short.valueOf(accountEntry.getKey());
if (accountService.getAccountById(accountId).getQuotaResourceType() == QuotaResourceType.ACCOUNT) {
long quota = (long) (quotaSource.getQuota(QuotaResource.fromAccount(accountService.getAccountById(accountId)), QuotaName.STORAGE_IN_GB).getQuotaValue()) * BYTES_IN_GB;
RestRequest restRequest = createRestRequest(accountService, accountId, (short) 1);
Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
assertEquals("Account id: " + accountEntry.getKey(), quota, quotaAndUsage.getFirst().longValue());
assertEquals(0L, quotaAndUsage.getSecond().longValue());
quotaAndUsage = enforcer.charge(restRequest, quota / 2);
assertEquals(quota, quotaAndUsage.getFirst().longValue());
assertEquals(quota / 2, quotaAndUsage.getSecond().longValue());
quotaAndUsage = enforcer.charge(restRequest, quota);
assertEquals(quota, quotaAndUsage.getFirst().longValue());
assertEquals(quota / 2 + quota, quotaAndUsage.getSecond().longValue());
} else {
for (Map.Entry<String, Long> containerEntry : accountEntry.getValue().entrySet()) {
short containerId = Short.valueOf(containerEntry.getKey());
long quota = containerEntry.getValue() * BYTES_IN_GB;
RestRequest restRequest = createRestRequest(accountService, accountId, containerId);
Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
assertEquals(quota, quotaAndUsage.getFirst().longValue());
assertEquals(0L, quotaAndUsage.getSecond().longValue());
quotaAndUsage = enforcer.charge(restRequest, quota / 2);
assertEquals(quota, quotaAndUsage.getFirst().longValue());
assertEquals(quota / 2, quotaAndUsage.getSecond().longValue());
quotaAndUsage = enforcer.charge(restRequest, quota);
assertEquals(quota, quotaAndUsage.getFirst().longValue());
assertEquals(quota / 2 + quota, quotaAndUsage.getSecond().longValue());
}
}
}
// Now create a restRequest that doesn't carry account and container
RestRequest restRequest = createRestRequest();
Pair<Long, Long> quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
assertEquals(-1L, quotaAndUsage.getFirst().longValue());
assertEquals(0L, quotaAndUsage.getSecond().longValue());
quotaAndUsage = enforcer.charge(restRequest, 100L);
assertEquals(-1L, quotaAndUsage.getFirst().longValue());
assertEquals(0L, quotaAndUsage.getSecond().longValue());
Account account = new AccountBuilder((short) 1000, String.valueOf(1000), Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 10000, String.valueOf(10000), Container.ContainerStatus.ACTIVE, "", (short) 1000).build()).build();
accountService.updateAccounts(Collections.singleton(account));
restRequest = createRestRequest(accountService, (short) 1000, (short) 10000);
quotaAndUsage = enforcer.getQuotaAndUsage(restRequest);
assertEquals(-1L, quotaAndUsage.getFirst().longValue());
assertEquals(0L, quotaAndUsage.getSecond().longValue());
quotaAndUsage = enforcer.charge(restRequest, 100L);
assertEquals(-1L, quotaAndUsage.getFirst().longValue());
assertEquals(0L, quotaAndUsage.getSecond().longValue());
}
use of com.github.ambry.account.InMemAccountService in project ambry by linkedin.
the class JsonCUQuotaDataProviderUtilTest method testGetCUQuotasFromJsonForEmptyJsonString.
@Test
public void testGetCUQuotasFromJsonForEmptyJsonString() throws IOException, AccountServiceException {
InMemAccountService accountService = new InMemAccountService(false, false);
ObjectMapper objectMapper = new ObjectMapper();
Map<String, JsonCUQuotaDataProviderUtil.MapOrQuota> testQuotas = objectMapper.readValue(EMPTY_JSON, new TypeReference<Map<String, JsonCUQuotaDataProviderUtil.MapOrQuota>>() {
});
for (String s : testQuotas.keySet()) {
accountService.updateAccounts(Collections.singletonList(createAccountForQuota(testQuotas.get(s), s)));
}
Map<String, CapacityUnit> quotas = JsonCUQuotaDataProviderUtil.getCUQuotasFromJson(EMPTY_JSON, accountService);
Assert.assertEquals(0, quotas.size());
}
use of com.github.ambry.account.InMemAccountService in project ambry by linkedin.
the class JSONStringStorageQuotaSourceTest method testJSONStringStorageQuotaSource.
@Test
public void testJSONStringStorageQuotaSource() throws Exception {
// Trick to create a string literal without escape.
String json = "{`10`: {`1`: 1000, `2`: 3000}, `20`: {`4`: 2000, `5`: 1000}, `30`: 4000, `40`: 5000, `50`: {`6`: 6000}}".replace("`", "\"");
Properties properties = new Properties();
properties.setProperty(StorageQuotaConfig.STORAGE_QUOTA_IN_JSON, json);
StorageQuotaConfig config = new StorageQuotaConfig(new VerifiableProperties(properties));
// Setting up accounts and account service
InMemAccountService accountService = new InMemAccountService(false, false);
Account account = new AccountBuilder((short) 10, "10", Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 1, "1", Container.ContainerStatus.ACTIVE, "", (short) 10).build()).addOrUpdateContainer(new ContainerBuilder((short) 2, "2", Container.ContainerStatus.ACTIVE, "", (short) 10).build()).build();
accountService.updateAccounts(Collections.singleton(account));
account = new AccountBuilder((short) 20, "20", Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 4, "4", Container.ContainerStatus.ACTIVE, "", (short) 20).build()).addOrUpdateContainer(new ContainerBuilder((short) 5, "5", Container.ContainerStatus.ACTIVE, "", (short) 20).build()).build();
accountService.updateAccounts(Collections.singleton(account));
account = new AccountBuilder((short) 30, "30", Account.AccountStatus.ACTIVE, QuotaResourceType.ACCOUNT).addOrUpdateContainer(new ContainerBuilder((short) 4, "4", Container.ContainerStatus.ACTIVE, "", (short) 30).build()).build();
accountService.updateAccounts(Collections.singleton(account));
account = new AccountBuilder((short) 40, "40", Account.AccountStatus.ACTIVE, QuotaResourceType.ACCOUNT).addOrUpdateContainer(new ContainerBuilder((short) 4, "4", Container.ContainerStatus.ACTIVE, "", (short) 40).build()).build();
accountService.updateAccounts(Collections.singleton(account));
account = new AccountBuilder((short) 50, "50", Account.AccountStatus.ACTIVE, QuotaResourceType.CONTAINER).addOrUpdateContainer(new ContainerBuilder((short) 6, "6", Container.ContainerStatus.ACTIVE, "", (short) 50).build()).build();
accountService.updateAccounts(Collections.singleton(account));
JSONStringStorageQuotaSource source = new JSONStringStorageQuotaSource(config, accountService);
QuotaResourceType resourceType = QuotaResourceType.CONTAINER;
Quota quota = source.getQuota(new QuotaResource("1000_1", resourceType), QuotaName.STORAGE_IN_GB);
assertNull(quota);
quota = source.getQuota(new QuotaResource("10_1", resourceType), QuotaName.STORAGE_IN_GB);
assertEquals(1000L, (long) quota.getQuotaValue());
quota = source.getQuota(new QuotaResource("10_2", resourceType), QuotaName.STORAGE_IN_GB);
assertEquals(3000L, (long) quota.getQuotaValue());
quota = source.getQuota(new QuotaResource("20_4", resourceType), QuotaName.STORAGE_IN_GB);
assertEquals(2000L, (long) quota.getQuotaValue());
quota = source.getQuota(new QuotaResource("20_5", resourceType), QuotaName.STORAGE_IN_GB);
assertEquals(1000L, (long) quota.getQuotaValue());
quota = source.getQuota(new QuotaResource("30", QuotaResourceType.ACCOUNT), QuotaName.STORAGE_IN_GB);
assertEquals(4000L, (long) quota.getQuotaValue());
quota = source.getQuota(new QuotaResource("40", QuotaResourceType.ACCOUNT), QuotaName.STORAGE_IN_GB);
assertEquals(5000L, (long) quota.getQuotaValue());
quota = source.getQuota(new QuotaResource("50_6", resourceType), QuotaName.STORAGE_IN_GB);
assertEquals(6000L, (long) quota.getQuotaValue());
}
use of com.github.ambry.account.InMemAccountService in project ambry by linkedin.
the class ReplicationTest method onReplicaAddedOrRemovedCallbackTest.
/**
* Test cluster map change callback in {@link ReplicationManager} when any remote replicas are added or removed.
* Test setup: attempt to add 3 replicas and remove 3 replicas respectively. The three replicas are picked as follows:
* (1) 1st replica on current node (should skip)
* (2) 2nd replica on remote node sharing partition with current one (should be added or removed)
* (3) 3rd replica on remote node but doesn't share partition with current one (should skip)
* @throws Exception
*/
@Test
public void onReplicaAddedOrRemovedCallbackTest() throws Exception {
MockClusterMap clusterMap = new MockClusterMap();
ClusterMapConfig clusterMapConfig = new ClusterMapConfig(verifiableProperties);
StoreConfig storeConfig = new StoreConfig(verifiableProperties);
// pick a node with no special partition as current node
Set<DataNodeId> specialPartitionNodes = clusterMap.getSpecialPartition().getReplicaIds().stream().map(ReplicaId::getDataNodeId).collect(Collectors.toSet());
DataNodeId currentNode = clusterMap.getDataNodes().stream().filter(d -> !specialPartitionNodes.contains(d)).findFirst().get();
MockStoreKeyConverterFactory storeKeyConverterFactory = new MockStoreKeyConverterFactory(null, null);
storeKeyConverterFactory.setConversionMap(new HashMap<>());
StorageManager storageManager = new StorageManager(storeConfig, new DiskManagerConfig(verifiableProperties), Utils.newScheduler(1, true), new MetricRegistry(), null, clusterMap, currentNode, null, null, new MockTime(), null, new InMemAccountService(false, false));
storageManager.start();
MockReplicationManager replicationManager = new MockReplicationManager(replicationConfig, clusterMapConfig, storeConfig, storageManager, clusterMap, currentNode, storeKeyConverterFactory, null);
ClusterMapChangeListener clusterMapChangeListener = clusterMap.getClusterMapChangeListener();
// find the special partition (not on current node) and get an irrelevant replica from it
PartitionId absentPartition = clusterMap.getSpecialPartition();
ReplicaId irrelevantReplica = absentPartition.getReplicaIds().get(0);
// find an existing replica on current node and one of its peer replicas on remote node
ReplicaId existingReplica = clusterMap.getReplicaIds(currentNode).get(0);
ReplicaId peerReplicaToRemove = existingReplica.getPartitionId().getReplicaIds().stream().filter(r -> r != existingReplica).findFirst().get();
// create a new node and place a peer of existing replica on it.
MockDataNodeId remoteNode = createDataNode(getListOfPorts(PLAIN_TEXT_PORT_START_NUMBER + 10, SSL_PORT_START_NUMBER + 10, HTTP2_PORT_START_NUMBER + 10), clusterMap.getDatacenterName((byte) 0), 3);
ReplicaId addedReplica = new MockReplicaId(remoteNode.getPort(), (MockPartitionId) existingReplica.getPartitionId(), remoteNode, 0);
// populate added replica and removed replica lists
List<ReplicaId> replicasToAdd = new ArrayList<>(Arrays.asList(existingReplica, addedReplica, irrelevantReplica));
List<ReplicaId> replicasToRemove = new ArrayList<>(Arrays.asList(existingReplica, peerReplicaToRemove, irrelevantReplica));
PartitionInfo partitionInfo = replicationManager.getPartitionToPartitionInfoMap().get(existingReplica.getPartitionId());
assertNotNull("PartitionInfo is not found", partitionInfo);
RemoteReplicaInfo peerReplicaInfo = partitionInfo.getRemoteReplicaInfos().stream().filter(info -> info.getReplicaId() == peerReplicaToRemove).findFirst().get();
// get the replica-thread for this peer replica
ReplicaThread peerReplicaThread = peerReplicaInfo.getReplicaThread();
// Test Case 1: replication manager encountered exception during startup (remote replica addition/removal will be skipped)
replicationManager.startWithException();
clusterMapChangeListener.onReplicaAddedOrRemoved(replicasToAdd, replicasToRemove);
// verify that PartitionInfo stays unchanged
verifyRemoteReplicaInfo(partitionInfo, addedReplica, false);
verifyRemoteReplicaInfo(partitionInfo, peerReplicaToRemove, true);
// Test Case 2: startup latch is interrupted
CountDownLatch initialLatch = replicationManager.startupLatch;
CountDownLatch mockLatch = Mockito.mock(CountDownLatch.class);
doThrow(new InterruptedException()).when(mockLatch).await();
replicationManager.startupLatch = mockLatch;
try {
clusterMapChangeListener.onReplicaAddedOrRemoved(replicasToAdd, replicasToRemove);
fail("should fail because startup latch is interrupted");
} catch (IllegalStateException e) {
// expected
}
replicationManager.startupLatch = initialLatch;
// Test Case 3: replication manager is successfully started
replicationManager.start();
clusterMapChangeListener.onReplicaAddedOrRemoved(replicasToAdd, replicasToRemove);
// verify that PartitionInfo has latest remote replica infos
verifyRemoteReplicaInfo(partitionInfo, addedReplica, true);
verifyRemoteReplicaInfo(partitionInfo, peerReplicaToRemove, false);
verifyRemoteReplicaInfo(partitionInfo, irrelevantReplica, false);
// verify new added replica is assigned to a certain thread
ReplicaThread replicaThread = replicationManager.getDataNodeIdToReplicaThreadMap().get(addedReplica.getDataNodeId());
assertNotNull("There is no ReplicaThread assocated with new replica", replicaThread);
Optional<RemoteReplicaInfo> findResult = replicaThread.getRemoteReplicaInfos().get(remoteNode).stream().filter(info -> info.getReplicaId() == addedReplica).findAny();
assertTrue("New added remote replica info should exist in corresponding thread", findResult.isPresent());
// verify the removed replica info's thread is null
assertNull("Thread in removed replica info should be null", peerReplicaInfo.getReplicaThread());
findResult = peerReplicaThread.getRemoteReplicaInfos().get(peerReplicaToRemove.getDataNodeId()).stream().filter(info -> info.getReplicaId() == peerReplicaToRemove).findAny();
assertFalse("Previous replica thread should not contain RemoteReplicaInfo that is already removed", findResult.isPresent());
storageManager.shutdown();
}
Aggregations