| // Copyright (C) 2022 The Android Open Source Project |
| // |
| // Licensed under the Apache License, Version 2.0 (the "License"); |
| // you may not use this file except in compliance with the License. |
| // You may obtain a copy of the License at |
| // |
| // http://www.apache.org/licenses/LICENSE-2.0 |
| // |
| // Unless required by applicable law or agreed to in writing, software |
| // distributed under the License is distributed on an "AS IS" BASIS, |
| // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| // See the License for the specific language governing permissions and |
| // limitations under the License. |
| package com.googlesource.gerrit.plugins.ratelimiter; |
| |
| import static com.googlesource.gerrit.plugins.ratelimiter.ListCommand.FORMAT; |
| import static com.googlesource.gerrit.plugins.ratelimiter.Module.UPLOAD_PACK_PER_HOUR; |
| |
| import com.google.common.cache.LoadingCache; |
| import com.google.gerrit.entities.Account; |
| import com.google.gerrit.extensions.restapi.ResourceNotFoundException; |
| import com.google.gerrit.server.account.AccountResolver; |
| import com.google.gson.*; |
| import com.google.inject.Inject; |
| import com.google.inject.name.Named; |
| import java.io.IOException; |
| import java.time.Duration; |
| import java.util.*; |
| import java.util.concurrent.TimeUnit; |
| import org.eclipse.jgit.errors.ConfigInvalidException; |
| |
| public class RateLimiterProcessing { |
| |
| private final LoadingCache<String, RateLimiter> uploadPackPerHour; |
| private final UserResolver userResolver; |
| private final AccountResolver accountResolver; |
| |
| @Inject |
| public RateLimiterProcessing( |
| @Named(UPLOAD_PACK_PER_HOUR) LoadingCache<String, RateLimiter> uploadPackPerHour, |
| UserResolver userResolver, |
| AccountResolver accountResolver) { |
| this.uploadPackPerHour = uploadPackPerHour; |
| this.userResolver = userResolver; |
| this.accountResolver = accountResolver; |
| } |
| |
| public String listPermits() { |
| return uploadPackPerHour.asMap().entrySet().stream() |
| .sorted(Map.Entry.comparingByValue()) |
| .map( |
| entry -> |
| String.format( |
| FORMAT, |
| getDisplayValue(entry.getKey(), userResolver), |
| permits(entry.getValue().permitsPerHour()), |
| permits(entry.getValue().availablePermits()), |
| permits(entry.getValue().usedPermits()), |
| Duration.ofSeconds(entry.getValue().remainingTime(TimeUnit.SECONDS)))) |
| .reduce("", String::concat); |
| } |
| |
| public String listPermitsAsJson() { |
| List<String> permitList = new ArrayList<>(); |
| Gson gson = new GsonBuilder().setPrettyPrinting().create(); |
| uploadPackPerHour.asMap().entrySet().stream() |
| .sorted(Map.Entry.comparingByValue()) |
| .map(this::getJsonObjectString) |
| .forEach(permitList::add); |
| return gson.toJson(new JsonParser().parse(permitList.toString())); |
| } |
| |
| private String getJsonObjectString(Map.Entry<String, RateLimiter> entry) { |
| JsonObject jsonObject = new JsonObject(); |
| jsonObject.addProperty("AccountId", getDisplayValue(entry.getKey(), userResolver)); |
| jsonObject.addProperty("permits_per_hour", permits(entry.getValue().permitsPerHour())); |
| jsonObject.addProperty("available_permits", permits(entry.getValue().availablePermits())); |
| jsonObject.addProperty("used_permit", permits(entry.getValue().usedPermits())); |
| jsonObject.addProperty( |
| "replenish_in", |
| Duration.ofSeconds(entry.getValue().remainingTime(TimeUnit.SECONDS)).toString()); |
| return jsonObject.toString(); |
| } |
| |
| private String permits(int value) { |
| return value == Integer.MAX_VALUE ? "unlimited" : Integer.toString(value); |
| } |
| |
| private String getDisplayValue(String key, UserResolver userResolver) { |
| return userResolver.getUserName(key).map(name -> key + " (" + name + ")").orElse(key); |
| } |
| |
| public void replenish(boolean all, List<Account.Id> accountIds, List<String> remoteHosts) { |
| if (all && (!accountIds.isEmpty() || !remoteHosts.isEmpty())) { |
| throw new IllegalArgumentException("cannot use --all with --user or --remotehost"); |
| } |
| if (all) { |
| uploadPackPerHour.asMap().values().forEach(RateLimiter::replenishPermits); |
| return; |
| } |
| accountIds.forEach(account -> replenishIfPresent(Integer.toString(account.get()))); |
| remoteHosts.forEach(this::replenishIfPresent); |
| } |
| |
| List<Account.Id> convertToAccountId(String[] usernames) |
| throws ConfigInvalidException, IOException, ResourceNotFoundException { |
| List<Account.Id> accountIds = new ArrayList<>(); |
| for (String user : usernames) { |
| AccountResolver.Result accountId = accountResolver.resolve(user); |
| if (accountId.asIdSet().isEmpty()) |
| throw new ResourceNotFoundException(String.format("User %s not found", user)); |
| accountIds.addAll(accountId.asIdSet()); |
| } |
| return accountIds; |
| } |
| |
| private void replenishIfPresent(String key) { |
| RateLimiter limiter = uploadPackPerHour.getIfPresent(key); |
| if (limiter != null) { |
| limiter.replenishPermits(); |
| } |
| } |
| } |