src/main/java/com/googlesource/gerrit/plugins/quota/RateLimitUploadListener.java

// Copyright (C) 2017 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.quota;

import static com.googlesource.gerrit.plugins.quota.Module.CACHE_NAME_ACCOUNTID;
import static com.googlesource.gerrit.plugins.quota.Module.CACHE_NAME_REMOTEHOST;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.cache.LoadingCache;
import com.google.common.util.concurrent.RateLimiter;
import com.google.gerrit.reviewdb.client.Account;
import com.google.gerrit.reviewdb.client.Project;
import com.google.gerrit.server.CurrentUser;
import com.google.gerrit.server.git.validators.UploadValidationListener;
import com.google.gerrit.server.validators.ValidationException;
import com.google.inject.Inject;
import com.google.inject.Provider;
import com.google.inject.name.Named;

import com.googlesource.gerrit.plugins.quota.Module.Holder;

import org.eclipse.jgit.lib.ObjectId;
import org.eclipse.jgit.lib.Repository;
import org.eclipse.jgit.transport.UploadPack;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.lang.reflect.Constructor;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.text.MessageFormat;
import java.util.Collection;
import java.util.concurrent.ExecutionException;

public class RateLimitUploadListener implements UploadValidationListener {
  private static final int SECONDS_PER_HOUR = 3600;
  private static final Logger log =
      LoggerFactory.getLogger(RateLimitUploadListener.class);
  private static final Method createStopwatchMethod;
  private static final Constructor<?> constructor;

  static {
    try {
      Class<?> sleepingStopwatchClass = Class.forName(
          "com.google.common.util.concurrent.RateLimiter$SleepingStopwatch");
      createStopwatchMethod =
          sleepingStopwatchClass.getDeclaredMethod("createFromSystemTimer");
      createStopwatchMethod.setAccessible(true);
      Class<?> burstyRateLimiterClass = Class.forName(
          "com.google.common.util.concurrent.SmoothRateLimiter$SmoothBursty");
      constructor = burstyRateLimiterClass.getDeclaredConstructors()[0];
      constructor.setAccessible(true);
    } catch (ClassNotFoundException | NoSuchMethodException e) {
      // shouldn't happen
      throw new RuntimeException(
          "Failed to prepare loading RateLimiter via reflection", e);
    }
  }

  /**
   * Create a custom instance of RateLimiter by accessing the non-public
   * constructor of the implementation class SmoothRateLimiter.SmoothBursty
   * through reflection.
   *
   * <p>
   * RateLimiter's implementation class SmoothRateLimiter.SmoothBursty allows to
   * collect permits during idle times which can be used to send bursts of
   * requests exceeding the average rate until the stored permits are consumed.
   * If the rate per second is 0.2 and you wait 20 seconds you can acquire 4
   * permits which in average matches the configured rate limit of 0.2
   * requests/second. If the permitted rate is smaller than 1 per second the
   * standard implementation doesn't allow any bursts since it hard-codes the
   * maximum time which can be used to collect stored permits to 1 second.
   *
   * <p>
   * Build jobs fetching updates from Gerrit are typically triggered by events
   * which can arrive in bursts. Hence the standard RateLimiter seems not to be
   * the right choice at least for fetch requests where we probably want to
   * limit the rate to less than 1 request per second per user.
   *
   * <p>
   * The used constructor can't be accessed through a public method yet hence
   * use reflection to instantiate it.
   *
   * @see "https://github.com/google/guava/issues/1974"
   * @param permitsPerSecond the new stable rate of this {@code RateLimiter}
   * @param maxBurstSeconds The maximum number of permits that can be saved.
   * @return a new RateLimiter
   */
  @VisibleForTesting
  static RateLimiter createSmoothBurstyRateLimiter(double permitsPerSecond,
      double maxBurstSeconds) {
    RateLimiter rl;
    try {
      Object stopwatch = createStopwatchMethod.invoke(null);
      rl = (RateLimiter) constructor.newInstance(stopwatch, maxBurstSeconds);
      rl.setRate(permitsPerSecond);
    } catch (InvocationTargetException | IllegalAccessException
        | InstantiationException e) {
      // shouldn't happen
      throw new RuntimeException(e);
    }
    return rl;
  }

  private final Provider<CurrentUser> user;
  private final LoadingCache<Account.Id, Holder> limitsPerAccount;
  private final LoadingCache<String, Holder> limitsPerRemoteHost;

  @Inject
  RateLimitUploadListener(Provider<CurrentUser> user,
      @Named(CACHE_NAME_ACCOUNTID) LoadingCache<Account.Id, Holder> limitsPerAccount,
      @Named(CACHE_NAME_REMOTEHOST) LoadingCache<String, Holder> limitsPerRemoteHost) {
    this.user = user;
    this.limitsPerAccount = limitsPerAccount;
    this.limitsPerRemoteHost = limitsPerRemoteHost;
  }

  @Override
  public void onBeginNegotiate(Repository repository, Project project,
      String remoteHost, UploadPack up, Collection<? extends ObjectId> wants,
      int cntOffered) throws ValidationException {
    RateLimiter limiter = null;
    CurrentUser u = user.get();
    if (u.isIdentifiedUser()) {
      Account.Id accountId = u.asIdentifiedUser().getAccountId();
      try {
        limiter = limitsPerAccount.get(accountId).get();
      } catch (ExecutionException e) {
        String msg = MessageFormat
            .format("Cannot get rate limits for account ''{}''", accountId);
        log.warn(msg, e);
      }
    } else {
      try {
        limiter = limitsPerRemoteHost.get(remoteHost).get();
      } catch (ExecutionException e) {
        String msg = MessageFormat.format(
            "Cannot get rate limits for anonymous access from remote host ''{0}''",
            remoteHost);
        log.warn(msg, e);
      }
    }
    if (limiter != null && !limiter.tryAcquire()) {
      throw new RateLimitException(
          MessageFormat.format("Exceeded rate limit of {0,number,##.##} fetch requests/hour",
              limiter.getRate() * SECONDS_PER_HOUR));
    }
  }

  @Override
  public void onPreUpload(Repository repository, Project project,
      String remoteHost, UploadPack up, Collection<? extends ObjectId> wants,
      Collection<? extends ObjectId> haves) throws ValidationException {
  }
}