package db import ( "context" "sort" "time" ) type CacheEntry struct { ID string StorageKey string OwnerDID string RepoDID string Engine string CacheKey string CacheHash string SizeBytes int64 State string CreatedAt time.Time LastUsedAt time.Time } const cacheEntryColumns = ` id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash, size_bytes, state, created_at, last_used_at` func (d *DB) InsertCacheEntry(ctx context.Context, entry CacheEntry) error { _, err := d.ExecContext(ctx, ` insert into cache_entries ( id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash, size_bytes, state, created_at, last_used_at ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, entry.ID, entry.StorageKey, entry.OwnerDID, entry.RepoDID, entry.Engine, entry.CacheKey, entry.CacheHash, entry.SizeBytes, entry.State, entry.CreatedAt.UnixNano(), entry.LastUsedAt.UnixNano(), ) return err } func (d *DB) MarkCacheEntryReady(ctx context.Context, id string, sizeBytes int64, now time.Time) ([]CacheEntry, error) { tx, err := d.BeginTx(ctx, nil) if err != nil { return nil, err } defer tx.Rollback() var repoDID, engine, key, hash string if err := tx.QueryRowContext(ctx, ` update cache_entries set state = 'ready', size_bytes = ?, last_used_at = ? where id = ? and state = 'pending' returning repo_did, engine, cache_key, cache_hash`, sizeBytes, now.UnixNano(), id).Scan(&repoDID, &engine, &key, &hash); err != nil { return nil, err } rows, err := tx.QueryContext(ctx, ` update cache_entries set state = 'deleting' where repo_did = ? and engine = ? and cache_key = ? and cache_hash = ? and state = 'ready' and id <> ? returning `+cacheEntryColumns, repoDID, engine, key, hash, id) if err != nil { return nil, err } var superseded []CacheEntry for rows.Next() { entry, err := scanCacheEntry(rows) if err != nil { rows.Close() return nil, err } superseded = append(superseded, *entry) } if err := rows.Close(); err != nil { return nil, err } if err := rows.Err(); err != nil { return nil, err } sort.Slice(superseded, func(i, j int) bool { return superseded[i].CreatedAt.After(superseded[j].CreatedAt) }) if err := tx.Commit(); err != nil { return nil, err } return superseded, nil } func (d *DB) FindCacheEntry(ctx context.Context, repoDID, engine, key, hash string) (*CacheEntry, error) { return scanCacheEntry(d.QueryRowContext(ctx, ` select `+cacheEntryColumns+` from cache_entries where repo_did = ? and engine = ? and cache_key = ? and cache_hash = ? and state = 'ready' order by created_at desc limit 1`, repoDID, engine, key, hash)) } func (d *DB) FindFallbackCacheEntry(ctx context.Context, repoDID, engine, key, excludeHash string) (*CacheEntry, error) { return scanCacheEntry(d.QueryRowContext(ctx, ` select `+cacheEntryColumns+` from cache_entries where repo_did = ? and engine = ? and cache_key = ? and cache_hash <> ? and cache_hash <> '' and state = 'ready' order by created_at desc limit 1`, repoDID, engine, key, excludeHash)) } func (d *DB) TouchCacheEntry(ctx context.Context, id string, now time.Time) error { _, err := d.ExecContext(ctx, ` update cache_entries set last_used_at = ? where id = ? and state = 'ready'`, now.UnixNano(), id) return err } func (d *DB) ClaimCacheEntry(ctx context.Context, id, expectedState string, expectedLastUsed time.Time) (bool, error) { result, err := d.ExecContext(ctx, ` update cache_entries set state = 'deleting' where id = ? and state = ? and last_used_at = ?`, id, expectedState, expectedLastUsed.UnixNano()) if err != nil { return false, err } changed, err := result.RowsAffected() return changed == 1, err } func (d *DB) RestoreCacheEntryState(ctx context.Context, id, state string) error { _, err := d.ExecContext(ctx, ` update cache_entries set state = ? where id = ? and state = 'deleting'`, state, id) return err } func (d *DB) ExpiredCacheEntries(ctx context.Context, readyBefore, pendingBefore time.Time, limit int) ([]CacheEntry, error) { ready, err := d.queryCacheEntries(ctx, ` select `+cacheEntryColumns+` from cache_entries where state = 'ready' and last_used_at < ? order by last_used_at limit ?`, readyBefore.UnixNano(), limit) if err != nil { return nil, err } recovery, err := d.queryCacheEntries(ctx, ` select `+cacheEntryColumns+` from cache_entries where state in ('pending', 'deleting') and created_at < ? order by created_at limit ?`, pendingBefore.UnixNano(), limit) if err != nil { return nil, err } entries := append(ready, recovery...) sort.Slice(entries, func(i, j int) bool { left, right := entries[i].CreatedAt, entries[j].CreatedAt if entries[i].State == "ready" { left = entries[i].LastUsedAt } if entries[j].State == "ready" { right = entries[j].LastUsedAt } return left.Before(right) }) if len(entries) > limit { entries = entries[:limit] } return entries, nil } func (d *DB) queryCacheEntries(ctx context.Context, query string, args ...any) ([]CacheEntry, error) { rows, err := d.QueryContext(ctx, query, args...) if err != nil { return nil, err } defer rows.Close() var entries []CacheEntry for rows.Next() { entry, err := scanCacheEntry(rows) if err != nil { return nil, err } entries = append(entries, *entry) } return entries, rows.Err() } func (d *DB) DeleteCacheEntry(ctx context.Context, id string) error { _, err := d.ExecContext(ctx, `delete from cache_entries where id = ?`, id) return err } func (d *DB) CacheUsageByOwner(ctx context.Context, ownerDID string) (bytes, count int64, err error) { err = d.QueryRowContext(ctx, ` select coalesce(sum(size_bytes), 0), count(*) from cache_entries where owner_did = ? and state in ('ready', 'deleting')`, ownerDID).Scan(&bytes, &count) return bytes, count, err } type cacheEntryScanner interface { Scan(dest ...any) error } func scanCacheEntry(row cacheEntryScanner) (*CacheEntry, error) { var entry CacheEntry var createdAt, lastUsedAt int64 if err := row.Scan( &entry.ID, &entry.StorageKey, &entry.OwnerDID, &entry.RepoDID, &entry.Engine, &entry.CacheKey, &entry.CacheHash, &entry.SizeBytes, &entry.State, &createdAt, &lastUsedAt, ); err != nil { return nil, err } entry.CreatedAt = time.Unix(0, createdAt) entry.LastUsedAt = time.Unix(0, lastUsedAt) return &entry, nil }