Add ability to remove keys by type or API call

Fixes falcondev-oss#79
Partially related to falcondev-oss#74

Signed-off-by: Oleksandr Porunov <[email protected]>

Author: porunov

docs/content/1.getting-started/1.index.md

@@ -89,6 +89,38 @@ The port the server should listen on.
 
 The directory to use for temporary files.
 
+#### `ENABLE_TYPED_KEY_PREFIX_REMOVAL`
+
+- Default: `false`
+
+If enabled then any specifically structured keys will be automatically removed if there are more keys of such type than 
+the amount specified via `MAX_STORED_KEYS_PER_TYPE`.
+Key type is defined by the following format:
+`{key_type}{TYPED_KEY_DELIMITER}{optional custom string}`
+
+Usually, it's recommended to use the following key structure:
+`{repository_name}_{branch}_{unique_key_name}{TYPED_KEY_DELIMITER}{hash or other unique data}`
+
+An example could look like:
+`falcondev-oss/github-actions-cache-server_dev_docker-cache@#@16ee33c5f9f59c0`
+
+In such case older keys with prefix `falcondev-oss/github-actions-cache-server_dev_docker-cache` will be automatically 
+removed whenever there are more records for this key type then the one specified via `MAX_STORED_KEYS_PER_TYPE`.
+
+#### `MAX_STORED_KEYS_PER_TYPE`
+
+- Default: `3`
+
+If `ENABLE_TYPED_KEY_PREFIX_REMOVAL` is `true` then this is the maximum amount of the most recent keys to keep per key type. 
+Any other older keys will be automatically removed.
+
+#### `TYPED_KEY_DELIMITER`
+
+- Default: `@#@`
+
+If `ENABLE_TYPED_KEY_PREFIX_REMOVAL` is `true` then this is the delimiter which is used to identify the key type. 
+Any prefix before `@#@` is considered to be the key type.
+
 ## 2. Setup with Self-Hosted Runners
 
 Set the `ACTIONS_RESULTS_URL` on your runner to the API URL (with a trailing slash).
@@ -148,3 +180,22 @@ variant: subtle
 There is no need to change any of your workflows! 🔥
 
 If you've set up your self-hosted runners correctly, they will automatically use the cache server for caching.
+
+## 4. Cache cleanup
+
+Cache cleanup can be triggered automatically via configured eviction policies like:
+- `CLEANUP_OLDER_THAN_DAYS` - eviction by time.
+- `ENABLE_TYPED_KEY_PREFIX_REMOVAL` - eviction by key type (key prefix).
+
+Moreover, an additional API exists which can be used to evict key records by using specified key prefix.
+API path: `/<token>/extra/clear_key_prefix`. API method: `POST`.
+Request body: `{ keyPrefix: string }`.
+
+Any keys with the prefix of `keyPrefix` will be removed and data will be cleared for such keys.
+Usually, it's useful to call keys removal for an archived branch or a closed PR where you know that the cache won't 
+be reused anymore. For such cases you could have a step which is called and clears all keys prefixed with the branch name. 
+For example: 
+```
+- name: Remove keys by prefix
+  run: curl --header "Content-Type: application/json" --request POST --data '{"keyPrefix":"${{ github.head_ref || github.ref_name }}"}' "${ACTIONS_CACHE_URL}extra/clear_key_prefix"
+```

lib/db/index.ts

@@ -221,6 +221,31 @@ export async function findStaleKeys(
     .execute()
 }
 
