Extend QuotaBackend and QuotaEnforcer extension point

The current design is complete from the perspective when the number
of tokens is known prior it is requested. However, there is at least
one case when a check should be performed against available tokens:
receive commit.

When a commit is transmitted its size is not known until it gets received
and quota is examined by PackParser.checkIfTooLarge(...) method after the
object is successfully read from the client. IOW object size is checked
against available space.

This change extends both QuotaEnforcer and QuotaBackend with methods
that allow retrieving the value of available tokens from implementations.

In case when the particular implementation doesn't provide input on
specific quota it should return QuotaResponse.noOp() object.
However, when more implementations return a value the minimum is returned
so that it is assured that it can be potentially requested.

Note that getting available tokens doesn't result in deducting them
from the tokens bucket, therefore, it is possible to pass over the
available quote however it will be eventually enforced.

Another implementation that was considered was to extend JGit with a
hook that would be called when object is received. It would have been
used to request object size value tokens. It could be performed this way
if/when JGit is modified.

Change-Id: I56091fe98148eff7ed1c1bcd8e4cf32c30081439
Signed-off-by: Jacek Centkowski <jcentkowski@collab.net>

6 files changed, 173 insertions, 0 deletions

diff --git a/java/com/google/gerrit/server/quota/DefaultQuotaBackend.java b/java/com/google/gerrit/server/quota/DefaultQuotaBackend.java
index 479bb529f3..d39e55ca80 100644
--- a/java/com/google/gerrit/server/quota/DefaultQuotaBackend.java
+++ b/java/com/google/gerrit/server/quota/DefaultQuotaBackend.java
@@ -24,6 +24,7 @@ import com.google.gerrit.reviewdb.client.Project;
 import com.google.gerrit.server.CurrentUser;
 import com.google.gerrit.server.plugincontext.PluginSetContext;
 import com.google.gerrit.server.plugincontext.PluginSetEntryContext;
+import com.google.gerrit.server.quota.QuotaResponse.Aggregated;
 import com.google.inject.Inject;
 import com.google.inject.Provider;
 import com.google.inject.Singleton;
@@ -100,6 +101,20 @@ public class DefaultQuotaBackend implements QuotaBackend {
     return QuotaResponse.Aggregated.create(ImmutableList.copyOf(responses));
   }
 
+  private static QuotaResponse.Aggregated availableTokens(
+      PluginSetContext<QuotaEnforcer> quotaEnforcers,
+      String quotaGroup,
+      QuotaRequestContext requestContext) {
+    // PluginSets can change their content when plugins (de-)register. Copy the currently registered
+    // plugins so that we can iterate twice on a stable list.
+    List<PluginSetEntryContext<QuotaEnforcer>> enforcers = ImmutableList.copyOf(quotaEnforcers);
+    List<QuotaResponse> responses = new ArrayList<>(enforcers.size());
+    for (PluginSetEntryContext<QuotaEnforcer> enforcer : enforcers) {
+      responses.add(enforcer.call(p -> p.availableTokens(quotaGroup, requestContext)));
+    }
+    return QuotaResponse.Aggregated.create(responses);
+  }
+
   private static void refillAfterErrorOrException(
       List<PluginSetEntryContext<QuotaEnforcer>> enforcers,
       List<QuotaResponse> collectedResponses,
@@ -158,5 +173,10 @@ public class DefaultQuotaBackend implements QuotaBackend {
       return DefaultQuotaBackend.request(
           quotaEnforcers, quotaGroup, requestContext, numTokens, false);
     }
+
+    @Override
+    public Aggregated availableTokens(String quotaGroup) {
+      return DefaultQuotaBackend.availableTokens(quotaEnforcers, quotaGroup, requestContext);
+    }
   }
 }
