Monorepo for Tangled
Something went wrong. Try again.
Go
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404package db
import ( "context" "database/sql" "fmt" "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) InsertCacheEntryWithinQuota(ctx context.Context, entry CacheEntry, maxEntries int64) (bool, error) { result, 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 ) select ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ? where ? <= 0 or ( select count(*) from cache_entries where owner_did = ? and state in ('pending', 'ready', 'deleting') ) < ?`, entry.ID, entry.StorageKey, entry.OwnerDID, entry.RepoDID, entry.Engine, entry.CacheKey, entry.CacheHash, entry.SizeBytes, entry.State, entry.CreatedAt.UnixNano(), entry.LastUsedAt.UnixNano(), maxEntries, entry.OwnerDID, maxEntries, ) if err != nil { return false, err } inserted, err := result.RowsAffected() return inserted == 1, err}
func (d *DB) MarkCacheEntryReady(ctx context.Context, id string, sizeBytes, maxBytes int64, now time.Time) ([]CacheEntry, error) { tx, err := d.BeginTx(ctx, nil) if err != nil { return nil, err } defer tx.Rollback() superseded, ready, err := markCacheEntryReady(ctx, tx, id, sizeBytes, maxBytes, now) if err != nil { return nil, err } if !ready { return nil, sql.ErrNoRows } if err := tx.Commit(); err != nil { return nil, err } return superseded, nil}
type cacheEntryTx interface { QueryRowContext(context.Context, string, ...any) *sql.Row QueryContext(context.Context, string, ...any) (*sql.Rows, error)}
func markCacheEntryReady(ctx context.Context, tx cacheEntryTx, id string, sizeBytes, maxBytes int64, now time.Time) ([]CacheEntry, bool, error) { var repoDID, engine, key, hash string err := tx.QueryRowContext(ctx, ` update cache_entries as current set state = 'ready', size_bytes = ?, last_used_at = ? where current.id = ? and current.state = 'pending' and (? <= 0 or ( select coalesce(sum(other.size_bytes), 0) from cache_entries as other where other.owner_did = current.owner_did and other.state in ('ready', 'deleting') and other.id <> current.id ) <= ? - ?) returning repo_did, engine, cache_key, cache_hash`, sizeBytes, now.UnixNano(), id, maxBytes, maxBytes, sizeBytes, ).Scan(&repoDID, &engine, &key, &hash) if err == sql.ErrNoRows { current, selectErr := scanCacheEntry(tx.QueryRowContext(ctx, `select `+cacheEntryColumns+` from cache_entries where id = ?`, id)) if selectErr == sql.ErrNoRows { return nil, false, nil } if selectErr != nil { return nil, false, selectErr } if current.State == "ready" { if current.SizeBytes != sizeBytes { return nil, false, fmt.Errorf("cache entry %q size changed from %d to %d", id, current.SizeBytes, sizeBytes) } return nil, true, nil } return nil, false, nil } if err != nil { return nil, false, 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, false, err } var superseded []CacheEntry for rows.Next() { entry, err := scanCacheEntry(rows) if err != nil { rows.Close() return nil, false, err } superseded = append(superseded, *entry) } if err := rows.Close(); err != nil { return nil, false, err } if err := rows.Err(); err != nil { return nil, false, err } return superseded, true, nil}
func (tx *EventBatchTx) MarkCacheEntryReady(ctx context.Context, id string, sizeBytes, maxBytes int64, now time.Time) ([]CacheEntry, bool, error) { return markCacheEntryReady(ctx, tx.tx, id, sizeBytes, maxBytes, now)}
func (tx *EventBatchTx) TouchCacheEntry(ctx context.Context, id string, now time.Time) error { _, err := tx.tx.ExecContext(ctx, ` update cache_entries set last_used_at = max(last_used_at, ?) where id = ? and state = 'ready'`, now.UnixNano(), id) return err}
func (tx *EventBatchTx) DiscardCacheEntry(ctx context.Context, id string, pendingOnly bool) (*CacheEntry, error) { stateFilter := "" if pendingOnly { stateFilter = " and state in ('pending', 'deleting')" } current, err := scanCacheEntry(tx.tx.QueryRowContext(ctx, ` update cache_entries set state = 'deleting' where id = ?`+stateFilter+` returning `+cacheEntryColumns, id)) if err == sql.ErrNoRows { return nil, nil } return current, err}
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) ClaimPendingCacheEntry(ctx context.Context, id string) (bool, error) { result, err := d.ExecContext(ctx, ` update cache_entries set state = 'deleting' where id = ? and state = 'pending'`, id) if err != nil { return false, err } changed, err := result.RowsAffected() return changed == 1, 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) { rows, err := d.QueryContext(ctx, ` select `+cacheEntryColumns+` from cache_entries where (state = 'ready' and last_used_at < ?) or (state in ('pending', 'deleting') and last_used_at < ?) order by last_used_at limit ?`, readyBefore.UnixNano(), pendingBefore.UnixNano(), limit, ) 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}
type MillCacheCapability struct { Action string CacheID string StorageKey string}
func (d *DB) SaveMillCacheCapabilities(leaseID string, capabilities []MillCacheCapability) error { tx, err := d.Begin() if err != nil { return err } defer tx.Rollback() if _, err := tx.Exec(`delete from mill_cache_capabilities where lease_id = ?`, leaseID); err != nil { return err } for _, capability := range capabilities { if capability.Action != "restore" && capability.Action != "save" { return fmt.Errorf("cache capability has invalid action %q", capability.Action) } if capability.CacheID == "" || capability.StorageKey == "" { return fmt.Errorf("cache capability has empty identity") } if _, err := tx.Exec(` insert into mill_cache_capabilities ( lease_id, action, cache_id, storage_key ) values (?, ?, ?, ?)`, leaseID, capability.Action, capability.CacheID, capability.StorageKey, ); err != nil { return err } } return tx.Commit()}
func (tx *EventBatchTx) ConsumeMillCacheCapability(ctx context.Context, leaseID, action, cacheID string) (string, bool, error) { var storageKey string err := tx.tx.QueryRowContext(ctx, ` delete from mill_cache_capabilities where lease_id = ? and action = ? and cache_id = ? returning storage_key`, leaseID, action, cacheID, ).Scan(&storageKey) if err == sql.ErrNoRows { return "", false, nil } if err != nil { return "", false, err } return storageKey, true, nil}
func (tx *EventBatchTx) QueueCacheObjectDeletion(ctx context.Context, storageKey string, now time.Time) error { _, err := tx.tx.ExecContext(ctx, ` insert into cache_object_deletions (storage_key, created_at) values (?, ?) on conflict(storage_key) do nothing`, storageKey, now.UnixNano(), ) return err}
func (d *DB) PendingCacheObjectDeletions(ctx context.Context, limit int) ([]string, error) { rows, err := d.QueryContext(ctx, ` select storage_key from cache_object_deletions order by created_at limit ?`, limit) if err != nil { return nil, err } defer rows.Close() var keys []string for rows.Next() { var key string if err := rows.Scan(&key); err != nil { return nil, err } keys = append(keys, key) } return keys, rows.Err()}
func (d *DB) CompleteCacheObjectDeletion(ctx context.Context, storageKey string) error { _, err := d.ExecContext(ctx, `delete from cache_object_deletions where storage_key = ?`, storageKey) return 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}