Fix rapid cache cleanup bugs: always run both index and bucket passes, check Redis index during bucket scan, evict prefix from Redis on bulk deletes

Author: cursoragent

packages/orchestrator/cmd/clean-rapid-cache/main.go

@@ -71,15 +71,21 @@ func clean(ctx context.Context, bucket string, prefix string, cutoff time.Time,
 		_ = client.Close()
 	}()
 
-	if deleted, err := cleanFromIndex(ctx, client, bucket, cutoff, maxDeletions, dryRun, index); err != nil {
+	deletedIndex, err := cleanFromIndex(ctx, client, bucket, cutoff, maxDeletions, dryRun, index)
+	if err != nil {
 		return err
-	} else if deleted > 0 {
-		log.Printf("summary dry_run=%t deleted=%d source=redis", dryRun, deleted)
+	}
+
+	if deletedIndex > 0 {
+		log.Printf("summary dry_run=%t deleted=%d source=redis", dryRun, deletedIndex)
+	}
 
+	remaining := maxDeletions - deletedIndex
+	if remaining <= 0 {
 		return nil
 	}
 
-	return cleanFromBucket(ctx, client, bucket, prefix, cutoff, maxDeletions, dryRun)
+	return cleanFromBucket(ctx, client, bucket, prefix, cutoff, remaining, dryRun, index)
 }
 
 func cleanFromIndex(ctx context.Context, client *gcs.Client, bucket string, cutoff time.Time, maxDeletions int, dryRun bool, index storage.RapidCacheIndex) (int, error) {
@@ -117,7 +123,7 @@ func cleanFromIndex(ctx context.Context, client *gcs.Client, bucket string, cuto
 	return deleted, nil
 }
 
-func cleanFromBucket(ctx context.Context, client *gcs.Client, bucket string, prefix string, cutoff time.Time, maxDeletions int, dryRun bool) error {
+func cleanFromBucket(ctx context.Context, client *gcs.Client, bucket string, prefix string, cutoff time.Time, maxDeletions int, dryRun bool, index storage.RapidCacheIndex) error {
 	var scanned, matched, deleted int
 	objects := client.Bucket(bucket).Objects(ctx, &gcs.Query{Prefix: prefix})
 	for {
@@ -133,6 +139,13 @@ func cleanFromBucket(ctx context.Context, client *gcs.Client, bucket string, pre
 		if !attrs.Updated.Before(cutoff) {
 			continue
 		}
+
+		if !dryRun {
+			if recent, err := index.IsRecent(ctx, attrs.Name); err == nil && recent {
+				continue
+			}
+		}
+
 		matched++
 		if deleted >= maxDeletions {
 			break
@@ -145,6 +158,7 @@ func cleanFromBucket(ctx context.Context, client *gcs.Client, bucket string, pre
 		if err := client.Bucket(bucket).Object(attrs.Name).Delete(ctx); err != nil {
 			return fmt.Errorf("delete cache object: %w", err)
 		}
+		_ = index.Evict(ctx, attrs.Name, attrs.Size)
 		deleted++
 	}
 

packages/shared/pkg/storage/storage_cache_rapid.go

@@ -27,6 +27,7 @@ func WrapInRapidBucketCache(_ context.Context, cache StorageProvider, inner Stor
 }
 
 func (p *rapidCacheProvider) DeleteObjectsWithPrefix(ctx context.Context, prefix string) error {
+	_ = p.index.EvictPrefix(ctx, prefix)
 	_ = p.cache.DeleteObjectsWithPrefix(ctx, "rapid-cache/"+prefix)
 
 	return p.inner.DeleteObjectsWithPrefix(ctx, prefix)

packages/shared/pkg/storage/storage_cache_rapid_index.go

@@ -16,6 +16,8 @@ type RapidCacheIndex interface {
 	Admit(ctx context.Context, path string, size int64) error
 	Evict(ctx context.Context, path string, size int64) error
 	Candidates(ctx context.Context, before time.Time, limit int64) ([]string, error)
+	IsRecent(ctx context.Context, path string) (bool, error)
+	EvictPrefix(ctx context.Context, prefix string) error
 }
 
 type noopRapidCacheIndex struct{}
@@ -28,6 +30,8 @@ func (noopRapidCacheIndex) Evict(context.Context, string, int64) error { return
 func (noopRapidCacheIndex) Candidates(context.Context, time.Time, int64) ([]string, error) {
 	return nil, nil
 }
+func (noopRapidCacheIndex) IsRecent(context.Context, string) (bool, error) { return false, nil }
+func (noopRapidCacheIndex) EvictPrefix(context.Context, string) error      { return nil }
 
 type redisRapidCacheIndex struct {
 	redis redis.UniversalClient
@@ -110,6 +114,61 @@ func (i *redisRapidCacheIndex) Candidates(ctx context.Context, before time.Time,
 	}).Result()
 }
 
+func (i *redisRapidCacheIndex) IsRecent(ctx context.Context, path string) (bool, error) {
+	score := i.redis.ZScore(ctx, i.keys.chunks, path)
+	if score.Err() == redis.Nil {
+		return false, nil
+	}
+	if score.Err() != nil {
+		return false, score.Err()
+	}
+
+	return true, nil
+}
+
+func (i *redisRapidCacheIndex) EvictPrefix(ctx context.Context, prefix string) error {
+	script := redis.NewScript(`
+local cursor = "0"
+local prefix = ARGV[1]
+local chunks_key = KEYS[1]
+local build_bytes_key = KEYS[2]
+local build_chunks_key = KEYS[3]
+
+repeat
+  local scan_result = redis.call('ZSCAN', chunks_key, cursor, 'MATCH', prefix .. '*', 'COUNT', 100)
+  cursor = scan_result[1]
+  local members = scan_result[2]
+  
+  for i = 1, #members, 2 do
+    local path = members[i]
+    redis.call('ZREM', chunks_key, path)
+    
+    local path_without_prefix = string.gsub(path, '^rapid%-cache/', '')
+    local build_id = string.match(path_without_prefix, '^([^/]+)')
+    if build_id then
+      local obj_result = redis.call('HGET', build_bytes_key, build_id)
+      if obj_result then
+        redis.call('HINCRBY', build_chunks_key, build_id, -1)
+        local chunks_left = redis.call('HGET', build_chunks_key, build_id)
+        if tonumber(chunks_left) <= 0 then
+          redis.call('HDEL', build_bytes_key, build_id)
+          redis.call('HDEL', build_chunks_key, build_id)
+        end
+      end
+    end
+  end
+until cursor == "0"
+
+return 0
+`)
+
+	return script.Run(ctx, i.redis, []string{
+		i.keys.chunks,
+		i.keys.buildBytes,
+		i.keys.buildChunks,
+	}, "rapid-cache/"+prefix).Err()
+}
+
 func rapidCacheBuildID(path string) string {
 	path = strings.TrimPrefix(path, rapidCachePrefix)
 	buildID, _ := SplitPath(path)