caching.ts

import * as fs from "fs";

import * as actionsCache from "@actions/cache";
import * as semver from "semver";

import {
  getRequiredInput,
  getWorkflowRunAttempt,
  getWorkflowRunID,
} from "../actions-util";
import { getAutomationID, listActionsCaches } from "../api-client";
import { createCacheKeyHash } from "../caching-utils";
import { type CodeQL } from "../codeql";
import { type Config } from "../config-utils";
import { getCommitOid } from "../git-utils";
import { type Language, parseBuiltInLanguage } from "../languages";
import { type Logger, withGroupAsync } from "../logging";
import {
  CleanupLevel,
  getBaseDatabaseOidsFilePath,
  getCodeQLDatabasePath,
  getErrorMessage,
  isInTestMode,
  tryGetFolderBytes,
  waitForResultWithTimeLimit,
} from "../util";

import { OverlayDatabaseMode } from "./overlay-database-mode";

/**
 * The maximum (uncompressed) size of the overlay base database that we will
 * upload. By default, the Actions Cache has an overall capacity of 10 GB, and
 * the Actions Cache client library uses zstd compression.
 *
 * Ideally we would apply a size limit to the compressed overlay-base database,
 * but we cannot do so because compression is handled transparently by the
 * Actions Cache client library. Instead we place a limit on the uncompressed
 * size of the overlay-base database.
 *
 * Assuming 2.5:1 compression ratio, the 7.5 GB limit on uncompressed data would
 * translate to a limit of around 3 GB after compression.
 */
const OVERLAY_BASE_DATABASE_MAX_UPLOAD_SIZE_MB = 7500;
const OVERLAY_BASE_DATABASE_MAX_UPLOAD_SIZE_BYTES =
  OVERLAY_BASE_DATABASE_MAX_UPLOAD_SIZE_MB * 1_000_000;

// Constants for database caching
const CACHE_VERSION = 1;
const CACHE_PREFIX = "codeql-overlay-base-database";

// The purpose of this ten-minute limit is to guard against the possibility
// that the cache service is unresponsive, which would otherwise cause the
// entire action to hang.  Normally we expect cache operations to complete
// within two minutes.
const MAX_CACHE_OPERATION_MS = 600_000;

/**
 * Checks that the overlay-base database is valid by checking for the
 * existence of the base database OIDs file.
 *
 * @param config The configuration object
 * @param logger The logger instance
 * @param warningPrefix Prefix for the check failure warning message
 * @returns True if the verification succeeded, false otherwise
 */
async function checkOverlayBaseDatabase(
  codeql: CodeQL,
  config: Config,
  logger: Logger,
  warningPrefix: string,
): Promise<boolean> {
  // An overlay-base database should contain the base database OIDs file.
  const baseDatabaseOidsFilePath = getBaseDatabaseOidsFilePath(config);
  if (!fs.existsSync(baseDatabaseOidsFilePath)) {
    logger.warning(
      `${warningPrefix}: ${baseDatabaseOidsFilePath} does not exist`,
    );
    return false;
  }

  for (const language of config.languages) {
    const dbPath = getCodeQLDatabasePath(config, language);
    try {
      const resolveDatabaseOutput = await codeql.resolveDatabase(dbPath);
      if (
        resolveDatabaseOutput === undefined ||
        !("overlayBaseSpecifier" in resolveDatabaseOutput)
      ) {
        logger.info(`${warningPrefix}: no overlayBaseSpecifier defined`);
        return false;
      } else {
        logger.debug(
          `Overlay base specifier for ${language} overlay-base database found: ` +
            `${resolveDatabaseOutput.overlayBaseSpecifier}`,
        );
      }
    } catch (e) {
      logger.warning(`${warningPrefix}: failed to resolve database: ${e}`);
      return false;
    }
  }

  return true;
}

/**
 * Uploads the overlay-base database to the GitHub Actions cache. If conditions
 * for uploading are not met, the function does nothing and returns false.
 *
 * This function uses the `checkout_path` input to determine the repository path
 * and works only when called from `analyze` or `upload-sarif`.
 *
 * @param codeql The CodeQL instance
 * @param config The configuration object
 * @param logger The logger instance
 * @returns A promise that resolves to true if the upload was performed and
 * successfully completed, or false otherwise
 */
