1#include <block/bcache.h>
2#include <block/block.h>
3#include <block/sched.h>
4#include <console/panic.h>
5#include <math/align.h>
6#include <mem/alloc.h>
7#include <mem/alloc_or_die.h>
8#include <stdbool.h>
9#include <stddef.h>
10#include <stdint.h>
11#include <thread/workqueue.h>
12
13static bool remove(struct bcache *cache, uint64_t key, uint64_t spb);
14
15static bool insert(struct bcache *cache, uint64_t key,
16 struct bcache_entry *value, bool already_locked);
17
18static struct bcache_entry *get(struct bcache *cache, uint64_t key);
19static bool write(struct block_device *d, struct bcache *cache,
20 struct bcache_entry *ent, uint64_t spb);
21
22/* prefetch is asynchronous */
23static enum errno prefetch(struct block_device *disk, struct bcache *cache,
24 uint64_t lba, uint64_t block_size, uint64_t spb);
25
26/* eviction must be explicitly and separately called */
27static bool insert(struct bcache *cache, uint64_t key,
28 struct bcache_entry *value, bool already_locked) {
29 enum irql irql;
30 if (!already_locked)
31 irql = spin_lock(lock: &cache->lock);
32
33 uint64_t index = bcache_hash(x: key, capacity: cache->capacity);
34 struct bcache_wrapper *head = cache->entries[index];
35
36 /* Key already exists */
37 for (struct bcache_wrapper *node = head; node; node = node->next) {
38 if (node->key == key) {
39 node->value = value;
40 node->value->access_time = bcache_get_ticks(cache);
41 if (!already_locked)
42 spin_unlock(lock: &cache->lock, old: irql);
43 return true;
44 }
45 }
46
47 /* New head */
48 struct bcache_wrapper *new_node = kmalloc(sizeof(struct bcache_wrapper));
49 if (!new_node) {
50 if (!already_locked)
51 spin_unlock(lock: &cache->lock, old: irql);
52 return false;
53 }
54
55 new_node->key = key;
56 new_node->value = value;
57 new_node->value->access_time = bcache_get_ticks(cache);
58 new_node->next = head;
59
60 cache->entries[index] = new_node;
61 cache->count++;
62
63 if (!already_locked)
64 spin_unlock(lock: &cache->lock, old: irql);
65 return true;
66}
67
68static struct bcache_entry *get(struct bcache *cache, uint64_t key) {
69 enum irql irql = spin_lock(lock: &cache->lock);
70
71 uint64_t index = bcache_hash(x: key, capacity: cache->capacity);
72 struct bcache_wrapper *node = cache->entries[index];
73
74 while (node) {
75 if (node->key == key) {
76 node->value->access_time = bcache_get_ticks(cache);
77 spin_unlock(lock: &cache->lock, old: irql);
78 return node->value;
79 }
80 node = node->next;
81 }
82
83 spin_unlock(lock: &cache->lock, old: irql);
84 return NULL;
85}
86
87static bool can_remove_lba_group(struct bcache *cache, uint64_t base_lba,
88 uint64_t spb) {
89 /* caller must already hold the lock */
90
91 for (uint64_t i = 0; i < spb; i++) {
92 uint64_t key = base_lba + i;
93 uint64_t index = bcache_hash(x: key, capacity: cache->capacity);
94 struct bcache_wrapper *node = cache->entries[index];
95 bool found = false;
96
97 while (node) {
98 if (node->key == key) {
99 found = true;
100 break;
101 }
102 node = node->next;
103 }
104
105 if (i == 0 && !found)
106 return false; /* base_lba must exist */
107 else if (i > 0 && found)
108 return false; /* another lba in the group is still cached */
109 }
110
111 return true;
112}
113
114static bool remove(struct bcache *cache, uint64_t key, uint64_t spb) {
115 enum irql irql = spin_lock(lock: &cache->lock);
116 uint64_t index = bcache_hash(x: key, capacity: cache->capacity);
117
118 struct bcache_wrapper *prev = NULL;
119 struct bcache_wrapper *node = cache->entries[index];
120
121 while (node) {
122 if (node->key == key) {
123 if (prev)
124 prev->next = node->next;
125 else
126 cache->entries[index] = node->next;
127
128 struct bcache_entry *val = node->value;
129 cache->count--;
130 kfree(node);
131
132 bool should_free = false;
133 if (val && key == val->lba && !val->no_evict) {
134 if (can_remove_lba_group(cache, base_lba: key, spb: val->size / spb)) {
135 should_free = true;
136 }
137 }
138
139 if (should_free) {
140 kfree_aligned(val->buffer);
141 kfree(val);
142 }
143
144 spin_unlock(lock: &cache->lock, old: irql);
145 return true;
146 }
147
148 prev = node;
149 node = node->next;
150 }
151
152 spin_unlock(lock: &cache->lock, old: irql);
153 return false;
154}
155
156struct bcache_pf_data {
157 struct bcache_entry *new_entry;
158 struct bcache *cache;
159};
160
161static void prefetch_callback(struct bio_request *bio) {
162 struct bcache_pf_data *data = bio->user_data;
163 insert(cache: data->cache, key: bio->lba, value: data->new_entry, false);
164
165 kfree(data);
166 kfree(bio);
167}
168
169static enum errno prefetch(struct block_device *disk, struct bcache *cache,
170 uint64_t lba, uint64_t block_size, uint64_t spb) {
171 uint64_t base_lba = ALIGN_DOWN(lba, spb);
172
173 /* no need to re-fetch existing entry */
174 if (get(cache, key: base_lba))
175 return ERR_EXIST;
176
177 struct bcache_pf_data *pf = kmalloc(sizeof(struct bcache_pf_data));
178 if (!pf)
179 return ERR_NO_MEM;
180
181 struct bio_request *req = bio_create_read(d: disk, lba: base_lba, sectors: spb, size: block_size,
182 cb: prefetch_callback, user: pf, NULL);
183 if (!req)
184 return ERR_NO_MEM;
185
186 pf->cache = cache;
187 pf->new_entry = kmalloc(sizeof(struct bcache_entry), ALLOC_FLAGS_ZERO);
188 if (!pf->new_entry)
189 return ERR_NO_MEM;
190
191 struct bcache_entry *ent = pf->new_entry;
192
193 ent->lba = lba;
194 ent->size = block_size;
195 ent->dirty = false;
196 ent->no_evict = false;
197 ent->buffer = req->buffer;
198
199 bio_sched_enqueue(disk, req);
200 return ERR_OK;
201}
202
203static bool evict(struct bcache *cache, uint64_t spb) {
204 enum irql irql = spin_lock(lock: &cache->lock);
205
206 uint64_t oldest = UINT64_MAX;
207 uint64_t target_key = 0;
208 bool found = false;
209
210 for (uint64_t i = 0; i < cache->capacity; i++) {
211 struct bcache_wrapper *node = cache->entries[i];
212 while (node) {
213 struct bcache_entry *entry = node->value;
214
215 if (!entry || entry->no_evict ||
216 atomic_load(&entry->refcount) > 0) {
217 node = node->next;
218 continue;
219 }
220
221 if (entry->access_time < oldest) {
222 oldest = entry->access_time;
223 target_key = node->key;
224 found = true;
225 }
226
227 node = node->next;
228 }
229 }
230
231 spin_unlock(lock: &cache->lock, old: irql);
232
233 if (found)
234 return remove(cache, key: target_key, spb);
235
236 return false;
237}
238
239static void stat(struct bcache *cache, uint64_t *total_dirty_out,
240 uint64_t *total_present_out) {
241 enum irql irql = spin_lock(lock: &cache->lock);
242
243 uint64_t total_dirty = 0;
244 uint64_t total_present = 0;
245
246 for (uint64_t i = 0; i < cache->capacity; i++) {
247 struct bcache_wrapper *node = cache->entries[i];
248 while (node) {
249 if (node->value) {
250 total_present++;
251 if (node->value->dirty)
252 total_dirty++;
253 }
254 node = node->next;
255 }
256 }
257
258 if (total_dirty_out)
259 *total_dirty_out = total_dirty;
260
261 if (total_present_out)
262 *total_present_out = total_present;
263
264 spin_unlock(lock: &cache->lock, old: irql);
265}
266
267/* TODO: writeback */
268static bool write(struct block_device *d, struct bcache *cache,
269 struct bcache_entry *ent, uint64_t spb) {
270 enum irql irql = spin_lock(lock: &cache->lock);
271
272 bool ret = d->write_sector(d, ent->lba, ent->buffer, spb);
273 uint64_t aligned = ALIGN_DOWN(ent->lba, spb);
274 if (aligned != ent->lba)
275 kfree(ent);
276
277 spin_unlock(lock: &cache->lock, old: irql);
278 return ret;
279}
280
281static void write_enqueue_cb(struct bio_request *req) {
282 struct bcache_entry *ent = req->user_data;
283
284 ent->dirty = false;
285 ent->request = NULL;
286 bcache_ent_unpin(ent);
287
288 kfree(req);
289}
290
291static void write_queue(struct block_device *d, struct bcache_entry *ent,
292 uint64_t spb, enum bio_request_priority prio) {
293
294 ent->dirty = true;
295
296 struct bio_request *req = bio_create_write(
297 d, lba: ent->lba, sectors: spb, size: ent->size, cb: write_enqueue_cb, user: ent, buf: ent->buffer);
298
299 req->priority = prio;
300 ent->request = req;
301
302 bcache_ent_pin(ent);
303 bio_sched_enqueue(disk: d, req);
304}
305
306void bcache_destroy(struct bcache *cache) {
307 for (uint64_t i = 0; i < cache->capacity; i++) {
308 struct bcache_wrapper *node = cache->entries[i];
309 while (node) {
310 struct bcache_wrapper *next = node->next;
311 if (node->value) {
312 kfree_aligned(node->value->buffer);
313 kfree(node->value);
314 }
315 kfree(node);
316 node = next;
317 }
318 }
319
320 kfree(cache->entries);
321 cache->entries = NULL;
322 cache->capacity = 0;
323 cache->count = 0;
324}
325
326void *bcache_get(struct block_device *disk, uint64_t lba, uint64_t block_size,
327 uint64_t spb, bool no_evict, struct bcache_entry **out_entry) {
328 uint64_t base_lba = ALIGN_DOWN(lba, spb);
329 struct bcache_entry *ent = get(cache: disk->cache, key: base_lba);
330
331 if (!ent)
332 return bcache_create_ent(disk, lba, block_size, sectors_per_block: spb, no_evict,
333 out_entry);
334
335 *out_entry = ent;
336
337 uint64_t offset = (lba - base_lba) * disk->sector_size;
338 return ent->buffer + offset;
339}
340
341bool bcache_insert(struct block_device *disk, uint64_t lba,
342 struct bcache_entry *ent, uint64_t spb) {
343 if (insert(cache: disk->cache, key: lba, value: ent, false)) {
344 return true;
345 } else {
346 evict(cache: disk->cache, spb);
347 return insert(cache: disk->cache, key: lba, value: ent, false);
348 }
349}
350
351bool bcache_evict(struct block_device *disk, uint64_t spb) {
352 return evict(cache: disk->cache, spb);
353}
354
355bool bcache_writethrough(struct block_device *disk, struct bcache_entry *ent,
356 uint64_t spb) {
357 return write(d: disk, cache: disk->cache, ent, spb);
358}
359
360void bcache_write_queue(struct block_device *disk, struct bcache_entry *ent,
361 uint64_t spb, enum bio_request_priority prio) {
362 write_queue(d: disk, ent, spb, prio);
363}
364
365void bcache_stat(struct block_device *disk, uint64_t *total_dirty_out,
366 uint64_t *total_present_out) {
367 stat(cache: disk->cache, total_dirty_out, total_present_out);
368}
369
370/* TODO: error code here */
371void *bcache_create_ent(struct block_device *disk, uint64_t lba,
372 uint64_t block_size, uint64_t sectors_per_block,
373 bool no_evict, struct bcache_entry **out_entry) {
374 uint64_t base_lba = ALIGN_DOWN(lba, sectors_per_block);
375
376 struct bcache_entry *ent = get(cache: disk->cache, key: base_lba);
377 if (!ent) {
378 uint8_t *buf = kmalloc_aligned(
379 PAGE_SIZE, PAGE_SIZE, ALLOC_FLAGS_PAGEABLE, ALLOC_BEHAVIOR_NORMAL);
380 if (!buf)
381 return NULL;
382
383 if (!disk->read_sector(disk, base_lba, buf, sectors_per_block)) {
384 kfree_aligned(buf);
385 *out_entry = NULL;
386 return NULL;
387 }
388
389 ent = kmalloc(sizeof(struct bcache_entry), ALLOC_FLAGS_ZERO);
390 if (!ent)
391 return NULL;
392
393 mutex_init(mtx: &ent->lock);
394
395 ent->buffer = buf;
396 ent->lba = base_lba;
397 ent->size = block_size;
398 ent->no_evict = no_evict;
399
400 bcache_insert(disk, lba: base_lba, ent, spb: sectors_per_block);
401 }
402
403 *out_entry = ent;
404
405 uint64_t offset = (lba - base_lba) * disk->sector_size;
406 return ent->buffer + offset;
407}
408
409enum errno bcache_prefetch_async(struct block_device *disk, uint64_t lba,
410 uint64_t block_size, uint64_t spb) {
411 return prefetch(disk, cache: disk->cache, ALIGN_DOWN(lba, spb), block_size, spb);
412}
413
414void bcache_init(struct bcache *cache, uint64_t capacity) {
415 spinlock_init(lock: &cache->lock);
416 cache->capacity = capacity;
417 cache->count = 0;
418 cache->entries = alloc_or_die(
419 kmalloc(sizeof(struct bcache_wrapper *) * capacity, ALLOC_FLAGS_ZERO));
420}
421