Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[risk=low][RW-13494] Renaming of free tier in many locations #9092

Draft
wants to merge 2 commits into
base: main
Choose a base branch
from
Draft
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,6 @@ void fireAdministrativeBypassTime(

void fireAcknowledgeTermsOfService(DbUser targetUser, Integer termsOfServiceVersion);

void fireSetFreeTierDollarLimitOverride(
void fireSetInitialCreditsOverride(
Long targetUserId, @Nullable Double previousDollarQuota, @Nullable Double newDollarQuota);
}
Original file line number Diff line number Diff line change
Expand Up @@ -117,7 +117,7 @@ public void fireAcknowledgeTermsOfService(DbUser targetUser, Integer termsOfServ
}

@Override
public void fireSetFreeTierDollarLimitOverride(
public void fireSetInitialCreditsOverride(
Long targetUserId, @Nullable Double previousDollarQuota, @Nullable Double newDollarQuota) {
DbUser adminUser = dbUserProvider.get();
Builder builder =
Expand All @@ -129,7 +129,7 @@ public void fireSetFreeTierDollarLimitOverride(
.actionId(actionIdProvider.get())
.actionType(ActionType.EDIT)
.targetType(TargetType.ACCOUNT)
.targetPropertyMaybe(AccountTargetProperty.FREE_TIER_DOLLAR_QUOTA.getPropertyName())
.targetPropertyMaybe(AccountTargetProperty.INITIAL_CREDITS_OVERRIDE.getPropertyName())
.targetIdMaybe(targetUserId);

if (previousDollarQuota != null) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,7 @@
public enum AccountTargetProperty implements SimpleTargetProperty {
IS_ENABLED("is_enabled"),
ACKNOWLEDGED_TOS_VERSION("acknowledged_tos_version"),
FREE_TIER_DOLLAR_QUOTA("free_tier_dollar_quota"),
INITIAL_CREDITS_OVERRIDE("initial_credits_override"),
ACCESS_TIERS("access_tiers");

private final String propertyName;
Expand Down
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
package org.pmiops.workbench.api;

import static org.pmiops.workbench.utils.BillingUtils.isInitialCredits;
import static org.pmiops.workbench.utils.CostComparisonUtils.getUserFreeTierDollarLimit;
import static org.pmiops.workbench.utils.CostComparisonUtils.getUserInitialCreditsLimit;

import com.google.common.collect.Sets;
import jakarta.inject.Provider;
Expand All @@ -24,7 +24,7 @@
import org.pmiops.workbench.exceptions.WorkbenchException;
import org.pmiops.workbench.leonardo.LeonardoApiClient;
import org.pmiops.workbench.mail.MailService;
import org.pmiops.workbench.model.ExpiredInitialCreditsEventRequest;
import org.pmiops.workbench.model.ExhaustedInitialCreditsEventRequest;
import org.pmiops.workbench.utils.CostComparisonUtils;
import org.pmiops.workbench.workspaces.WorkspaceService;
import org.slf4j.Logger;
Expand All @@ -39,32 +39,32 @@ public class CloudTaskInitialCreditsExhaustionController
private static final Logger logger =
LoggerFactory.getLogger(CloudTaskInitialCreditsExhaustionController.class);

private final WorkspaceDao workspaceDao;
private final WorkspaceService workspaceService;
private final UserDao userDao;
private final Provider<WorkbenchConfig> workbenchConfig;
private final LeonardoApiClient leonardoApiClient;
private final MailService mailService;
private final Provider<WorkbenchConfig> workbenchConfig;
private final UserDao userDao;
private final WorkspaceDao workspaceDao;
private final WorkspaceService workspaceService;

CloudTaskInitialCreditsExhaustionController(
WorkspaceDao workspaceDao,
WorkspaceService workspaceService,
UserDao userDao,
Provider<WorkbenchConfig> workbenchConfig,
LeonardoApiClient leonardoApiClient,
MailService mailService) {
this.workspaceDao = workspaceDao;
this.workspaceService = workspaceService;
this.userDao = userDao;
this.workbenchConfig = workbenchConfig;
MailService mailService,
Provider<WorkbenchConfig> workbenchConfig,
UserDao userDao,
WorkspaceDao workspaceDao,
WorkspaceService workspaceService) {
this.leonardoApiClient = leonardoApiClient;
this.mailService = mailService;
this.userDao = userDao;
this.workbenchConfig = workbenchConfig;
this.workspaceDao = workspaceDao;
this.workspaceService = workspaceService;
}

@SuppressWarnings("unchecked")
@Override
public ResponseEntity<Void> handleInitialCreditsExhaustionBatch(
ExpiredInitialCreditsEventRequest request) {
ExhaustedInitialCreditsEventRequest request) {

if (request.getUsers().isEmpty()) {
logger.warn("users are empty");
Expand All @@ -85,11 +85,12 @@ public ResponseEntity<Void> handleInitialCreditsExhaustionBatch(
Map<String, Double> stringKeyLiveCostMap = (Map<String, Double>) request.getLiveCostByCreator();
Map<Long, Double> liveCostByCreator = convertMapKeysToLong(stringKeyLiveCostMap);

var newlyExpiredUsers = getNewlyExpiredUsers(usersSet, dbCostByCreator, liveCostByCreator);
var newlyExhaustedUsers = getNewlyExhaustedUsers(usersSet, dbCostByCreator, liveCostByCreator);

handleExpiredUsers(newlyExpiredUsers);
handleExhaustedUsers(newlyExhaustedUsers);

alertUsersBasedOnTheThreshold(usersSet, dbCostByCreator, liveCostByCreator, newlyExpiredUsers);
alertUsersBasedOnTheThreshold(
usersSet, dbCostByCreator, liveCostByCreator, newlyExhaustedUsers);

logger.info(
"handleInitialCreditsExhaustionBatch: Finished processing request for users: {}",
Expand All @@ -98,19 +99,20 @@ public ResponseEntity<Void> handleInitialCreditsExhaustionBatch(
return ResponseEntity.noContent().build();
}

private void handleExpiredUsers(Set<DbUser> newlyExpiredUsers) {
newlyExpiredUsers.forEach(
private void handleExhaustedUsers(Set<DbUser> newlyExhaustedUsers) {
newlyExhaustedUsers.forEach(
user -> {
logger.info(
"Free tier Billing Service: handling user with expired credits {}",
"handleInitialCreditsExhaustionBatch: handling user with exhausted credits {}",
user.getUsername());
workspaceService.updateInitialCreditsExhaustion(user, true);
// delete apps and runtimes
deleteAppsAndRuntimesInFreeTierWorkspaces(user);
deleteAppsAndRuntimesInInitialCreditsWorkspaces(user);
try {
mailService.alertUserInitialCreditsExhausted(user);
} catch (MessagingException e) {
logger.warn("failed to send free tier expiration email to {}", user.getUsername(), e);
logger.warn(
"failed to send initial credits exhaustion email to {}", user.getUsername(), e);
}
});
}
Expand All @@ -119,7 +121,7 @@ private void alertUsersBasedOnTheThreshold(
Set<DbUser> users,
Map<Long, Double> dbCostByCreator,
Map<Long, Double> liveCostByCreator,
Set<DbUser> newlyExpiredUsers) {
Set<DbUser> newlyExhaustedUsers) {
final List<Double> costThresholdsInDescOrder =
workbenchConfig.get().billing.freeTierCostAlertThresholds;
costThresholdsInDescOrder.sort(Comparator.reverseOrder());
Expand All @@ -129,10 +131,10 @@ private void alertUsersBasedOnTheThreshold(
.filter(u -> liveCostByCreator.containsKey(u.getUserId()))
.collect(Collectors.toMap(DbUser::getUserId, Function.identity()));

// Filter out the users who have recently expired because we already alerted them
// Filter out the users who have recently exhausted because we already alerted them
Map<Long, Double> filteredLiveCostByCreator =
liveCostByCreator.entrySet().stream()
.filter(entry -> !newlyExpiredUsers.contains(usersCache.get(entry.getKey())))
.filter(entry -> !newlyExhaustedUsers.contains(usersCache.get(entry.getKey())))
.collect(Collectors.toMap(Map.Entry::getKey, Map.Entry::getValue));

logger.info("Handling cost alerts for users: {}", usersCache.keySet());
Expand All @@ -153,29 +155,30 @@ private void alertUsersBasedOnTheThreshold(
}

/**
* Get the list of newly expired users (who exceeded their free tier limit) and mark all their
* workspaces as inactive
* Get the list of newly exhausted users (those who exceeded their initial credits limit) and mark
* all their workspaces as inactive
*
* @param allUsers set of all users to filter them whether they have active free tier workspace
* @param allUsers set of all users to filter whether they have active initial credits workspaces
* @param dbCostByCreator Map of userId->dbCost
* @param liveCostByCreator Map of userId->liveCost
* @return a {@link Set} of newly expired users
* @return a {@link Set} of newly exhausted users
*/
private Set<DbUser> getNewlyExpiredUsers(
private Set<DbUser> getNewlyExhaustedUsers(
final Set<DbUser> allUsers,
Map<Long, Double> dbCostByCreator,
Map<Long, Double> liveCostByCreator) {

final Map<Long, DbUser> dbUsersWithChangedCosts =
findDbUsersWithChangedCosts(allUsers, dbCostByCreator, liveCostByCreator);
Set<DbUser> freeTierUsers = getFreeTierActiveWorkspaceCreatorsIn(allUsers);
Set<DbUser> creatorsWithInitialCredits =
filterToWorkspaceCreatorsWithActiveInitialCredits(allUsers);

// Find users who exceeded their free tier limit
// Find users who exceeded their initial credits limit
// Here costs in liveCostByCreator could be outdated because we're filtering on active or
// recently deleted workspaces in previous steps.
// However, dbCostByCreator will contain the up-to-date costs for all the
// other workspaces. This is why Math.max is used
final Set<DbUser> expiredUsers =
final Set<DbUser> exhaustedUsers =
dbUsersWithChangedCosts.entrySet().stream()
.filter(
e ->
Expand All @@ -187,14 +190,15 @@ private Set<DbUser> getNewlyExpiredUsers(
.map(Map.Entry::getValue)
.collect(Collectors.toSet());

final Set<DbUser> newlyExpiredFreeTierUsers = Sets.intersection(expiredUsers, freeTierUsers);
final Set<DbUser> newlyExhaustedCreatorsWithInitialCredits =
Sets.intersection(exhaustedUsers, creatorsWithInitialCredits);

logger.info(
String.format(
"Found %d users exceeding their free tier limit, out of which, %d are new",
expiredUsers.size(), newlyExpiredFreeTierUsers.size()));
"Found %d users exceeding their initial credits limit, out of which, %d are new",
exhaustedUsers.size(), newlyExhaustedCreatorsWithInitialCredits.size()));

return newlyExpiredFreeTierUsers;
return newlyExhaustedCreatorsWithInitialCredits;
}

/**
Expand Down Expand Up @@ -229,12 +233,12 @@ private Map<Long, DbUser> findDbUsersWithChangedCosts(
return dbUsersWithChangedCosts;
}

private Set<DbUser> getFreeTierActiveWorkspaceCreatorsIn(Set<DbUser> users) {
private Set<DbUser> filterToWorkspaceCreatorsWithActiveInitialCredits(Set<DbUser> users) {
return workspaceDao.findCreatorsByActiveInitialCredits(
List.of(workbenchConfig.get().billing.initialCreditsBillingAccountName()), users);
}

private void deleteAppsAndRuntimesInFreeTierWorkspaces(DbUser user) {
private void deleteAppsAndRuntimesInInitialCreditsWorkspaces(DbUser user) {
logger.info("Deleting apps and runtimes for user {}", user.getUsername());

workspaceDao.findAllByCreator(user).stream()
Expand All @@ -258,7 +262,7 @@ private void deleteAppsAndRuntimesInFreeTierWorkspaces(DbUser user) {
* Has this user passed a cost threshold between this check and the previous run?
*
* <p>Compare this user's total cost with that of the previous run, and trigger an alert if this
* is the run which pushed it over a free credits threshold.
* is the run which pushed it over an initial credits threshold.
*
* @param user The user to check
* @param currentCost The current total cost incurred by this user, according to BigQuery
Expand All @@ -270,7 +274,7 @@ private void maybeAlertOnCostThresholds(
DbUser user, double currentCost, double previousCost, List<Double> thresholdsInDescOrder) {

final double limit =
getUserFreeTierDollarLimit(
getUserInitialCreditsLimit(
user, workbenchConfig.get().billing.defaultFreeCreditsDollarLimit);
final double remainingBalance = limit - currentCost;

Expand All @@ -279,7 +283,7 @@ private void maybeAlertOnCostThresholds(
if (CostComparisonUtils.compareCosts(currentCost, previousCost) < 0) {
String msg =
String.format(
"User %s (%s) has %f in total free tier spending in BigQuery, "
"User %s (%s) has %f in total initial credits spending in BigQuery, "
+ "which is less than the %f previous spending we have recorded in the DB",
user.getUsername(),
Optional.ofNullable(user.getContactEmail()).orElse("NULL"),
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -57,21 +57,21 @@ public class CloudTaskUserController implements CloudTaskUserApiDelegate {

private final AccessModuleService accessModuleService;
private final CloudResourceManagerService cloudResourceManagerService;
private final InitialCreditsBatchUpdateService freeTierBillingUpdateService;
private final InitialCreditsBatchUpdateService initialCreditsBatchUpdateService;
private final InitialCreditsService initialCreditsService;
private final Provider<Stopwatch> stopwatchProvider;
private final UserService userService;

CloudTaskUserController(
AccessModuleService accessModuleService,
CloudResourceManagerService cloudResourceManagerService,
InitialCreditsBatchUpdateService freeTierBillingUpdateService,
InitialCreditsBatchUpdateService initialCreditsBatchUpdateService,
InitialCreditsService initialCreditsService,
Provider<Stopwatch> stopwatchProvider,
UserService userService) {
this.accessModuleService = accessModuleService;
this.cloudResourceManagerService = cloudResourceManagerService;
this.freeTierBillingUpdateService = freeTierBillingUpdateService;
this.initialCreditsBatchUpdateService = initialCreditsBatchUpdateService;
this.initialCreditsService = initialCreditsService;
this.stopwatchProvider = stopwatchProvider;
this.userService = userService;
Expand Down Expand Up @@ -131,8 +131,7 @@ private int auditOneUser(DbUser user) {
* Takes in batch of user Ids check whether users have incurred sufficient cost in their
* workspaces to trigger alerts due to passing thresholds or exceeding limits
*
* @param userIds : Batch of user IDs from cloud task queue: freeTierBillingQueue
* @return
* @param userIds : Batch of user IDs from cloud task queue: initialCreditsUsageQueue
*/
@Override
public ResponseEntity<Void> checkAndAlertFreeTierBillingUsageBatch(List<Long> userIds) {
Expand All @@ -143,7 +142,7 @@ public ResponseEntity<Void> checkAndAlertFreeTierBillingUsageBatch(List<Long> us
return processUserIdBatch(
userIds,
"alerting for initial credits usage",
freeTierBillingUpdateService::checkAndAlertFreeTierBillingUsage);
initialCreditsBatchUpdateService::checkAndAlertInitialCreditsUsage);
}

@Override
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,6 @@

import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;
import org.pmiops.workbench.cloudtasks.TaskQueueService;
import org.pmiops.workbench.db.dao.GoogleProjectPerCostDao;
import org.pmiops.workbench.db.dao.UserService;
Expand All @@ -15,34 +14,31 @@
@RestController
public class OfflineBillingController implements OfflineBillingApiDelegate {

private final InitialCreditsBatchUpdateService freeTierBillingService;
private final GoogleProjectPerCostDao googleProjectPerCostDao;
private final InitialCreditsBatchUpdateService initialCreditsBatchUpdateService;
private final TaskQueueService taskQueueService;

private final UserService userService;

@Autowired
OfflineBillingController(
InitialCreditsBatchUpdateService freeTierBillingService,
GoogleProjectPerCostDao googleProjectPerCostDao,
UserService userService,
TaskQueueService taskQueueService) {
this.freeTierBillingService = freeTierBillingService;
InitialCreditsBatchUpdateService initialCreditsBatchUpdateService,
TaskQueueService taskQueueService,
UserService userService) {
this.googleProjectPerCostDao = googleProjectPerCostDao;
this.initialCreditsBatchUpdateService = initialCreditsBatchUpdateService;
this.taskQueueService = taskQueueService;
this.userService = userService;
this.googleProjectPerCostDao = googleProjectPerCostDao;
}

@Override
public ResponseEntity<Void> checkFreeTierBillingUsage() {
// Get cost for all workspace from BQ
Map<String, Double> freeTierForAllWorkspace =
freeTierBillingService.getFreeTierWorkspaceCostsFromBQ();
Map<String, Double> workspaceCostsFromBQ =
initialCreditsBatchUpdateService.getWorkspaceCostsFromBQ();

List<DbGoogleProjectPerCost> googleProjectCostList =
freeTierForAllWorkspace.entrySet().stream()
.map(DbGoogleProjectPerCost::new)
.collect(Collectors.toList());
workspaceCostsFromBQ.entrySet().stream().map(DbGoogleProjectPerCost::new).toList();

// Clear table googleproject_cost and then insert all entries from BQ
googleProjectPerCostDao.deleteAll();
Expand All @@ -51,8 +47,8 @@ public ResponseEntity<Void> checkFreeTierBillingUsage() {

List<Long> allUserIds = userService.getAllUserIds();

taskQueueService.groupAndPushFreeTierBilling(allUserIds);
log.info("Pushed all users to Cloud Task for Free Tier Billing");
taskQueueService.groupAndPushInitialCreditsUsage(allUserIds);
log.info("Pushed all users to the Initial Credits Usage Cloud Task");

return ResponseEntity.noContent().build();
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -341,7 +341,7 @@ private boolean notifyForUnusedDisk(ListPersistentDiskResponse disk, int daysUnu
if (BillingUtils.isInitialCredits(
workspace.get().getBillingAccountName(), configProvider.get())) {
initialCreditsRemaining =
initialCreditsService.getWorkspaceCreatorFreeCreditsRemaining(workspace.get());
initialCreditsService.getWorkspaceCreatorInitialCreditsRemaining(workspace.get());
}

mailService.alertUsersUnusedDiskWarningThreshold(
Expand Down
Loading