[HUD] Button for updating Dr. CI (#5844)

Add button on top right where logged in users can update Dr. CI
https://torchci-git-csl-drcibutton-fbopensource.vercel.app/pr/138513

I am failing to log in on the preview, but it works locally, probably
due to prod vs preview env vars

Rate limit users to 10 calls / hr

Is this a good idea? No clue

Should this be limited to people with write access?

torchci/components/DrCIButton.tsx

@@ -0,0 +1,94 @@
+import { Button, CircularProgress, Tooltip } from "@mui/material";
+import { useSession } from "next-auth/react";
+import { useEffect, useState } from "react";
+
+export default function DrCIButton({
+  owner,
+  repo,
+  prNumber,
+}: {
+  owner: string;
+  repo: string;
+  prNumber: number;
+}) {
+  const session = useSession();
+  const loggedIn = session.status === "authenticated" && session.data !== null;
+  // loading, clickable, failed, rateLimited
+  const [buttonState, setButtonState] = useState("clickable");
+
+  const url = `/api/drci/drci?prNumber=${prNumber}`;
+  if (buttonState == "loading" && loggedIn) {
+    fetch(url, {
+      method: "POST",
+      body: JSON.stringify({ repo }),
+      headers: {
+        Authorization: session.data!["accessToken"],
+        "Cache-Control": "no-cache",
+        "Content-Type": "application/json",
+      },
+    }).then((res) => {
+      if (res.status == 429) {
+        setButtonState("rateLimited");
+        return;
+      }
+      if (!res.ok) {
+        setButtonState("failed");
+        return;
+      }
+      setButtonState("clickable");
+      return res.json();
+    });
+  }
+
+  useEffect(() => {
+    if (buttonState == "failed" || buttonState == "rateLimited") {
+      setTimeout(() => {
+        setButtonState("clickable");
+      }, 5000);
+    }
+  }, [buttonState]);
+
+  return (
+    <Tooltip
+      title={
+        owner == "pytorch"
+          ? loggedIn
+            ? "Click to update Dr. CI.  This might take a while."
+            : "You must be logged in to update Dr. CI"
+          : "Dr. CI is only available for pytorch org PRs"
+      }
+    >
+      <span>
+        <Button
+          variant="contained"
+          disableElevation
+          disabled={
+            !loggedIn || buttonState != "clickable" || owner != "pytorch"
+          }
+          onClick={() => {
+            setButtonState("loading");
+          }}
+        >
+          {buttonState == "loading" && (
+            <CircularProgress
+              size={20}
+              sx={{
+                color: "primary",
+                position: "absolute",
+                top: "50%",
+                left: "50%",
+                marginTop: "-10px",
+                marginLeft: "-10px",
+              }}
+            />
+          )}
+          {buttonState == "rateLimited"
+            ? "Exceeded Rate Limit"
+            : buttonState == "failed"
+            ? "Failed to Update"
+            : "Update Dr. CI"}
+        </Button>
+      </span>
+    </Tooltip>
+  );
+}

torchci/lib/clickhouse.ts

@@ -17,6 +17,14 @@ export function getClickhouseClient() {
   });
 }
 
