Line data Source code
1 : // Copyright (C) 2018 The Android Open Source Project
2 : //
3 : // Licensed under the Apache License, Version 2.0 (the "License");
4 : // you may not use this file except in compliance with the License.
5 : // You may obtain a copy of the License at
6 : //
7 : // http://www.apache.org/licenses/LICENSE-2.0
8 : //
9 : // Unless required by applicable law or agreed to in writing, software
10 : // distributed under the License is distributed on an "AS IS" BASIS,
11 : // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 : // See the License for the specific language governing permissions and
13 : // limitations under the License.
14 :
15 : package com.google.gerrit.server.quota;
16 :
17 : import static com.google.common.base.Preconditions.checkState;
18 :
19 : import com.google.common.collect.ImmutableList;
20 : import com.google.common.flogger.FluentLogger;
21 : import com.google.gerrit.entities.Account;
22 : import com.google.gerrit.entities.Change;
23 : import com.google.gerrit.entities.Project;
24 : import com.google.gerrit.server.CurrentUser;
25 : import com.google.gerrit.server.plugincontext.PluginSetContext;
26 : import com.google.gerrit.server.plugincontext.PluginSetEntryContext;
27 : import com.google.gerrit.server.quota.QuotaResponse.Aggregated;
28 : import com.google.inject.Inject;
29 : import com.google.inject.Provider;
30 : import com.google.inject.Singleton;
31 : import java.util.ArrayList;
32 : import java.util.List;
33 :
34 : @Singleton
35 : public class DefaultQuotaBackend implements QuotaBackend {
36 138 : private static final FluentLogger logger = FluentLogger.forEnclosingClass();
37 :
38 : private final Provider<CurrentUser> userProvider;
39 : private final PluginSetContext<QuotaEnforcer> quotaEnforcers;
40 :
41 : @Inject
42 : DefaultQuotaBackend(
43 138 : Provider<CurrentUser> userProvider, PluginSetContext<QuotaEnforcer> quotaEnforcers) {
44 138 : this.userProvider = userProvider;
45 138 : this.quotaEnforcers = quotaEnforcers;
46 138 : }
47 :
48 : @Override
49 : public WithUser currentUser() {
50 27 : return new WithUser(quotaEnforcers, userProvider.get());
51 : }
52 :
53 : @Override
54 : public WithUser user(CurrentUser user) {
55 97 : return new WithUser(quotaEnforcers, user);
56 : }
57 :
58 : private static QuotaResponse.Aggregated request(
59 : PluginSetContext<QuotaEnforcer> quotaEnforcers,
60 : String quotaGroup,
61 : QuotaRequestContext requestContext,
62 : long numTokens,
63 : boolean deduct) {
64 32 : checkState(numTokens > 0, "numTokens must be a positive, non-zero long");
65 :
66 : // PluginSets can change their content when plugins (de-)register. Copy the currently registered
67 : // plugins so that we can iterate twice on a stable list.
68 32 : List<PluginSetEntryContext<QuotaEnforcer>> enforcers = ImmutableList.copyOf(quotaEnforcers);
69 32 : List<QuotaResponse> responses = new ArrayList<>(enforcers.size());
70 32 : for (PluginSetEntryContext<QuotaEnforcer> enforcer : enforcers) {
71 : try {
72 1 : if (deduct) {
73 1 : responses.add(enforcer.call(p -> p.requestTokens(quotaGroup, requestContext, numTokens)));
74 : } else {
75 1 : responses.add(enforcer.call(p -> p.dryRun(quotaGroup, requestContext, numTokens)));
76 : }
77 1 : } catch (RuntimeException e) {
78 : // Roll back the quota request for all enforcers that deducted the quota. Rethrow the
79 : // exception to adhere to the API contract.
80 1 : if (deduct) {
81 1 : refillAfterErrorOrException(enforcers, responses, quotaGroup, requestContext, numTokens);
82 : }
83 1 : throw e;
84 1 : }
85 1 : }
86 :
87 32 : if (deduct && responses.stream().anyMatch(r -> r.status().isError())) {
88 : // Roll back the quota request for all enforcers that deducted the quota (= the request
89 : // succeeded). Don't touch failed enforcers as the interface contract said that failed
90 : // requests should not be deducted.
91 1 : refillAfterErrorOrException(enforcers, responses, quotaGroup, requestContext, numTokens);
92 : }
93 :
94 32 : logger.atFine().log(
95 : "Quota request for %s with %s (deduction=%s) for %s token returned %s",
96 : quotaGroup,
97 : requestContext,
98 32 : deduct ? "(deduction=yes)" : "(deduction=no)",
99 32 : numTokens,
100 : responses);
101 32 : return QuotaResponse.Aggregated.create(ImmutableList.copyOf(responses));
102 : }
103 :
104 : private static QuotaResponse.Aggregated availableTokens(
105 : PluginSetContext<QuotaEnforcer> quotaEnforcers,
106 : String quotaGroup,
107 : QuotaRequestContext requestContext) {
108 : // PluginSets can change their content when plugins (de-)register. Copy the currently registered
109 : // plugins so that we can iterate twice on a stable list.
110 97 : List<PluginSetEntryContext<QuotaEnforcer>> enforcers = ImmutableList.copyOf(quotaEnforcers);
111 97 : List<QuotaResponse> responses = new ArrayList<>(enforcers.size());
112 97 : for (PluginSetEntryContext<QuotaEnforcer> enforcer : enforcers) {
113 1 : responses.add(enforcer.call(p -> p.availableTokens(quotaGroup, requestContext)));
114 1 : }
115 97 : return QuotaResponse.Aggregated.create(responses);
116 : }
117 :
118 : private static void refillAfterErrorOrException(
119 : List<PluginSetEntryContext<QuotaEnforcer>> enforcers,
120 : List<QuotaResponse> collectedResponses,
121 : String quotaGroup,
122 : QuotaRequestContext requestContext,
123 : long numTokens) {
124 1 : for (int i = 0; i < collectedResponses.size(); i++) {
125 1 : if (collectedResponses.get(i).status().isOk()) {
126 1 : enforcers.get(i).run(p -> p.refill(quotaGroup, requestContext, numTokens));
127 : }
128 : }
129 1 : }
130 :
131 : static class WithUser extends WithResource implements QuotaBackend.WithUser {
132 : WithUser(PluginSetContext<QuotaEnforcer> quotaEnforcers, CurrentUser user) {
133 104 : super(quotaEnforcers, QuotaRequestContext.builder().user(user).build());
134 104 : }
135 :
136 : @Override
137 : public QuotaBackend.WithResource account(Account.Id account) {
138 6 : QuotaRequestContext ctx = requestContext.toBuilder().account(account).build();
139 6 : return new WithResource(quotaEnforcers, ctx);
140 : }
141 :
142 : @Override
143 : public QuotaBackend.WithResource project(Project.NameKey project) {
144 102 : QuotaRequestContext ctx = requestContext.toBuilder().project(project).build();
145 102 : return new WithResource(quotaEnforcers, ctx);
146 : }
147 :
148 : @Override
149 : public QuotaBackend.WithResource change(Change.Id change, Project.NameKey project) {
150 13 : QuotaRequestContext ctx = requestContext.toBuilder().change(change).project(project).build();
151 13 : return new WithResource(quotaEnforcers, ctx);
152 : }
153 : }
154 :
155 : static class WithResource implements QuotaBackend.WithResource {
156 : protected final QuotaRequestContext requestContext;
157 : protected final PluginSetContext<QuotaEnforcer> quotaEnforcers;
158 :
159 : private WithResource(
160 104 : PluginSetContext<QuotaEnforcer> quotaEnforcers, QuotaRequestContext quotaRequestContext) {
161 104 : this.quotaEnforcers = quotaEnforcers;
162 104 : this.requestContext = quotaRequestContext;
163 104 : }
164 :
165 : @Override
166 : public QuotaResponse.Aggregated requestTokens(String quotaGroup, long numTokens) {
167 32 : return DefaultQuotaBackend.request(
168 : quotaEnforcers, quotaGroup, requestContext, numTokens, true);
169 : }
170 :
171 : @Override
172 : public QuotaResponse.Aggregated dryRun(String quotaGroup, long numTokens) {
173 1 : return DefaultQuotaBackend.request(
174 : quotaEnforcers, quotaGroup, requestContext, numTokens, false);
175 : }
176 :
177 : @Override
178 : public Aggregated availableTokens(String quotaGroup) {
179 97 : return DefaultQuotaBackend.availableTokens(quotaEnforcers, quotaGroup, requestContext);
180 : }
181 : }
182 : }
|
