From 81bf2e6231c9cad40a66138337e9761ba66801dc Mon Sep 17 00:00:00 2001
From: alexesprit <alex.esprit@gmail.com>
Date: Wed, 14 Oct 2020 18:24:09 +0300
Subject: [PATCH] Add throttling plugin for Octokit

---
 package-lock.json | 24 ++++++++++++++++++++
 package.json      |  1 +
 src/octokit.ts    | 56 +++++++++++++++++++++++++++++++++++++++++++++++
 src/update.ts     |  8 ++++---
 4 files changed, 86 insertions(+), 3 deletions(-)
 create mode 100644 src/octokit.ts

diff --git a/package-lock.json b/package-lock.json
index 1e91411..5374110 100644
--- a/package-lock.json
+++ b/package-lock.json
@@ -174,6 +174,25 @@
         "deprecation": "^2.3.1"
       }
     },
+    "@octokit/plugin-throttling": {
+      "version": "3.3.1",
+      "resolved": "https://registry.npmjs.org/@octokit/plugin-throttling/-/plugin-throttling-3.3.1.tgz",
+      "integrity": "sha512-NdUmn52Wr2YCbF8dwjIFQLOcyBlvMJuCKMtybUGhFGA+F7kZtW2yyN73J4hvVL+mfBsjZVqa5vmsP9O8NX9+fg==",
+      "requires": {
+        "@octokit/types": "^5.5.0",
+        "bottleneck": "^2.15.3"
+      },
+      "dependencies": {
+        "@octokit/types": {
+          "version": "5.5.0",
+          "resolved": "https://registry.npmjs.org/@octokit/types/-/types-5.5.0.tgz",
+          "integrity": "sha512-UZ1pErDue6bZNjYOotCNveTXArOMZQFG6hKJfOnGnulVCMcVVi7YIIuuR4WfBhjo7zgpmzn/BkPDnUXtNx+PcQ==",
+          "requires": {
+            "@types/node": ">= 8"
+          }
+        }
+      }
+    },
     "@octokit/request": {
       "version": "5.4.7",
       "resolved": "https://registry.npmjs.org/@octokit/request/-/request-5.4.7.tgz",
@@ -412,6 +431,11 @@
       "resolved": "https://registry.npmjs.org/before-after-hook/-/before-after-hook-2.1.0.tgz",
       "integrity": "sha512-IWIbu7pMqyw3EAJHzzHbWa85b6oud/yfKYg5rqB5hNE8CeMi3nX+2C2sj0HswfblST86hpVEOAb9x34NZd6P7A=="
     },
+    "bottleneck": {
+      "version": "2.19.5",
+      "resolved": "https://registry.npmjs.org/bottleneck/-/bottleneck-2.19.5.tgz",
+      "integrity": "sha512-VHiNCbI1lKdl44tGrhNfU3lup0Tj/ZBMJB5/2ZbNXRCPuRCO7ed2mgcK4r17y+KB2EfuYuRaVlwNbAeaWGSpbw=="
+    },
     "brace-expansion": {
       "version": "1.1.11",
       "resolved": "https://registry.npmjs.org/brace-expansion/-/brace-expansion-1.1.11.tgz",
diff --git a/package.json b/package.json
index 5cd5ff2..3f748ef 100644
--- a/package.json
+++ b/package.json
@@ -16,6 +16,7 @@
   "dependencies": {
     "@actions/core": "1.2.6",
     "@actions/github": "4.0.0",
+    "@octokit/plugin-throttling": "3.3.1",
     "fast-glob": "3.2.4"
   },
   "devDependencies": {
diff --git a/src/octokit.ts b/src/octokit.ts
new file mode 100644
index 0000000..2c8858b
--- /dev/null
+++ b/src/octokit.ts
@@ -0,0 +1,56 @@
+import { Octokit } from '@octokit/core';
+import { throttling } from '@octokit/plugin-throttling';
+
+import { GitHub, getOctokitOptions } from '@actions/github/lib/utils';
+
+const OctokitConstructor = GitHub.plugin(throttling);
+
+const maxRetryCount = 1;
+
+interface RequestOptions {
+	method: string;
+	url: string;
+	request: RequestInfo;
+}
+
+interface RequestInfo {
+	retryCount: number;
+}
+
+export function createOctokit(
+	accessToken: string
+): InstanceType<typeof GitHub> {
+	const octokit = new OctokitConstructor(
+		getOctokitOptions(accessToken, {
+			throttle: { onRateLimit, onAbuseLimit },
+		})
+	);
+
+	return octokit;
+}
+
+function onRateLimit(
+	retryAfter: number,
+	options: RequestOptions,
+	octokit: Octokit
+): boolean {
+	octokit.log.warn(
+		`Request quota exhausted for request ${options.method} ${options.url}`
+	);
+	if (options.request.retryCount < maxRetryCount) {
+		octokit.log.info(`Retrying after ${retryAfter} seconds!`);
+		return true;
+	}
+
+	return false;
+}
+
+function onAbuseLimit(
+	_: number,
+	options: RequestOptions,
+	octokit: Octokit
+): void {
+	octokit.log.warn(
+		`Abuse detected for request ${options.method} ${options.url}`
+	);
+}
diff --git a/src/update.ts b/src/update.ts
index 7ba9ed6..c51d096 100644
--- a/src/update.ts
+++ b/src/update.ts
@@ -3,9 +3,11 @@
 import { readFile, existsSync } from 'fs';
 import { promisify } from 'util';
 
-import { getOctokit, context } from '@actions/github';
+import { GitHub } from '@actions/github/lib/utils';
+import { context } from '@actions/github';
 
 import { UpdaterOptions, isNotNull } from './util';
+import { createOctokit } from './octokit';
 
 const readFileAsync = promisify(readFile);
 
@@ -27,12 +29,12 @@ interface UpdateResult {
 }
 
 export class Updater {
-	private octokit: ReturnType<typeof getOctokit>;
+	private octokit: InstanceType<typeof GitHub>;
 	private message: string;
 	private defaultBranch: string | null;
 
 	constructor(options: UpdaterOptions) {
-		this.octokit = getOctokit(options.token);
+		this.octokit = createOctokit(options.token);
 
 		this.message = options.message;
 		this.defaultBranch = options.branch || null;