export async function cleanupAndUploadOverlayBaseDatabaseToCache(
  codeql: CodeQL,
  config: Config,
  logger: Logger,
): Promise<boolean> {
  const overlayDatabaseMode = config.overlayDatabaseMode;
  if (overlayDatabaseMode !== OverlayDatabaseMode.OverlayBase) {
    logger.debug(
      `Overlay database mode is ${overlayDatabaseMode}. ` +
        "Skip uploading overlay-base database to cache.",
    );
    return false;
  }
  if (!config.useOverlayDatabaseCaching) {
    logger.debug(
      "Overlay database caching is disabled. " +
        "Skip uploading overlay-base database to cache.",
    );
    return false;
  }
  if (isInTestMode()) {
    logger.debug(
      "In test mode. Skip uploading overlay-base database to cache.",
    );
    return false;
  }

  const databaseIsValid = await checkOverlayBaseDatabase(
    codeql,
    config,
    logger,
    "Abort uploading overlay-base database to cache",
  );
  if (!databaseIsValid) {
    return false;
  }

  // Clean up the database using the overlay cleanup level.
  await withGroupAsync("Cleaning up databases", async () => {
    await codeql.databaseCleanupCluster(config, CleanupLevel.Overlay);
  });

  const dbLocation = config.dbLocation;

  const databaseSizeBytes = await tryGetFolderBytes(dbLocation, logger);
  if (databaseSizeBytes === undefined) {
    logger.warning(
      "Failed to determine database size. " +
        "Skip uploading overlay-base database to cache.",
    );
    return false;
  }

  if (databaseSizeBytes > OVERLAY_BASE_DATABASE_MAX_UPLOAD_SIZE_BYTES) {
    const databaseSizeMB = Math.round(databaseSizeBytes / 1_000_000);
    logger.warning(
      `Database size (${databaseSizeMB} MB) ` +
        `exceeds maximum upload size (${OVERLAY_BASE_DATABASE_MAX_UPLOAD_SIZE_MB} MB). ` +
        "Skip uploading overlay-base database to cache.",
    );
    return false;
  }

  const codeQlVersion = (await codeql.getVersion()).version;
  const checkoutPath = getRequiredInput("checkout_path");
  const cacheSaveKey = await getCacheSaveKey(
    config,
    codeQlVersion,
    checkoutPath,
    logger,
  );
  logger.info(
    `Uploading overlay-base database to Actions cache with key ${cacheSaveKey}`,
  );

  try {
    const cacheId = await waitForResultWithTimeLimit(
      MAX_CACHE_OPERATION_MS,
      actionsCache.saveCache([dbLocation], cacheSaveKey),
      () => {},
    );
    if (cacheId === undefined) {
      logger.warning("Timed out while uploading overlay-base database");
      return false;
    }
  } catch (error) {
    logger.warning(
      "Failed to upload overlay-base database to cache: " +
        `${error instanceof Error ? error.message : String(error)}`,
    );
    return false;
  }
  logger.info(`Successfully uploaded overlay-base database from ${dbLocation}`);
  return true;
}

export interface OverlayBaseDatabaseDownloadStats {
  databaseSizeBytes: number;
  databaseDownloadDurationMs: number;
}

/**
 * Downloads the overlay-base database from the GitHub Actions cache. If conditions
 * for downloading are not met, the function does nothing and returns false.
 *
 * @param codeql The CodeQL instance
 * @param config The configuration object
 * @param logger The logger instance
 * @returns A promise that resolves to download statistics if an overlay-base
 * database was successfully downloaded, or undefined if the download was
 * either not performed or failed.
 */