+export function getClickhouseClientWritable() {
+  return createClient({
+    host: process.env.CLICKHOUSE_HUD_USER_URL ?? "http://localhost:8123",
+    username: process.env.CLICKHOUSE_HUD_USER_WRITE_USERNAME ?? "default",
+    password: process.env.CLICKHOUSE_HUD_USER_WRITE_PASSWORD ?? "",
+  });
+}
+
 export async function queryClickhouse(
   query: string,
   params: Record<string, unknown>

torchci/lib/rateLimit.ts

@@ -0,0 +1,39 @@
+// Rate limit users
+import dayjs from "dayjs";
+import { getClickhouseClientWritable, queryClickhouse } from "./clickhouse";
+
+async function checkRateLimit(user: string, key: string) {
+  const res = await queryClickhouse(
+    `
+select count() as count from misc.rate_limit
+where user = {user: String} and key = {key: String} and time_inserted > {timestamp: DateTime}`,
+    {
+      user,
+      key,
+      timestamp: dayjs()
+        .utc()
+        .subtract(1, "hour")
+        .format("YYYY-MM-DD HH:mm:ss"),
+    }
+  );
+  if (res.length == 0) {
+    return 0;
+  }
+  return res[0].count;
+}
+
+async function incrementRateLimit(user: string, key: string) {
+  await getClickhouseClientWritable().insert({
+    table: "misc.rate_limit",
+    values: [[user, key, dayjs().utc().format("YYYY-MM-DD HH:mm:ss")]],
+  });
+}
+
+export async function drCIRateLimitExceeded(user: string) {
+  const rateLimit = 10;
+  return (await checkRateLimit(user, "DrCI")) >= rateLimit;
+}
+
+export async function incrementDrCIRateLimit(user: string) {
+  return await incrementRateLimit(user, "DrCI");
+}

torchci/pages/[repoOwner]/[repoName]/pull/[prNumber].tsx

@@ -1,4 +1,6 @@
+import { Stack } from "@mui/material";
 import CommitStatus from "components/CommitStatus";
+import DrCIButton from "components/DrCIButton";
 import { useSetTitle } from "components/DynamicTitle";
 import ErrorBoundary from "components/ErrorBoundary";
 import { useCHContext } from "components/UseClickhouseProvider";
@@ -133,16 +135,30 @@ function Page() {
   }
   return (
     <div>
-      <h1>
-        {prData.title}{" "}
-        <code>
-          <a
-            href={`https://github.com/${repoOwner}/${repoName}/pull/${prNumber}`}
-          >
-            #{prNumber}
-          </a>
-        </code>
-      </h1>
+      <Stack
+        direction="row"
+        spacing={0}
+        sx={{
+          justifyContent: "space-between",
+          alignItems: "flex-start",
+        }}
+      >
+        <h1>
+          {prData.title}{" "}
+          <code>
+            <a
+              href={`https://github.com/${repoOwner}/${repoName}/pull/${prNumber}`}
+            >
+              #{prNumber}
+            </a>
+          </code>
+        </h1>
+        <DrCIButton
+          prNumber={prNumber ? parseInt(prNumber as string) : 0}
+          owner={repoOwner as string}
+          repo={repoName as string}
+        />
+      </Stack>
       <CommitHeader
         repoOwner={repoOwner as string}
         repoName={repoName as string}

torchci/pages/api/drci/drci.ts

@@ -31,7 +31,7 @@ import {
   fetchFailedJobsFromCommits,
   fetchRecentWorkflows,
 } from "lib/fetchRecentWorkflows";
-import { getOctokit } from "lib/github";
+import { getOctokit, getOctokitWithUserToken } from "lib/github";
 import {
   backfillMissingLog,
   getDisabledTestIssues,
@@ -44,6 +44,7 @@ import {
   removeCancelledJobAfterRetry,
   removeJobNameSuffix,
 } from "lib/jobUtils";
+import { drCIRateLimitExceeded, incrementDrCIRateLimit } from "lib/rateLimit";
 import { getS3Client } from "lib/s3";
 import { IssueData, PRandJobs, RecentWorkflowsData } from "lib/types";
 import _ from "lodash";
@@ -67,20 +68,38 @@ export default async function handler(
 ) {
   const authorization = req.headers.authorization;
 
-  if (authorization === process.env.DRCI_BOT_KEY) {
+  if (authorization == process.env.DRCI_BOT_KEY) {
+    // Dr. CI bot key is used to update the comment, probably called from the
+    // update Dr. CI workflow
+  } else if (authorization) {
+    // Authorization provided, probably a user calling it.
+    // Check that they are only updating a single PR
     const { prNumber } = req.query;
-    const { repo }: UpdateCommentBody = req.body;
-    const octokit = await getOctokit(OWNER, repo);
-
-    const failures = await updateDrciComments(
-      octokit,
-      repo,
-      prNumber ? [parseInt(prNumber as string)] : []
-    );
-    res.status(200).json(failures);
+    if (prNumber === undefined) {
+      return res.status(403).end();
+    }
+    // Check if they exceed the rate limit
+    const userOctokit = await getOctokitWithUserToken(authorization as string);
+    const user = await userOctokit.rest.users.getAuthenticated();
+    if (await drCIRateLimitExceeded(user.data.login)) {
+      return res.status(429).end();
+    }
+    incrementDrCIRateLimit(user.data.login);
+  } else {
+    // No authorization provided, return 403
+    return res.status(403).end();
   }
 
-  res.status(403).end();
+  const { prNumber } = req.query;
+  const { repo }: UpdateCommentBody = req.body;
+  const octokit = await getOctokit(OWNER, repo);
+
+  const failures = await updateDrciComments(
+    octokit,
+    repo,
+    prNumber ? [parseInt(prNumber as string)] : []
+  );
+  res.status(200).json(failures);
 }
 
 export async function updateDrciComments(