Cache Git repository size(s), update cached size after receive-pack

Computing the occupied disk size for a repository is expensive. Cache
the repository size to avoid computing it for every receive-pack.
Implement the PostReceiveHook to get notified when a new pack is
received and increase the cached repository size for the size of the
received pack.

The cached repository size is currently only increased for the size of
the received pack. The pack size is the most significant contributor to
the increased repository size and the quota doesn't need to be
byte-precise. The size of the newly created references or reflog entries
is not added to the cached value of the repository size. This could be
improved in a follow up change.

Change-Id: Ia6be19c7f40b31301cf0431efd366566d7254a43
diff --git a/BUCK b/BUCK
index 6760b75..c625ce6 100644
--- a/BUCK
+++ b/BUCK
@@ -1,6 +1,6 @@
 include_defs('//lib/maven.defs')
 
-API_VERSION = '2.9-SNAPSHOT'
+API_VERSION = '2.10-SNAPSHOT'
 REPO = MAVEN_LOCAL
 
 gerrit_plugin(
diff --git a/src/main/java/com/googlesource/gerrit/plugins/quota/MaxRepositorySizeQuota.java b/src/main/java/com/googlesource/gerrit/plugins/quota/MaxRepositorySizeQuota.java
index d27641d..6e8b0cb 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/quota/MaxRepositorySizeQuota.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/quota/MaxRepositorySizeQuota.java
@@ -14,11 +14,18 @@
 
 package com.googlesource.gerrit.plugins.quota;
 
+import com.google.common.cache.CacheLoader;
+import com.google.common.cache.LoadingCache;
 import com.google.gerrit.reviewdb.client.Project;
+import com.google.gerrit.server.cache.CacheModule;
+import com.google.gerrit.server.config.GerritServerConfig;
+import com.google.gerrit.server.config.SitePaths;
 import com.google.gerrit.server.git.ReceivePackInitializer;
 import com.google.gerrit.server.git.GitRepositoryManager;
 import com.google.inject.Inject;
+import com.google.inject.Module;
 import com.google.inject.Singleton;
+import com.google.inject.name.Named;
 
 import java.io.File;
 import java.io.IOException;
@@ -27,26 +34,51 @@
 import java.nio.file.Path;
 import java.nio.file.SimpleFileVisitor;
 import java.nio.file.attribute.BasicFileAttributes;
+import java.util.Collection;
+import java.util.concurrent.ExecutionException;
+import java.util.concurrent.TimeUnit;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.apache.commons.lang.mutable.MutableLong;
 import org.eclipse.jgit.errors.RepositoryNotFoundException;
+import org.eclipse.jgit.lib.Config;
 import org.eclipse.jgit.lib.Repository;
+import org.eclipse.jgit.transport.PostReceiveHook;
+import org.eclipse.jgit.transport.ReceiveCommand;
 import org.eclipse.jgit.transport.ReceivePack;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 @Singleton
-class MaxRepositorySizeQuota implements ReceivePackInitializer {
+class MaxRepositorySizeQuota implements ReceivePackInitializer, PostReceiveHook {
   private static final Logger log = LoggerFactory
       .getLogger(MaxRepositorySizeQuota.class);
 
+  private static final String CACHE_NAME = "repo_size";
+
+  static Module module() {
+    return new CacheModule() {
+      protected void configure() {
+        persist(CACHE_NAME, Project.NameKey.class, AtomicLong.class)
+            .loader(Loader.class)
+            .expireAfterWrite(1, TimeUnit.DAYS);
+      }
+    };
+  }
+
   private final QuotaFinder quotaFinder;
-  private final GitRepositoryManager gitManager;
+  private final LoadingCache<Project.NameKey, AtomicLong> cache;
+  private final SitePaths site;
+  private final Path basePath;
 
   @Inject
-  MaxRepositorySizeQuota(QuotaFinder quotaFinder, GitRepositoryManager gitManager) {
+  MaxRepositorySizeQuota(QuotaFinder quotaFinder,
+      @Named(CACHE_NAME) LoadingCache<Project.NameKey, AtomicLong> cache,
+      SitePaths site, @GerritServerConfig final Config cfg) {
     this.quotaFinder = quotaFinder;
-    this.gitManager = gitManager;
+    this.cache = cache;
+    this.site = site;
+    basePath = site.resolve(cfg.getString("gerrit", null, "basePath")).toPath();
   }
 
   @Override
@@ -62,35 +94,69 @@
     }
 
     try {
-      long maxPackSize = Math.max(0, maxRepoSize - getDiskUsage(project));
+      long maxPackSize = Math.max(0, maxRepoSize - cache.get(project).get());
       rp.setMaxPackSizeLimit(maxPackSize);
-    } catch (IOException e) {
+    } catch (ExecutionException e) {
       log.warn("Couldn't setMaxPackSizeLimit on receive-pack for "
           + project.get(), e);
     }
   }
 
-  private long getDiskUsage(Project.NameKey project) throws IOException {
-    Repository git = gitManager.openRepository(project);
+  @Override
+  public void onPostReceive(ReceivePack rp, Collection<ReceiveCommand> commands) {
+    Project.NameKey project = projectName(rp);
     try {
-      return getDiskUsage(git.getDirectory());
-    } finally {
-      git.close();
+      cache.get(project).getAndAdd(rp.getPackSize());
+    } catch (ExecutionException e) {
+      log.warn("Couldn't process onPostReceive for " + project.get(), e);
     }
   }
 
-  private static long getDiskUsage(File dir) throws IOException {
-    final MutableLong size = new MutableLong();
-    Files.walkFileTree(dir.toPath(), new SimpleFileVisitor<Path>() {
-      @Override
-      public FileVisitResult visitFile(Path path, BasicFileAttributes attrs)
-          throws IOException {
-        if (attrs.isRegularFile()) {
-          size.add(attrs.size());
-        }
-        return FileVisitResult.CONTINUE;
+  private Project.NameKey projectName(ReceivePack rp) {
+    Path gitDir = rp.getRepository().getDirectory().toPath();
+    if (gitDir.startsWith(basePath)) {
+      String p = basePath.relativize(gitDir).toString();
+      String n = p.substring(0, p.length() - ".git".length());
+      return new Project.NameKey(n);
+    } else {
+      log.warn("Couldn't determine the project name from " + gitDir);
+      return null;
+    }
+  }
+
+  @Singleton
+  static class Loader extends CacheLoader<Project.NameKey, AtomicLong> {
+
+    private final GitRepositoryManager gitManager;
+
+    @Inject
+    Loader(GitRepositoryManager gitManager) {
+      this.gitManager = gitManager;
+    }
+
+    @Override
+    public AtomicLong load(Project.NameKey project) throws IOException {
+      Repository git = gitManager.openRepository(project);
+      try {
+        return new AtomicLong(getDiskUsage(git.getDirectory()));
+      } finally {
+        git.close();
       }
-    });
-    return size.longValue();
+    }
+
+    private static long getDiskUsage(File dir) throws IOException {
+      final MutableLong size = new MutableLong();
+      Files.walkFileTree(dir.toPath(), new SimpleFileVisitor<Path>() {
+        @Override
+        public FileVisitResult visitFile(Path path, BasicFileAttributes attrs)
+            throws IOException {
+          if (attrs.isRegularFile()) {
+            size.add(attrs.size());
+          }
+          return FileVisitResult.CONTINUE;
+        }
+      });
+      return size.longValue();
+    }
   }
 }
diff --git a/src/main/java/com/googlesource/gerrit/plugins/quota/Module.java b/src/main/java/com/googlesource/gerrit/plugins/quota/Module.java
index c51f670..c7c851c 100644
--- a/src/main/java/com/googlesource/gerrit/plugins/quota/Module.java
+++ b/src/main/java/com/googlesource/gerrit/plugins/quota/Module.java
@@ -19,6 +19,8 @@
 import com.google.gerrit.server.validators.ProjectCreationValidationListener;
 import com.google.inject.AbstractModule;
 
+import org.eclipse.jgit.transport.PostReceiveHook;
+
 class Module extends AbstractModule {
   @Override
   protected void configure() {
@@ -26,5 +28,8 @@
         .to(MaxRepositoriesQuotaValidator.class);
     DynamicSet.bind(binder(), ReceivePackInitializer.class)
         .to(MaxRepositorySizeQuota.class);
+    DynamicSet.bind(binder(), PostReceiveHook.class)
+        .to(MaxRepositorySizeQuota.class);
+    install(MaxRepositorySizeQuota.module());
   }
 }