export async function downloadOverlayBaseDatabaseFromCache(
  codeql: CodeQL,
  config: Config,
  logger: Logger,
): Promise<OverlayBaseDatabaseDownloadStats | undefined> {
  const overlayDatabaseMode = config.overlayDatabaseMode;
  if (overlayDatabaseMode !== OverlayDatabaseMode.Overlay) {
    logger.debug(
      `Overlay database mode is ${overlayDatabaseMode}. ` +
        "Skip downloading overlay-base database from cache.",
    );
    return undefined;
  }
  if (!config.useOverlayDatabaseCaching) {
    logger.debug(
      "Overlay database caching is disabled. " +
        "Skip downloading overlay-base database from cache.",
    );
    return undefined;
  }
  if (isInTestMode()) {
    logger.debug(
      "In test mode. Skip downloading overlay-base database from cache.",
    );
    return undefined;
  }

  const dbLocation = config.dbLocation;
  const codeQlVersion = (await codeql.getVersion()).version;
  const cacheRestoreKeyPrefix = await getCacheRestoreKeyPrefix(
    config,
    codeQlVersion,
  );

  logger.info(
    "Looking in Actions cache for overlay-base database with " +
      `restore key ${cacheRestoreKeyPrefix}`,
  );

  let databaseDownloadDurationMs = 0;
  try {
    const databaseDownloadStart = performance.now();
    const foundKey = await waitForResultWithTimeLimit(
      // This ten-minute limit for the cache restore operation is mainly to
      // guard against the possibility that the cache service is unresponsive
      // and hangs outside the data download.
      //
      // Data download (which is normally the most time-consuming part of the
      // restore operation) should not run long enough to hit this limit. Even
      // for an extremely large 10GB database, at a download speed of 40MB/s
      // (see below), the download should complete within five minutes. If we
      // do hit this limit, there are likely more serious problems other than
      // mere slow download speed.
      //
      // This is important because we don't want any ongoing file operations
      // on the database directory when we do hit this limit. Hitting this
      // time limit takes us to a fallback path where we re-initialize the
      // database from scratch at dbLocation, and having the cache restore
      // operation continue to write into dbLocation in the background would
      // really mess things up. We want to hit this limit only in the case
      // of a hung cache service, not just slow download speed.
      MAX_CACHE_OPERATION_MS,
      actionsCache.restoreCache(
        [dbLocation],
        cacheRestoreKeyPrefix,
        undefined,
        {
          // Azure SDK download (which is the default) uses 128MB segments; see
          // https://github.com/actions/toolkit/blob/main/packages/cache/README.md.
          // Setting segmentTimeoutInMs to 3000 translates to segment download
          // speed of about 40 MB/s, which should be achievable unless the
          // download is unreliable (in which case we do want to abort).
          segmentTimeoutInMs: 3000,
        },
      ),
      () => {
        logger.info("Timed out downloading overlay-base database from cache");
      },
    );
    databaseDownloadDurationMs = Math.round(
      performance.now() - databaseDownloadStart,
    );

    if (foundKey === undefined) {
      logger.info("No overlay-base database found in Actions cache");
      return undefined;
    }

    logger.info(
      `Downloaded overlay-base database in cache with key ${foundKey}`,
    );
  } catch (error) {
    logger.warning(
      "Failed to download overlay-base database from cache: " +
        `${error instanceof Error ? error.message : String(error)}`,
    );
    return undefined;
  }

  const databaseIsValid = await checkOverlayBaseDatabase(
    codeql,
    config,
    logger,
    "Downloaded overlay-base database is invalid",
  );
  if (!databaseIsValid) {
    logger.warning("Downloaded overlay-base database failed validation");
    return undefined;
  }

  const databaseSizeBytes = await tryGetFolderBytes(dbLocation, logger);
  if (databaseSizeBytes === undefined) {
    logger.info(
      "Filesystem error while accessing downloaded overlay-base database",
    );
    // The problem that warrants reporting download failure is not that we are
    // unable to determine the size of the database. Rather, it is that we
    // encountered a filesystem error while accessing the database, which
    // indicates that an overlay analysis will likely fail.
    return undefined;
  }

  logger.info(`Successfully downloaded overlay-base database to ${dbLocation}`);
  return {
    databaseSizeBytes: Math.round(databaseSizeBytes),
    databaseDownloadDurationMs,
  };
}

/**
 * Computes the cache key for saving the overlay-base database to the GitHub
 * Actions cache.
 *
 * The key consists of the restore key prefix (which does not include the
 * commit SHA) and the commit SHA of the current checkout.
 */
export async function getCacheSaveKey(
  config: Config,
  codeQlVersion: string,
  checkoutPath: string,
  logger: Logger,
): Promise<string> {
  let runId = 1;
  let attemptId = 1;
  try {
    runId = getWorkflowRunID();
    attemptId = getWorkflowRunAttempt();
  } catch (e) {
    logger.warning(
      `Failed to get workflow run ID or attempt ID. Reason: ${getErrorMessage(e)}`,
    );
  }
  const sha = await getCommitOid(checkoutPath);
  const restoreKeyPrefix = await getCacheRestoreKeyPrefix(
    config,
    codeQlVersion,
  );
  return `${restoreKeyPrefix}${sha}-${runId}-${attemptId}`;
}

