package db
import (
"context"
"database/sql"
"encoding/json"
"slices"
"strings"
"time"
"tangled.org/core/spindle/quota"
)
var _ quota.NixUseRecorder = (*QuotaStore)(nil)
const (
nixUseBatch = 500
nixDropBatch = 500
)
type NixObjectRow struct {
Cache, StoreHash, RepoDID, OwnerDID string
NarSize int64
LastUsedAt time.Time
Gen int64
}
type NixEvictFilter struct {
RepoDID, OwnerDID string
}
const nixObjectColumns = `cache, store_hash, repo_did, owner_did, nar_size, last_used_at, gen`
func scanNixObjects(rows *sql.Rows) ([]NixObjectRow, error) {
defer rows.Close()
var out []NixObjectRow
for rows.Next() {
var r NixObjectRow
var lastUsed int64
if err := rows.Scan(&r.Cache, &r.StoreHash, &r.RepoDID, &r.OwnerDID, &r.NarSize, &lastUsed, &r.Gen); err != nil {
return nil, err
}
r.LastUsedAt = time.Unix(lastUsed, 0)
out = append(out, r)
}
return out, rows.Err()
}
func (d *QuotaStore) upsertNixObject(ctx context.Context, conn execQuerier, repoDID, ownerDID, encoded string) error {
var obj quota.NixObject
if err := json.Unmarshal([]byte(encoded), &obj); err != nil {
return err
}
now := d.now().Unix()
if _, err := conn.ExecContext(ctx, `
INSERT INTO nix_objects (cache, store_hash, repo_did, owner_did, store_path, nar_size, seq, last_used_at, state, gen)
VALUES (?, ?, ?, ?, ?, ?, (SELECT coalesce(max(seq), 0) + 1 FROM nix_objects), ?, 'live', 1)
ON CONFLICT(cache, store_hash) DO UPDATE SET
repo_did = excluded.repo_did,
owner_did = excluded.owner_did,
store_path = excluded.store_path,
nar_size = excluded.nar_size,
last_used_at = excluded.last_used_at,
state = 'live',
gen = gen + 1
`, obj.Cache, obj.StoreHash, repoDID, ownerDID, obj.StorePath, obj.NarSize, now); err != nil {
return err
}
if _, err := conn.ExecContext(ctx, `DELETE FROM nix_refs WHERE cache = ? AND from_hash = ?`, obj.Cache, obj.StoreHash); err != nil {
return err
}
for _, ref := range obj.References {
hash, _, _ := strings.Cut(ref, "-")
if hash == "" || hash == obj.StoreHash {
continue
}
if _, err := conn.ExecContext(ctx, `
INSERT OR IGNORE INTO nix_refs (cache, from_hash, to_hash) VALUES (?, ?, ?)
`, obj.Cache, obj.StoreHash, hash); err != nil {
return err
}
// bumping gen makes the pending finish skip it. if attic already deleted it, the reference dangles and nix builds it
if _, err := conn.ExecContext(ctx, `
UPDATE nix_objects SET state = 'live', gen = gen + 1
WHERE cache = ? AND store_hash = ? AND state = 'evicting'
`, obj.Cache, hash); err != nil {
return err
}
}
return nil
}
func (d *QuotaStore) RecordNixUse(ctx context.Context, cache string, storeHashes []string) error {
return d.withTx(ctx, func(conn *sql.Conn) error {
now := d.now().Unix()
for batch := range slices.Chunk(storeHashes, nixUseBatch) {
args := []any{now, cache}
for _, hash := range batch {
args = append(args, hash)
}
if _, err := conn.ExecContext(ctx, `
UPDATE nix_objects SET last_used_at = max(last_used_at, ?)
WHERE cache = ? AND store_hash IN (`+placeholders(len(batch))+`)
`, args...); err != nil {
return err
}
}
return nil
})
}
// NixCacheUsage counts evicting rows too because their bytes stay in the cache until the delete lands
func (d *QuotaStore) NixCacheUsage(ctx context.Context) (byRepo, byOwner map[string]int64, total int64, err error) {
byRepo, byOwner = map[string]int64{}, map[string]int64{}
err = d.withReadTx(ctx, func(tx *sql.Tx) error {
rows, err := tx.QueryContext(ctx, `SELECT repo_did, owner_did, SUM(nar_size) FROM nix_objects GROUP BY repo_did, owner_did`)
if err != nil {
return err
}
defer rows.Close()
for rows.Next() {
var repo, owner string
var size int64
if err := rows.Scan(&repo, &owner, &size); err != nil {
return err
}
byRepo[repo] += size
byOwner[owner] += size
total += size
}
return rows.Err()
})
return byRepo, byOwner, total, err
}
// a referrer pins its dependencies only if committed after them. a cycle can't pin itself that way
// an out of order upload can leave a dangling reference that nix builds
const noReferrer = `NOT EXISTS (
SELECT 1 FROM nix_refs r
JOIN nix_objects p ON p.cache = r.cache AND p.store_hash = r.from_hash AND p.seq > o.seq
WHERE r.cache = o.cache AND r.to_hash = o.store_hash
)`
func (d *QuotaStore) NixEvictable(ctx context.Context, f NixEvictFilter, usedBefore time.Time, limit int) ([]NixObjectRow, error) {
var out []NixObjectRow
err := d.withReadTx(ctx, func(tx *sql.Tx) error {
where, args := `state = 'live' AND last_used_at < ?`, []any{usedBefore.Unix()}
if f.RepoDID != "" {
where += ` AND repo_did = ?`
args = append(args, f.RepoDID)
}
if f.OwnerDID != "" {
where += ` AND owner_did = ?`
args = append(args, f.OwnerDID)
}
rows, err := tx.QueryContext(ctx, `
SELECT `+nixObjectColumns+` FROM nix_objects o
WHERE `+where+` AND `+noReferrer+`
ORDER BY last_used_at, cache, store_hash
LIMIT ?
`, append(args, limit)...)
if err != nil {
return err
}
out, err = scanNixObjects(rows)
return err
})
return out, err
}
func (d *QuotaStore) MarkNixEvicting(ctx context.Context, rows []NixObjectRow, usedBefore time.Time) ([]NixObjectRow, error) {
var marked []NixObjectRow
err := d.withTx(ctx, func(conn *sql.Conn) error {
marked = nil
for _, r := range rows {
res, err := conn.ExecContext(ctx, `
UPDATE nix_objects AS o SET state = 'evicting'
WHERE cache = ? AND store_hash = ? AND state = 'live' AND gen = ? AND last_used_at < ? AND `+noReferrer,
r.Cache, r.StoreHash, r.Gen, usedBefore.Unix())
if err != nil {
return err
}
if n, err := res.RowsAffected(); err != nil {
return err
} else if n == 1 {
marked = append(marked, r)
}
}
return nil
})
if err != nil {
return nil, err
}
return marked, nil
}
func (d *QuotaStore) NixEvicting(ctx context.Context) ([]NixObjectRow, error) {
var out []NixObjectRow
err := d.withReadTx(ctx, func(tx *sql.Tx) error {
rows, err := tx.QueryContext(ctx, `
SELECT `+nixObjectColumns+` FROM nix_objects WHERE state = 'evicting' ORDER BY cache, store_hash
`)
if err != nil {
return err
}
out, err = scanNixObjects(rows)
return err
})
return out, err
}
// FinishNixEviction skips a row that was uploaded again, because that bumps its gen
// one tx per chunk, and a failure leaves the rest evicting for the next pass
func (d *QuotaStore) FinishNixEviction(ctx context.Context, cache string, gone []NixObjectRow) error {
for chunk := range slices.Chunk(gone, nixDropBatch) {
err := d.withTx(ctx, func(conn *sql.Conn) error {
for _, r := range chunk {
res, err := conn.ExecContext(ctx, `
DELETE FROM nix_objects WHERE cache = ? AND store_hash = ? AND state = 'evicting' AND gen = ?`,
cache, r.StoreHash, r.Gen)
if err != nil {
return err
}
if n, err := res.RowsAffected(); err != nil {
return err
} else if n != 1 {
continue
}
if _, err := conn.ExecContext(ctx, `DELETE FROM nix_refs WHERE cache = ? AND from_hash = ?`, cache, r.StoreHash); err != nil {
return err
}
if _, err := conn.ExecContext(ctx, `DELETE FROM quota_allocations WHERE repo_did = ? AND kind = ? AND key = ?`,
r.RepoDID, quota.KindNixCache, quota.NixCacheKey(cache, r.StoreHash)); err != nil {
return err
}
}
return nil
})
if err != nil {
return err
}
}
return nil
}
func placeholders(n int) string {
return strings.TrimSuffix(strings.Repeat("?,", n), ",")
}
func (d *QuotaStore) RecordNixGrant(ctx context.Context, expires time.Time, life time.Duration) error {
return d.withTx(ctx, func(conn *sql.Conn) error {
if _, err := conn.ExecContext(ctx, `DELETE FROM nix_grants WHERE expires_at <= ?`, d.now().Unix()); err != nil {
return err
}
_, err := conn.ExecContext(ctx, `INSERT INTO nix_grants (expires_at, life) VALUES (?, ?)`,
expires.Add(time.Second-1).Unix(), int64((life+time.Second-1)/time.Second))
return err
})
}
func (d *QuotaStore) LongestLiveNixGrant(ctx context.Context, now time.Time) (time.Duration, error) {
var life int64
err := d.withReadTx(ctx, func(tx *sql.Tx) error {
return tx.QueryRowContext(ctx, `SELECT coalesce(max(life), 0) FROM nix_grants WHERE expires_at > ?`, now.Unix()).Scan(&life)
})
return time.Duration(life) * time.Second, err
}