+export async function findPrefixedKeysForRemoval(
+  db: DB,
+  { keyPrefix, skipRecentKeysLimit }: { keyPrefix: string; skipRecentKeysLimit?: number; },
+) {
+
+  let query = db
+    .selectFrom('cache_keys')
+    .where('key', 'like', `${keyPrefix}%`)
+
+  if (skipRecentKeysLimit && skipRecentKeysLimit > 0){
+    query = query.where(({ eb, selectFrom, not }) => not(eb(
+      'id',
+      'in',
+      selectFrom('cache_keys')
+        .select('cache_keys.id')
+        .orderBy('cache_keys.accessed_at desc')
+        .limit(skipRecentKeysLimit)
+    )))
+  }
+
+  return query
+    .selectAll()
+    .execute()
+}
+
 export async function createKey(
   db: DB,
   { key, version, date }: { key: string; version: string; date?: Date },

lib/env.ts

@@ -15,6 +15,9 @@ const envSchema = z.object({
   DEBUG: booleanSchema.default('false'),
   NITRO_PORT: portSchema.default(3000),
   TEMP_DIR: z.string().default(tmpdir()),
+  ENABLE_TYPED_KEY_PREFIX_REMOVAL: booleanSchema.default('false'),
+  MAX_STORED_KEYS_PER_TYPE: z.coerce.number().int().min(0).default(3),
+  TYPED_KEY_DELIMITER: z.string().default('@#@'),
 })
 
 const parsedEnv = envSchema.safeParse(process.env)

lib/storage/index.ts

@@ -13,6 +13,7 @@ import {
   touchKey,
   updateOrCreateKey,
   useDB,
+  findPrefixedKeysForRemoval,
 } from '~/lib/db'
 import { ENV } from '~/lib/env'
 import { logger } from '~/lib/logger'
@@ -21,6 +22,7 @@ import { getStorageDriver } from '~/lib/storage/drivers'
 import { getObjectNameFromKey } from '~/lib/utils'
 
 export interface Storage {
+  pruneCacheByKeyPrefix: (keyPrefix: string) => Promise<void>
   getCacheEntry: (
     keys: string[],
     version: string,
@@ -68,9 +70,44 @@ export async function initializeStorage() {
       const driver = await driverSetup()
       const db = await useDB()
 
-      storage = {
+    storage = {
+
+      async pruneCacheKeys(keysForRemoval){
+        if (keysForRemoval.length === 0) {
+          logger.debug('Prune: No caches to prune')
+          return
+        }
+
+        await driver.delete({
+          objectNames: keysForRemoval.map((key) => getObjectNameFromKey(key.key, key.version)),
+        })
+
+        await pruneKeys(db, keysForRemoval)
+      },
+
+      async pruneStaleCacheByType(key){
+        if (!ENV.ENABLE_TYPED_KEY_PREFIX_REMOVAL){
+          return
+        }
+        let keyTypeIndex = key.indexOf(ENV.TYPED_KEY_DELIMITER)
+        if (keyTypeIndex < 1){
+          return
+        }
+        let keyType = key.substring(0, keyTypeIndex)
+        logger.debug(`Prune by type is called: Type [${keyType}]. Full key [${key}].`)
+        let keysForRemoval = await findPrefixedKeysForRemoval(db, { keyPrefix: keyType, skipRecentKeysLimit: ENV.MAX_STORED_KEYS_PER_TYPE } )
+        logger.debug(`Removing ${keysForRemoval.length} keys for prefix [${keyType}].`)
+        await this.pruneCacheKeys(keysForRemoval)
+      },
+
+      async pruneCacheByKeyPrefix(keyPrefix){
+        let keysForRemoval = await findPrefixedKeysForRemoval(db, { keyPrefix: keyPrefix, skipRecentKeysLimit: 0 } )
+        logger.debug(`Removing ${keysForRemoval.length} keys for prefix [${keyPrefix}].`)
+        await this.pruneCacheKeys(keysForRemoval)
+      },
+
         async reserveCache(key, version, totalSize) {
-          logger.debug('Reserve:', { key, version })
+            logger.debug('Reserve:', { key, version })
 
           if (await getUpload(db, { key, version })) {
             logger.debug(`Reserve: Already reserved. Ignoring...`, { key, version })
@@ -103,6 +140,8 @@ export async function initializeStorage() {
             uploadId,
           })
 
+          this.pruneStaleCacheByType(key)
+
           return {
             cacheId: uploadId,
           }

routes/[token]/extra/clear_key_prefix.post.ts

@@ -0,0 +1,25 @@
+import { z } from 'zod'
+
+import { useStorageAdapter } from '~/lib/storage'
+
+const queryParamSchema = z.object({
+  keyPrefix: z.string().min(1),
+})
+
+export default defineEventHandler(async (event) => {
+  const parsedQuery = queryParamSchema.safeParse(getQuery(event))
+  if (!parsedQuery.success)
+    throw createError({
+      statusCode: 400,
+      statusMessage: `Invalid query parameters: ${parsedQuery.error.message}`,
+    })
+
+  const { keyPrefix } = parsedQuery.data
+
+  const adapter = await useStorageAdapter()
+  await adapter.pruneCacheByKeyPrefix(keyPrefix)
+
+  return {
+    ok: true,
+  }
+})