/**
 * Computes the cache key prefix for restoring the overlay-base database from
 * the GitHub Actions cache.
 *
 * Actions cache supports using multiple restore keys to indicate preference,
 * and this function could in principle take advantage of that feature by
 * returning a list of restore key prefixes. However, since overlay-base
 * databases are built from the default branch and used in PR analysis, it is
 * exceedingly unlikely that the commit SHA will ever be the same.
 *
 * Therefore, this function returns only a single restore key prefix, which does
 * not include the commit SHA. This allows us to restore the most recent
 * compatible overlay-base database.
 */
export async function getCacheRestoreKeyPrefix(
  config: Config,
  codeQlVersion: string,
): Promise<string> {
  return `${await getCacheKeyPrefixBase(config.languages)}${codeQlVersion}-`;
}

/**
 * Computes the cache key prefix for overlay-base databases, excluding the
 * CodeQL version.
 */
async function getCacheKeyPrefixBase(
  parsedLanguages: Language[],
): Promise<string> {
  const languagesComponent = [...parsedLanguages].sort().join("_");

  const cacheKeyComponents = {
    automationID: await getAutomationID(),
    // Add more components here as needed in the future
  };
  const componentsHash = createCacheKeyHash(cacheKeyComponents);

  // CACHE_PREFIX: distinguishes overlay-base databases from other cache objects
  // CACHE_VERSION: cache format version
  // componentsHash: hash of additional components (see above for details)
  // languagesComponent: the languages included in the overlay-base database
  //
  // Technically we can also include languages in the componentsHash, but
  // including them explicitly in the cache key makes it easier to debug and
  // understand the cache key structure.
  return `${CACHE_PREFIX}-${CACHE_VERSION}-${componentsHash}-${languagesComponent}-`;
}

/**
 * Searches the GitHub Actions cache for overlay-base databases matching the given languages, and
 * returns all stable CodeQL versions found across matching cache entries.
 *
 * @returns Unique stable CodeQL versions found in cached overlay-base databases, sorted from latest to
 * earliest, or undefined if one of the languages is not a built-in language.
 */
export async function getCodeQlVersionsForOverlayBaseDatabases(
  rawLanguages: string[],
  logger: Logger,
): Promise<string[] | undefined> {
  const languages = rawLanguages.map(parseBuiltInLanguage);
  if (languages.includes(undefined)) {
    logger.warning(
      "One or more provided languages are not recognized as built-in languages. " +
        "Skipping searching for overlay-base databases in cache.",
    );
    return undefined;
  }
  const cacheKeyPrefix = await getCacheKeyPrefixBase(
    languages.filter((l) => l !== undefined),
  );

  logger.debug(
    `Searching for overlay-base databases in Actions cache with ` +
      `prefix ${cacheKeyPrefix}`,
  );

  const caches = await listActionsCaches(cacheKeyPrefix);

  if (caches.length === 0) {
    logger.info("No overlay-base databases found in Actions cache.");
    return [];
  }

  logger.info(
    `Found ${caches.length} overlay-base ` +
      `${caches.length === 1 ? "database" : "databases"} in the Actions cache.`,
  );

  // Parse CodeQL versions from cache keys, matching only stable releases.
  //
  // After the prefix, the remaining key format starts with `${codeQlVersion}-`. Nightlies will have
  // a suffix like `+202604201548` that will break the match.
  //
  // Caveat: this relies on the fact that we haven't released any CodeQL bundles with the
  // `x.y.z-<pre-release>` semver format which does not interact well with the current overlay base
  // DB cache key format.
  const versionRegex = /^([\d.]+)-/;
  const versionSet = new Set<string>();

  for (const cache of caches) {
    if (!cache.key) continue;
    const suffix = cache.key.substring(cacheKeyPrefix.length);
    const match = suffix.match(versionRegex);
    if (match && semver.valid(match[1])) {
      versionSet.add(match[1]);
    }
  }

  if (versionSet.size === 0) {
    logger.info(
      "Could not parse any CodeQL versions from overlay-base database " +
        "cache keys.",
    );
    return [];
  }

  const versions = [...versionSet].sort(semver.rcompare);

  logger.info(
    `Found overlay databases for the following CodeQL versions in the Actions cache: ${versions.join(", ")}`,
  );

  return versions;
}