diff --git a/java/com/google/gerrit/server/quota/QuotaBackend.java b/java/com/google/gerrit/server/quota/QuotaBackend.java
index d4dc46dd78..11ce61f74e 100644
--- a/java/com/google/gerrit/server/quota/QuotaBackend.java
+++ b/java/com/google/gerrit/server/quota/QuotaBackend.java
@@ -82,5 +82,22 @@ public interface QuotaBackend {
      * not to deduct any quota yet. Can be used to do pre-flight requests where necessary
      */
     QuotaResponse.Aggregated dryRun(String quotaGroup, long tokens);
+
+    /**
+     * Requests a minimum number of tokens available in implementations. This is a pre-flight check
+     * for the exceptional case when the requested number of tokens is not known in advance but
+     * boundary can be specified. For instance, when the commit is received its size is not known
+     * until the transfer happens however one can specify how many bytes can be accepted to meet the
+     * repository size quota.
+     *
+     * <p>By definition, this is not an allocating request, therefore, it should be followed by the
+     * call to {@link #requestTokens(String, long)} when the size gets determined so that quota
+     * could be properly adjusted. It is in developer discretion to ensure that it gets called.
+     * There might be a case when particular quota gets temporarily overbooked when multiple
+     * requests are performed but the following calls to {@link #requestTokens(String, long)} will
+     * fail at the moment when a quota is exhausted. It is not a subject of quota backend to reclaim
+     * tokens that were used due to overbooking.
+     */
+    QuotaResponse.Aggregated availableTokens(String quotaGroup);
   }
 }
