Merge pull request #717 from salesforce/rate-limit

DFIU to support rate limiting the PR creation in a given amount of time.

README.md

@@ -141,7 +141,11 @@ named arguments:
   -s {true,false}, --skipprcreation {true,false}
                          Only update image tag store. Skip creating PRs
   -x X                   comment snippet mentioned in line just before FROM instruction for ignoring a child image. Defaults to 'no-dfiu'
-
+  -r, --rate-limit-pr-creations 
+                         Enable rateLimiting for throttling the number of PRs DFIU will cut over a period of time. 
+                         The argument value should be in format "<positive_integer>-<ISO-8601_formatted_time>". For example "--rate-limit-pr-creations 60-PT1H" to create 60 PRs per hour.
+                         Default is not set, this means no ratelimiting is imposed.
+                         
 subcommands:
   Specify which feature to perform
 
@@ -213,6 +217,41 @@ Example:
 FROM imagename:imagetag # no-dfiu
 ```
 
+### PR throttling
+
+In case you want to throttle the number of PRs cut by DFIU over a period of time,
+set --rate-limit-pr-creations with appropriate value. 
+
+##### Default case:
+
+By default, this feature is disabled. This will be enabled when argument ``--rate-limit-pr-creations`` will be passed
+with appropriate value.
+
+```
+example: dockerfile-image-update all image-tag-store-repo-falcon //throttling will be disabled by default
+```
+
+##### Configuring the rate limit:
+
+Below are some examples that will throttle the number of PRs cut based on values passed to the
+argument ``--rate-limit-pr-creations``
+The argument value should be in format ``<positive_integer>-<ISO-8601_formatted_time>``.
+For example ``--rate-limit-pr-creations 60-PT1H`` would mean the tool will cut 60 PRs every hour and the rate of adding
+a new PR will be (PT1H/60) i.e. one minute.
+This will distribute the load uniformly and avoid sudden spikes, The process will go in waiting state until next PR
+could be sent.
+
+Below are some more examples:
+
+```
+Usage: 
+    dockerfile-image-update --rate-limit-pr-creations 60-PT1H all image-tag-store-repo-falcon //DFIU can send up to 60 PRs per hour.
+    dockerfile-image-update --rate-limit-pr-creations 500-PT1H all image-tag-store-repo-falcon //DFIU can send up to 500 PRs per hour.
+    dockerfile-image-update --rate-limit-pr-creations 86400-PT24H all image-tag-store-repo-falcon //DFIU can send up to 1 PRs per second.
+    dockerfile-image-update --rate-limit-pr-creations 1-PT1S all image-tag-store-repo-falcon //Same as above. DFIU can send up to 1 PRs per second.
+    dockerfile-image-update --rate-limit-pr-creations 5000 all image-tag-store-repo-falcon //rate limiting will be disabled because argument is not in correct format.
+```
+
 ## Developer Guide
 
 ### Building

dockerfile-image-update/pom.xml

@@ -139,6 +139,11 @@
             <version>3.24.0</version>
             <scope>compile</scope>
         </dependency>
+        <dependency>
+            <groupId>com.bucket4j</groupId>
+            <artifactId>bucket4j-core</artifactId>
+            <version>8.1.1</version>
+        </dependency>
     </dependencies>
 
     <build>

dockerfile-image-update/src/main/java/com/salesforce/dockerfileimageupdate/CommandLine.java

@@ -26,6 +26,7 @@
 import java.util.Set;
 import java.util.TreeSet;
 import java.util.stream.Collectors;
+import static com.salesforce.dockerfileimageupdate.utils.Constants.*;
 
 /**
  * Created by minho-park on 6/29/2016.
@@ -84,6 +85,12 @@ static ArgumentParser getArgumentParser() {
         parser.addArgument("-x")
                 .help("comment snippet mentioned in line just before FROM instruction for ignoring a child image. " +
                         "Defaults to 'no-dfiu'");
+        parser.addArgument("-r", "--" + RATE_LIMIT_PR_CREATION)
+                .type(String.class)
+                .setDefault("")
+                .required(false)
+                .help("Use RateLimiting when sending PRs. RateLimiting is enabled only if this value is set it's disabled by default.");
+
         return parser;
     }
 

dockerfile-image-update/src/main/java/com/salesforce/dockerfileimageupdate/subcommands/impl/All.java

@@ -19,6 +19,7 @@
 import com.salesforce.dockerfileimageupdate.utils.ImageStoreUtil;
 import com.salesforce.dockerfileimageupdate.utils.ProcessingErrors;
 import com.salesforce.dockerfileimageupdate.utils.PullRequests;
+import com.salesforce.dockerfileimageupdate.utils.RateLimiter;
 import net.sourceforge.argparse4j.inf.Namespace;
 import org.kohsuke.github.*;
 import org.slf4j.Logger;
@@ -37,20 +38,22 @@ public class All implements ExecutableWithNamespace {
     @Override
     public void execute(final Namespace ns, final DockerfileGitHubUtil dockerfileGitHubUtil) throws Exception {
         loadDockerfileGithubUtil(dockerfileGitHubUtil);
+        RateLimiter rateLimiter =  RateLimiter.getInstance(ns);
         String store = ns.get(Constants.STORE);
         try {
             ImageTagStore imageTagStore = ImageStoreUtil.initializeImageTagStore(this.dockerfileGitHubUtil, store);
             List<ImageTagStoreContent> imageNamesWithTag = imageTagStore.getStoreContent(dockerfileGitHubUtil, store);
             Integer numberOfImagesToProcess = imageNamesWithTag.size();
-            List<ProcessingErrors> imagesThatCouldNotBeProcessed = processImagesWithTag(ns, imageNamesWithTag);
+            List<ProcessingErrors> imagesThatCouldNotBeProcessed = processImagesWithTag(ns, imageNamesWithTag, rateLimiter);
             printSummary(imagesThatCouldNotBeProcessed, numberOfImagesToProcess);
         } catch (Exception e) {
             log.error("Encountered issues while initializing the image tag store or getting its contents. Cannot continue. Exception: ", e);
             System.exit(2);
         }
     }
 
-    protected List<ProcessingErrors> processImagesWithTag(Namespace ns, List<ImageTagStoreContent> imageNamesWithTag) {
+    protected List<ProcessingErrors> processImagesWithTag(Namespace ns, List<ImageTagStoreContent> imageNamesWithTag,
+                                                          RateLimiter rateLimiter) {
         Integer gitApiSearchLimit = ns.get(Constants.GIT_API_SEARCH_LIMIT);
         Map<String, Boolean> orgsToIncludeInSearch = new HashMap<>();
         if (ns.get(Constants.GIT_ORG) != null) {
@@ -64,13 +67,15 @@ protected List<ProcessingErrors> processImagesWithTag(Namespace ns, List<ImageTa
         for (ImageTagStoreContent content : imageNamesWithTag) {
             String image = content.getImageName();
             String tag = content.getTag();
-            failureMessage = processImageWithTag(image, tag, ns, orgsToIncludeInSearch, gitApiSearchLimit);
+            failureMessage = processImageWithTag(image, tag, ns, orgsToIncludeInSearch, gitApiSearchLimit,
+                    rateLimiter);
             failureMessage.ifPresent(message -> imagesThatCouldNotBeProcessed.add(processErrorMessages(image, tag, Optional.of(message))));
         }
         return imagesThatCouldNotBeProcessed;
     }
 
-    protected Optional<Exception> processImageWithTag(String image, String tag, Namespace ns, Map<String, Boolean> orgsToIncludeInSearch, Integer gitApiSearchLimit) {
+    protected Optional<Exception> processImageWithTag(String image, String tag, Namespace ns, Map<String, Boolean> orgsToIncludeInSearch,
+                                                      Integer gitApiSearchLimit, RateLimiter rateLimiter) {
         Optional<Exception> failureMessage = Optional.empty();
         try {
             PullRequests pullRequests = getPullRequests();
@@ -85,7 +90,7 @@ protected Optional<Exception> processImageWithTag(String image, String tag, Name
                 while (it.hasNext()){
                     try {
                         pullRequests.prepareToCreate(ns, pullRequestSender,
-                                it.next(), gitForkBranch, dockerfileGitHubUtil);
+                                it.next(), gitForkBranch, dockerfileGitHubUtil, rateLimiter);
                     } catch (IOException e) {
                         log.error("Could not send pull request for image {}.", image);
                         failureMessage = Optional.of(e);

dockerfile-image-update/src/main/java/com/salesforce/dockerfileimageupdate/subcommands/impl/Child.java

@@ -16,6 +16,7 @@
 import com.salesforce.dockerfileimageupdate.utils.Constants;
 import com.salesforce.dockerfileimageupdate.utils.DockerfileGitHubUtil;
 import com.salesforce.dockerfileimageupdate.utils.ImageStoreUtil;
+import com.salesforce.dockerfileimageupdate.utils.RateLimiter;
 import net.sourceforge.argparse4j.inf.Namespace;
 import org.kohsuke.github.GHRepository;
 import org.slf4j.Logger;
@@ -36,6 +37,7 @@ public void execute(final Namespace ns, final DockerfileGitHubUtil dockerfileGit
         String forceTag = ns.get(Constants.FORCE_TAG);
         String store = ns.get(Constants.STORE);
         ImageStoreUtil imageStoreUtil = getImageStoreUtil();
+        RateLimiter rateLimiter =  RateLimiter.getInstance(ns);
         try {
             ImageTagStore imageTagStore = imageStoreUtil.initializeImageTagStore(dockerfileGitHubUtil, store);
             /* Updates store if a store is specified. */
@@ -68,7 +70,8 @@ public void execute(final Namespace ns, final DockerfileGitHubUtil dockerfileGit
         dockerfileGitHubUtil.createPullReq(repo,
                 gitForkBranch.getBranchName(),
                 fork,
-                pullRequestInfo);
+                pullRequestInfo,
+                rateLimiter);
     }
 
     protected ImageStoreUtil getImageStoreUtil(){

dockerfile-image-update/src/main/java/com/salesforce/dockerfileimageupdate/subcommands/impl/Parent.java

@@ -18,6 +18,7 @@
 import com.salesforce.dockerfileimageupdate.utils.DockerfileGitHubUtil;
 import com.salesforce.dockerfileimageupdate.utils.ImageStoreUtil;
 import com.salesforce.dockerfileimageupdate.utils.PullRequests;
+import com.salesforce.dockerfileimageupdate.utils.RateLimiter;
 import net.sourceforge.argparse4j.inf.Namespace;
 import org.kohsuke.github.GHContent;
 import org.kohsuke.github.PagedSearchIterable;
@@ -59,6 +60,7 @@ public void execute(final Namespace ns, DockerfileGitHubUtil dockerfileGitHubUti
         PullRequests pullRequests = getPullRequests();
         GitHubPullRequestSender pullRequestSender = getPullRequestSender(dockerfileGitHubUtil, ns);
         GitForkBranch gitForkBranch = getGitForkBranch(ns);
+        RateLimiter rateLimiter = RateLimiter.getInstance(ns);
         log.info("Finding Dockerfiles with the given image...");
 
         Integer gitApiSearchLimit = ns.get(Constants.GIT_API_SEARCH_LIMIT);
@@ -69,7 +71,8 @@ public void execute(final Namespace ns, DockerfileGitHubUtil dockerfileGitHubUti
             for (int i = 0; i < contentsFoundWithImage.size(); i++ ) {
                 try {
                     pullRequests.prepareToCreate(ns, pullRequestSender,
-                            contentsFoundWithImage.get(i), gitForkBranch, dockerfileGitHubUtil);
+                            contentsFoundWithImage.get(i), gitForkBranch,
+                            dockerfileGitHubUtil, rateLimiter);
                 } catch (IOException e) {
                     log.error("Could not send pull request.", e);
                 }

dockerfile-image-update/src/main/java/com/salesforce/dockerfileimageupdate/utils/Constants.java

@@ -9,6 +9,8 @@
 package com.salesforce.dockerfileimageupdate.utils;
 
 
+import java.time.Duration;
+
 /**
  * @author minho-park
  */
@@ -37,5 +39,13 @@ private Constants() {
     public static final String GIT_API_SEARCH_LIMIT = "ghapisearchlimit";
     public static final String SKIP_PR_CREATION = "skipprcreation";
     public static final String IGNORE_IMAGE_STRING = "x";
+    public static final String RATE_LIMIT_PR_CREATION = "rate-limit-pr-creations";
+    //max number of PRs to be sent (or tokens to be added)  per DEFAULT_RATE_LIMIT_DURATION(per hour in this case)
+    public static final long DEFAULT_RATE_LIMIT = 60;
+
+    public static final long DEFAULT_CONSUMING_TOKEN_RATE = 1;
+    public static final Duration DEFAULT_RATE_LIMIT_DURATION = Duration.ofMinutes(DEFAULT_RATE_LIMIT);
+    //token adding rate(here:a token added every 2 minutes in the bucket)
+    public static final Duration DEFAULT_TOKEN_ADDING_RATE = Duration.ofMinutes(DEFAULT_CONSUMING_TOKEN_RATE);
 
 }

dockerfile-image-update/src/main/java/com/salesforce/dockerfileimageupdate/utils/DockerfileGitHubUtil.java

@@ -357,10 +357,22 @@ public GitHubJsonStore getGitHubJsonStore(String store) {
     }
 
     public void createPullReq(GHRepository origRepo,
-                              String branch, GHRepository forkRepo,
-                              PullRequestInfo pullRequestInfo) throws InterruptedException, IOException {
+                              String branch,
+                              GHRepository forkRepo,
+                              PullRequestInfo pullRequestInfo,
+                              RateLimiter rateLimiter) throws InterruptedException, IOException {
         // TODO: This may loop forever in the event of constant -1 pullRequestExitCodes...
         while (true) {
+            // TODO: accept rateLimiter Optional with option to get no-op rateLimiter
+            // where it's not required.
+            if(rateLimiter != null) {
+                log.info("Trying to consume a token before creating pull request..");
+                // Consume a token from the token bucket.
+                // If a token is not available this method will block until
+                // the refill adds one to the bucket.
+                rateLimiter.consume();
+                log.info("Token consumed, proceeding with PR creation..");
+            }
             int pullRequestExitCode = gitHubUtil.createPullReq(origRepo,
                     branch, forkRepo, pullRequestInfo.getTitle(), pullRequestInfo.getBody());
             if (pullRequestExitCode == 0) {
@@ -484,7 +496,8 @@ public void changeDockerfiles(Namespace ns,
                                      Multimap<String, GitHubContentToProcess> pathToDockerfilesInParentRepo,
                                      GitHubContentToProcess gitHubContentToProcess,
                                      List<String> skippedRepos,
-                                     GitForkBranch gitForkBranch) throws IOException,
+                                     GitForkBranch gitForkBranch,
+                                     RateLimiter rateLimiter) throws IOException,
             InterruptedException {
         // Should we skip doing a getRepository just to fill in the parent value? We already know this to be the parent...
         GHRepository parent = gitHubContentToProcess.getParent();
@@ -527,7 +540,8 @@ public void changeDockerfiles(Namespace ns,
             createPullReq(parent,
                     gitForkBranch.getBranchName(),
                     forkedRepo,
-                    pullRequestInfo);
+                    pullRequestInfo,
+                    rateLimiter);
         }
     }
 }

dockerfile-image-update/src/main/java/com/salesforce/dockerfileimageupdate/utils/PullRequests.java

@@ -16,7 +16,8 @@ public void prepareToCreate(final Namespace ns,
                                             GitHubPullRequestSender pullRequestSender,
                                             PagedSearchIterable<GHContent> contentsFoundWithImage,
                                             GitForkBranch gitForkBranch,
-                                            DockerfileGitHubUtil dockerfileGitHubUtil) throws IOException {
+                                            DockerfileGitHubUtil dockerfileGitHubUtil,
+                                            RateLimiter rateLimiter) throws IOException {
         Multimap<String, GitHubContentToProcess> pathToDockerfilesInParentRepo =
                 pullRequestSender.forkRepositoriesFoundAndGetPathToDockerfiles(contentsFoundWithImage, gitForkBranch);
         List<IOException> exceptions = new ArrayList<>();
@@ -28,7 +29,8 @@ public void prepareToCreate(final Namespace ns,
                 try {
                     dockerfileGitHubUtil.changeDockerfiles(ns,
                             pathToDockerfilesInParentRepo,
-                            forkWithContentPaths.get(), skippedRepos, gitForkBranch);
+                            forkWithContentPaths.get(), skippedRepos,
+                            gitForkBranch, rateLimiter);
                 } catch (IOException | InterruptedException e) {
                     log.error(String.format("Error changing Dockerfile for %s", forkWithContentPaths.get().getParent().getFullName()), e);
                     exceptions.add((IOException) e);