alice.pds.demo.boltless.dev
This repository has no description
1package db
2
3import (
4 "context"
5 "sort"
6 "time"
7)
8
9type CacheEntry struct {
10 ID string
11 StorageKey string
12 OwnerDID string
13 RepoDID string
14 Engine string
15 CacheKey string
16 CacheHash string
17 SizeBytes int64
18 State string
19 CreatedAt time.Time
20 LastUsedAt time.Time
21}
22
23const cacheEntryColumns = `
24 id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash,
25 size_bytes, state, created_at, last_used_at`
26
27func (d *DB) InsertCacheEntry(ctx context.Context, entry CacheEntry) error {
28 _, err := d.ExecContext(ctx, `
29 insert into cache_entries (
30 id, storage_key, owner_did, repo_did, engine, cache_key, cache_hash,
31 size_bytes, state, created_at, last_used_at
32 ) values (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`,
33 entry.ID,
34 entry.StorageKey,
35 entry.OwnerDID,
36 entry.RepoDID,
37 entry.Engine,
38 entry.CacheKey,
39 entry.CacheHash,
40 entry.SizeBytes,
41 entry.State,
42 entry.CreatedAt.UnixNano(),
43 entry.LastUsedAt.UnixNano(),
44 )
45 return err
46}
47
48func (d *DB) MarkCacheEntryReady(ctx context.Context, id string, sizeBytes int64, now time.Time) ([]CacheEntry, error) {
49 tx, err := d.BeginTx(ctx, nil)
50 if err != nil {
51 return nil, err
52 }
53 defer tx.Rollback()
54
55 var repoDID, engine, key, hash string
56 if err := tx.QueryRowContext(ctx, `
57 update cache_entries
58 set state = 'ready', size_bytes = ?, last_used_at = ?
59 where id = ? and state = 'pending'
60 returning repo_did, engine, cache_key, cache_hash`,
61 sizeBytes, now.UnixNano(), id).Scan(&repoDID, &engine, &key, &hash); err != nil {
62 return nil, err
63 }
64
65 rows, err := tx.QueryContext(ctx, `
66 update cache_entries
67 set state = 'deleting'
68 where repo_did = ? and engine = ? and cache_key = ? and cache_hash = ?
69 and state = 'ready' and id <> ?
70 returning `+cacheEntryColumns, repoDID, engine, key, hash, id)
71 if err != nil {
72 return nil, err
73 }
74 var superseded []CacheEntry
75 for rows.Next() {
76 entry, err := scanCacheEntry(rows)
77 if err != nil {
78 rows.Close()
79 return nil, err
80 }
81 superseded = append(superseded, *entry)
82 }
83 if err := rows.Close(); err != nil {
84 return nil, err
85 }
86 if err := rows.Err(); err != nil {
87 return nil, err
88 }
89 sort.Slice(superseded, func(i, j int) bool {
90 return superseded[i].CreatedAt.After(superseded[j].CreatedAt)
91 })
92 if err := tx.Commit(); err != nil {
93 return nil, err
94 }
95 return superseded, nil
96}
97
98func (d *DB) FindCacheEntry(ctx context.Context, repoDID, engine, key, hash string) (*CacheEntry, error) {
99 return scanCacheEntry(d.QueryRowContext(ctx, `
100 select `+cacheEntryColumns+`
101 from cache_entries
102 where repo_did = ? and engine = ? and cache_key = ? and cache_hash = ? and state = 'ready'
103 order by created_at desc
104 limit 1`, repoDID, engine, key, hash))
105}
106
107func (d *DB) FindFallbackCacheEntry(ctx context.Context, repoDID, engine, key, excludeHash string) (*CacheEntry, error) {
108 return scanCacheEntry(d.QueryRowContext(ctx, `
109 select `+cacheEntryColumns+`
110 from cache_entries
111 where repo_did = ? and engine = ? and cache_key = ?
112 and cache_hash <> ? and cache_hash <> '' and state = 'ready'
113 order by created_at desc
114 limit 1`, repoDID, engine, key, excludeHash))
115}
116
117func (d *DB) TouchCacheEntry(ctx context.Context, id string, now time.Time) error {
118 _, err := d.ExecContext(ctx, `
119 update cache_entries set last_used_at = ? where id = ? and state = 'ready'`, now.UnixNano(), id)
120 return err
121}
122
123func (d *DB) ClaimCacheEntry(ctx context.Context, id, expectedState string, expectedLastUsed time.Time) (bool, error) {
124 result, err := d.ExecContext(ctx, `
125 update cache_entries
126 set state = 'deleting'
127 where id = ? and state = ? and last_used_at = ?`,
128 id, expectedState, expectedLastUsed.UnixNano())
129 if err != nil {
130 return false, err
131 }
132 changed, err := result.RowsAffected()
133 return changed == 1, err
134}
135
136func (d *DB) RestoreCacheEntryState(ctx context.Context, id, state string) error {
137 _, err := d.ExecContext(ctx, `
138 update cache_entries set state = ? where id = ? and state = 'deleting'`, state, id)
139 return err
140}
141
142func (d *DB) ExpiredCacheEntries(ctx context.Context, readyBefore, pendingBefore time.Time, limit int) ([]CacheEntry, error) {
143 ready, err := d.queryCacheEntries(ctx, `
144 select `+cacheEntryColumns+`
145 from cache_entries
146 where state = 'ready' and last_used_at < ?
147 order by last_used_at
148 limit ?`, readyBefore.UnixNano(), limit)
149 if err != nil {
150 return nil, err
151 }
152 recovery, err := d.queryCacheEntries(ctx, `
153 select `+cacheEntryColumns+`
154 from cache_entries
155 where state in ('pending', 'deleting') and created_at < ?
156 order by created_at
157 limit ?`, pendingBefore.UnixNano(), limit)
158 if err != nil {
159 return nil, err
160 }
161 entries := append(ready, recovery...)
162 sort.Slice(entries, func(i, j int) bool {
163 left, right := entries[i].CreatedAt, entries[j].CreatedAt
164 if entries[i].State == "ready" {
165 left = entries[i].LastUsedAt
166 }
167 if entries[j].State == "ready" {
168 right = entries[j].LastUsedAt
169 }
170 return left.Before(right)
171 })
172 if len(entries) > limit {
173 entries = entries[:limit]
174 }
175 return entries, nil
176}
177
178func (d *DB) queryCacheEntries(ctx context.Context, query string, args ...any) ([]CacheEntry, error) {
179 rows, err := d.QueryContext(ctx, query, args...)
180 if err != nil {
181 return nil, err
182 }
183 defer rows.Close()
184 var entries []CacheEntry
185 for rows.Next() {
186 entry, err := scanCacheEntry(rows)
187 if err != nil {
188 return nil, err
189 }
190 entries = append(entries, *entry)
191 }
192 return entries, rows.Err()
193}
194
195func (d *DB) DeleteCacheEntry(ctx context.Context, id string) error {
196 _, err := d.ExecContext(ctx, `delete from cache_entries where id = ?`, id)
197 return err
198}
199
200func (d *DB) CacheUsageByOwner(ctx context.Context, ownerDID string) (bytes, count int64, err error) {
201 err = d.QueryRowContext(ctx, `
202 select coalesce(sum(size_bytes), 0), count(*)
203 from cache_entries
204 where owner_did = ? and state in ('ready', 'deleting')`, ownerDID).Scan(&bytes, &count)
205 return bytes, count, err
206}
207
208type cacheEntryScanner interface {
209 Scan(dest ...any) error
210}
211
212func scanCacheEntry(row cacheEntryScanner) (*CacheEntry, error) {
213 var entry CacheEntry
214 var createdAt, lastUsedAt int64
215 if err := row.Scan(
216 &entry.ID,
217 &entry.StorageKey,
218 &entry.OwnerDID,
219 &entry.RepoDID,
220 &entry.Engine,
221 &entry.CacheKey,
222 &entry.CacheHash,
223 &entry.SizeBytes,
224 &entry.State,
225 &createdAt,
226 &lastUsedAt,
227 ); err != nil {
228 return nil, err
229 }
230 entry.CreatedAt = time.Unix(0, createdAt)
231 entry.LastUsedAt = time.Unix(0, lastUsedAt)
232 return &entry, nil
233}