1 /* 2 * Copyright (C) 2024 The Android Open Source Project 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 17 package com.android.adservices.service.measurement; 18 19 import com.google.common.collect.ImmutableSet; 20 21 import java.util.HashMap; 22 import java.util.Map; 23 import java.util.Objects; 24 import java.util.Optional; 25 26 /** Class for keeping track of aggregatable named budgets */ 27 public class AggregatableNamedBudgets { 28 private Map<String, BudgetAndContribution> mNameToBudgetAndAggregateContribution = 29 new HashMap<>(); 30 31 public static class BudgetAndContribution { 32 private int mAggregateContribution; 33 private final int mBudget; 34 BudgetAndContribution(int budget)35 public BudgetAndContribution(int budget) { 36 mAggregateContribution = 0; 37 mBudget = budget; 38 } 39 getBudget()40 public int getBudget() { 41 return mBudget; 42 } 43 getAggregateContribution()44 public int getAggregateContribution() { 45 return mAggregateContribution; 46 } 47 setAggregateContribution(int contribution)48 public void setAggregateContribution(int contribution) { 49 mAggregateContribution = contribution; 50 } 51 52 @Override hashCode()53 public int hashCode() { 54 return Objects.hash(mAggregateContribution, mBudget); 55 } 56 57 @Override equals(Object obj)58 public boolean equals(Object obj) { 59 if (!(obj instanceof BudgetAndContribution)) { 60 return false; 61 } 62 BudgetAndContribution budgetAndContribution = (BudgetAndContribution) obj; 63 return budgetAndContribution.mAggregateContribution == mAggregateContribution 64 && budgetAndContribution.mBudget == mBudget; 65 } 66 } 67 68 /** 69 * @return the budget 70 */ maybeGetBudget(String name)71 public Optional<Integer> maybeGetBudget(String name) { 72 if (mNameToBudgetAndAggregateContribution.containsKey(name)) { 73 return Optional.of(mNameToBudgetAndAggregateContribution.get(name).getBudget()); 74 } 75 return Optional.empty(); 76 } 77 78 /** 79 * @return the aggregate contribution towards the named budget 80 */ maybeGetContribution(String name)81 public Optional<Integer> maybeGetContribution(String name) { 82 if (mNameToBudgetAndAggregateContribution.containsKey(name)) { 83 return Optional.of( 84 mNameToBudgetAndAggregateContribution.get(name).getAggregateContribution()); 85 } 86 return Optional.empty(); 87 } 88 89 /** 90 * @return whether the contribution was successfully added to the named budget or not 91 */ setContribution(String name, int contribution)92 public boolean setContribution(String name, int contribution) { 93 if (contribution > mNameToBudgetAndAggregateContribution.get(name).getBudget()) { 94 return false; 95 } 96 97 // Set mAggregateContribution to contribution for the named budget. 98 mNameToBudgetAndAggregateContribution.get(name).setAggregateContribution(contribution); 99 return true; 100 } 101 102 /** 103 * @return the names of the budgets 104 */ getBudgetNames()105 public ImmutableSet<String> getBudgetNames() { 106 return ImmutableSet.copyOf(mNameToBudgetAndAggregateContribution.keySet()); 107 } 108 109 /** Set budget for the named budget's aggregate contributions */ createContributionBudget(String name, int maxBudget)110 public void createContributionBudget(String name, int maxBudget) { 111 BudgetAndContribution budgetAndContribution = new BudgetAndContribution(maxBudget); 112 113 mNameToBudgetAndAggregateContribution.put(name, budgetAndContribution); 114 } 115 116 @Override hashCode()117 public int hashCode() { 118 return Objects.hash(mNameToBudgetAndAggregateContribution); 119 } 120 121 @Override equals(Object obj)122 public boolean equals(Object obj) { 123 if (!(obj instanceof AggregatableNamedBudgets)) { 124 return false; 125 } 126 AggregatableNamedBudgets aggregatableNamedBudgets = (AggregatableNamedBudgets) obj; 127 return Objects.equals( 128 mNameToBudgetAndAggregateContribution, 129 aggregatableNamedBudgets.mNameToBudgetAndAggregateContribution); 130 } 131 } 132