diff --git a/java/com/google/gerrit/server/quota/QuotaEnforcer.java b/java/com/google/gerrit/server/quota/QuotaEnforcer.java
index 9c55e11f0c..4c1de14197 100644
--- a/java/com/google/gerrit/server/quota/QuotaEnforcer.java
+++ b/java/com/google/gerrit/server/quota/QuotaEnforcer.java
@@ -60,6 +60,19 @@ public interface QuotaEnforcer {
   QuotaResponse dryRun(String quotaGroup, QuotaRequestContext ctx, long numTokens);
 
   /**
+   * Returns available tokens that can be later requested.
+   *
+   * <p>This is used as a pre-flight check for the exceptional case when the requested number of
+   * tokens is not known in advance. Implementation should not deduct tokens from a bucket. It
+   * should be followed by a call to {@link #requestTokens(String, QuotaRequestContext, long)} with
+   * the number of tokens that were eventually used. It is in {@link QuotaBackend} callers
+   * discretion to ensure that {@link
+   * com.google.gerrit.server.quota.QuotaBackend.WithResource#requestTokens(String, long)} is
+   * called.
+   */
+  QuotaResponse availableTokens(String quotaGroup, QuotaRequestContext ctx);
+
+  /**
    * A previously requested and deducted quota has to be refilled (if possible) because the request
    * failed other quota checks. Implementations can choose to leave this a no-op in case they are
    * the first line of defence (e.g. always deduct HTTP quota even if the request failed for other
diff --git a/java/com/google/gerrit/server/quota/QuotaResponse.java b/java/com/google/gerrit/server/quota/QuotaResponse.java
index 99a7c1d86a..940f731327 100644
--- a/java/com/google/gerrit/server/quota/QuotaResponse.java
+++ b/java/com/google/gerrit/server/quota/QuotaResponse.java
@@ -22,6 +22,7 @@ import com.google.common.collect.ImmutableList;
 import com.google.common.collect.Streams;
 import java.util.Collection;
 import java.util.Optional;
+import java.util.OptionalLong;
 import java.util.stream.Collectors;
 
 @AutoValue
@@ -55,6 +56,10 @@ public abstract class QuotaResponse {
     return new AutoValue_QuotaResponse.Builder().status(Status.OK).build();
   }
 
+  public static QuotaResponse ok(long tokens) {
+    return new AutoValue_QuotaResponse.Builder().status(Status.OK).availableTokens(tokens).build();
+  }
+
   public static QuotaResponse noOp() {
     return new AutoValue_QuotaResponse.Builder().status(Status.NO_OP).build();
   }
@@ -65,12 +70,16 @@ public abstract class QuotaResponse {
 
   public abstract Status status();
 
+  public abstract Optional<Long> availableTokens();
+
   public abstract Optional<String> message();
 
   @AutoValue.Builder
   public abstract static class Builder {
     public abstract QuotaResponse.Builder status(Status status);
 
+    public abstract QuotaResponse.Builder availableTokens(Long tokens);
+
     public abstract QuotaResponse.Builder message(String message);
 
     public abstract QuotaResponse build();
@@ -96,6 +105,13 @@ public abstract class QuotaResponse {
       return responses().stream().filter(r -> r.status().isOk()).collect(toImmutableList());
     }
 
+    public OptionalLong availableTokens() {
+      return responses().stream()
+          .filter(r -> r.status().isOk() && r.availableTokens().isPresent())
+          .mapToLong(r -> r.availableTokens().get())
+          .min();
+    }
+
     public ImmutableList<QuotaResponse> error() {
       return responses().stream().filter(r -> r.status().isError()).collect(toImmutableList());
     }
diff --git a/javatests/com/google/gerrit/acceptance/server/quota/DefaultQuotaBackendIT.java b/javatests/com/google/gerrit/acceptance/server/quota/DefaultQuotaBackendIT.java
index 286e5ae06a..3e902839b4 100644
--- a/javatests/com/google/gerrit/acceptance/server/quota/DefaultQuotaBackendIT.java
+++ b/javatests/com/google/gerrit/acceptance/server/quota/DefaultQuotaBackendIT.java
@@ -125,6 +125,58 @@ public class DefaultQuotaBackendIT extends AbstractDaemonTest {
   }
 
   @Test
+  public void availableTokensForUserAndAccount() {
+    QuotaRequestContext ctx =
+        QuotaRequestContext.builder().user(identifiedAdmin).account(user.id()).build();
+    QuotaResponse r = QuotaResponse.ok(10L);
+    expect(quotaEnforcer.availableTokens("testGroup", ctx)).andReturn(r);
+    replay(quotaEnforcer);
+    assertThat(quotaBackend.user(identifiedAdmin).account(user.id()).availableTokens("testGroup"))
+        .isEqualTo(singletonAggregation(r));
+  }
+
+  @Test
+  public void availableTokensForUserAndProject() {
+    QuotaRequestContext ctx =
+        QuotaRequestContext.builder().user(identifiedAdmin).project(project).build();
+    QuotaResponse r = QuotaResponse.ok(10L);
+    expect(quotaEnforcer.availableTokens("testGroup", ctx)).andReturn(r);
+    replay(quotaEnforcer);
+    assertThat(quotaBackend.user(identifiedAdmin).project(project).availableTokens("testGroup"))
+        .isEqualTo(singletonAggregation(r));
+  }
+
+  @Test
+  public void availableTokensForUserAndChange() throws Exception {
+    Change.Id changeId = createChange().getChange().getId();
+    QuotaRequestContext ctx =
+        QuotaRequestContext.builder()
+            .user(identifiedAdmin)
+            .change(changeId)
+            .project(project)
+            .build();
+    QuotaResponse r = QuotaResponse.ok(10L);
+    expect(quotaEnforcer.availableTokens("testGroup", ctx)).andReturn(r);
+    replay(quotaEnforcer);
+    assertThat(
+            quotaBackend
+                .user(identifiedAdmin)
+                .change(changeId, project)
+                .availableTokens("testGroup"))
+        .isEqualTo(singletonAggregation(r));
+  }
+
+  @Test
+  public void availableTokens() {
+    QuotaRequestContext ctx = QuotaRequestContext.builder().user(identifiedAdmin).build();
+    QuotaResponse r = QuotaResponse.ok(10L);
+    expect(quotaEnforcer.availableTokens("testGroup", ctx)).andReturn(r);
+    replay(quotaEnforcer);
+    assertThat(quotaBackend.user(identifiedAdmin).availableTokens("testGroup"))
+        .isEqualTo(singletonAggregation(r));
+  }
+
+  @Test
   public void requestTokenError() throws Exception {
     QuotaRequestContext ctx = QuotaRequestContext.builder().user(identifiedAdmin).build();
     expect(quotaEnforcer.requestTokens("testGroup", ctx, 1))
@@ -139,6 +191,20 @@ public class DefaultQuotaBackendIT extends AbstractDaemonTest {
   }
 
   @Test
+  public void availableTokensError() throws Exception {
+    QuotaRequestContext ctx = QuotaRequestContext.builder().user(identifiedAdmin).build();
+    expect(quotaEnforcer.availableTokens("testGroup", ctx))
+        .andReturn(QuotaResponse.error("failed"));
+    replay(quotaEnforcer);
+    QuotaResponse.Aggregated result =
+        quotaBackend.user(identifiedAdmin).availableTokens("testGroup");
+    assertThat(result).isEqualTo(singletonAggregation(QuotaResponse.error("failed")));
+    exception.expect(QuotaException.class);
+    exception.expectMessage("failed");
+    result.throwOnError();
+  }
+
+  @Test
   public void requestTokenPluginThrowsAndRethrows() {
     QuotaRequestContext ctx = QuotaRequestContext.builder().user(identifiedAdmin).build();
     expect(quotaEnforcer.requestTokens("testGroup", ctx, 1)).andThrow(new NullPointerException());
@@ -148,6 +214,16 @@ public class DefaultQuotaBackendIT extends AbstractDaemonTest {
     quotaBackend.user(identifiedAdmin).requestToken("testGroup");
   }
 
+  @Test
+  public void availableTokensPluginThrowsAndRethrows() {
+    QuotaRequestContext ctx = QuotaRequestContext.builder().user(identifiedAdmin).build();
+    expect(quotaEnforcer.availableTokens("testGroup", ctx)).andThrow(new NullPointerException());
+    replay(quotaEnforcer);
+
+    exception.expect(NullPointerException.class);
+    quotaBackend.user(identifiedAdmin).availableTokens("testGroup");
+  }
+
   private static QuotaResponse.Aggregated singletonAggregation(QuotaResponse response) {
     return QuotaResponse.Aggregated.create(Collections.singleton(response));
   }
diff --git a/javatests/com/google/gerrit/acceptance/server/quota/MultipleQuotaPluginsIT.java b/javatests/com/google/gerrit/acceptance/server/quota/MultipleQuotaPluginsIT.java
index 8b9ffc5f77..f5f11610c6 100644
--- a/javatests/com/google/gerrit/acceptance/server/quota/MultipleQuotaPluginsIT.java
+++ b/javatests/com/google/gerrit/acceptance/server/quota/MultipleQuotaPluginsIT.java
@@ -32,6 +32,7 @@ import com.google.gerrit.server.quota.QuotaRequestContext;
 import com.google.gerrit.server.quota.QuotaResponse;
 import com.google.inject.Inject;
 import com.google.inject.Module;
+import java.util.OptionalLong;
 import org.easymock.EasyMock;
 import org.junit.Before;
 import org.junit.Test;
@@ -123,4 +124,34 @@ public class MultipleQuotaPluginsIT extends AbstractDaemonTest {
             QuotaResponse.Aggregated.create(
                 ImmutableList.of(QuotaResponse.error("fail"), QuotaResponse.noOp())));
   }
+
+  @Test
+  public void minimumAvailableTokens() {
+    QuotaRequestContext ctx = QuotaRequestContext.builder().user(identifiedAdmin).build();
+    expect(quotaEnforcerA.availableTokens("testGroup", ctx)).andReturn(QuotaResponse.ok(20L));
+    expect(quotaEnforcerB.availableTokens("testGroup", ctx)).andReturn(QuotaResponse.ok(10L));
+
+    replay(quotaEnforcerA);
+    replay(quotaEnforcerB);
+
+    OptionalLong tokens =
+        quotaBackend.user(identifiedAdmin).availableTokens("testGroup").availableTokens();
+    assertThat(tokens.isPresent()).isTrue();
+    assertThat(tokens.getAsLong()).isEqualTo(10L);
+  }
+
+  @Test
+  public void ignoreNoOpForAvailableTokens() {
+    QuotaRequestContext ctx = QuotaRequestContext.builder().user(identifiedAdmin).build();
+    expect(quotaEnforcerA.availableTokens("testGroup", ctx)).andReturn(QuotaResponse.noOp());
+    expect(quotaEnforcerB.availableTokens("testGroup", ctx)).andReturn(QuotaResponse.ok(20L));
+
+    replay(quotaEnforcerA);
+    replay(quotaEnforcerB);
+
+    OptionalLong tokens =
+        quotaBackend.user(identifiedAdmin).availableTokens("testGroup").availableTokens();
+    assertThat(tokens.isPresent()).isTrue();
+    assertThat(tokens.getAsLong()).isEqualTo(20L);
+  }
 }