758 lines
23 KiB
C
758 lines
23 KiB
C
/*
|
|
* disk-cache.c
|
|
* Copyright (C) 2020 Kovid Goyal <kovid at kovidgoyal.net>
|
|
*
|
|
* Distributed under terms of the GPL3 license.
|
|
*/
|
|
|
|
#define EXTRA_INIT if (PyModule_AddFunctions(module, module_methods) != 0) return false;
|
|
#define MAX_KEY_SIZE 256u
|
|
#if __linux__
|
|
#define HAS_SENDFILE
|
|
#endif
|
|
|
|
#include "disk-cache.h"
|
|
#include "uthash.h"
|
|
#include "loop-utils.h"
|
|
#include "threading.h"
|
|
#include "cross-platform-random.h"
|
|
#include <structmember.h>
|
|
#include <stdlib.h>
|
|
#include <sys/stat.h>
|
|
#include <fcntl.h>
|
|
#include <unistd.h>
|
|
#ifdef HAS_SENDFILE
|
|
#include <sys/sendfile.h>
|
|
#endif
|
|
|
|
|
|
typedef struct {
|
|
void *hash_key;
|
|
uint8_t *data;
|
|
size_t data_sz;
|
|
unsigned short hash_keylen;
|
|
bool written_to_disk;
|
|
off_t pos_in_cache_file;
|
|
uint8_t encryption_key[64];
|
|
UT_hash_handle hh;
|
|
} CacheEntry;
|
|
|
|
|
|
typedef struct {
|
|
PyObject_HEAD
|
|
char *cache_dir;
|
|
int cache_file_fd;
|
|
pthread_mutex_t lock;
|
|
pthread_t write_thread;
|
|
bool thread_started, lock_inited, loop_data_inited, shutting_down, fully_initialized;
|
|
LoopData loop_data;
|
|
CacheEntry *entries, currently_writing;
|
|
unsigned long long total_size;
|
|
} DiskCache;
|
|
|
|
static void
|
|
xor_data(const uint8_t* restrict key, const size_t key_sz, uint8_t* restrict data, const size_t data_sz) {
|
|
size_t unaligned_sz = data_sz % key_sz;
|
|
size_t aligned_sz = data_sz - unaligned_sz;
|
|
for (size_t offset = 0; offset < aligned_sz; offset += key_sz) {
|
|
for (size_t i = 0; i < key_sz; i++) data[offset + i] ^= key[i];
|
|
}
|
|
for (size_t i = 0; i < unaligned_sz; i++) data[aligned_sz + i] ^= key[i];
|
|
}
|
|
|
|
|
|
void
|
|
free_cache_entry(CacheEntry *e) {
|
|
if (e->hash_key) { free(e->hash_key); e->hash_key = NULL; }
|
|
if (e->data) { free(e->data); e->data = NULL; }
|
|
free(e);
|
|
}
|
|
|
|
#define mutex(op) pthread_mutex_##op(&self->lock)
|
|
|
|
static PyObject*
|
|
new(PyTypeObject *type, PyObject UNUSED *args, PyObject UNUSED *kwds) {
|
|
DiskCache *self;
|
|
self = (DiskCache*)type->tp_alloc(type, 0);
|
|
if (self) {
|
|
self->cache_file_fd = -1;
|
|
}
|
|
return (PyObject*) self;
|
|
}
|
|
|
|
static int
|
|
open_cache_file(const char *cache_path) {
|
|
size_t sz = strlen(cache_path) + 16;
|
|
char *buf = calloc(1, sz);
|
|
if (!buf) { errno = ENOMEM; return -1; }
|
|
snprintf(buf, sz - 1, "%s/XXXXXXXXXXXX", cache_path);
|
|
int fd = -1;
|
|
while (fd < 0) {
|
|
fd = mkostemp(buf, O_CLOEXEC);
|
|
if (fd > -1 || errno != EINTR) break;
|
|
}
|
|
if (fd > -1) unlink(buf);
|
|
free(buf);
|
|
return fd;
|
|
}
|
|
|
|
// Write loop {{{
|
|
static bool
|
|
copy_between_files(int infd, int outfd, off_t in_pos, size_t len, uint8_t *buf, size_t bufsz) {
|
|
#ifdef HAS_SENDFILE
|
|
(void)buf; (void)bufsz;
|
|
while (len) {
|
|
off_t r = in_pos;
|
|
ssize_t n = sendfile(outfd, infd, &r, len);
|
|
if (n < 0) {
|
|
if (errno != EAGAIN) return false;
|
|
continue;
|
|
}
|
|
in_pos += n; len -= n;
|
|
}
|
|
#else
|
|
while (len) {
|
|
ssize_t amt_read = pread(infd, buf, MIN(len, bufsz), in_pos);
|
|
if (amt_read < 0) {
|
|
if (errno == EINTR || errno == EAGAIN) continue;
|
|
return false;
|
|
}
|
|
if (amt_read == 0) {
|
|
errno = EIO;
|
|
return false;
|
|
}
|
|
len -= amt_read;
|
|
in_pos += amt_read;
|
|
uint8_t *p = buf;
|
|
while(amt_read) {
|
|
ssize_t amt_written = write(outfd, p, amt_read);
|
|
if (amt_written < 0) {
|
|
if (errno == EINTR || errno == EAGAIN) continue;
|
|
return false;
|
|
}
|
|
if (amt_written == 0) {
|
|
errno = EIO;
|
|
return false;
|
|
}
|
|
amt_read -= amt_written;
|
|
p += amt_written;
|
|
}
|
|
}
|
|
#endif
|
|
return true;
|
|
}
|
|
|
|
static inline off_t
|
|
size_of_cache_file(DiskCache *self) {
|
|
return lseek(self->cache_file_fd, 0, SEEK_END);
|
|
}
|
|
|
|
size_t
|
|
disk_cache_size_on_disk(PyObject *self) {
|
|
off_t ans = size_of_cache_file((DiskCache*)self);
|
|
return MAX(0, ans);
|
|
}
|
|
|
|
typedef struct {
|
|
uint8_t hash_key[MAX_KEY_SIZE];
|
|
unsigned short hash_keylen;
|
|
off_t old_offset, new_offset;
|
|
size_t data_sz;
|
|
} DefragEntry;
|
|
|
|
static void
|
|
defrag(DiskCache *self) {
|
|
int new_cache_file = -1;
|
|
DefragEntry *defrag_entries = NULL;
|
|
uint8_t *buf = NULL;
|
|
const size_t bufsz = 1024 * 1024;
|
|
bool lock_released = false, ok = false;
|
|
|
|
off_t size_on_disk = size_of_cache_file(self);
|
|
if (size_on_disk <= 0) goto cleanup;
|
|
size_t num_entries = HASH_COUNT(self->entries);
|
|
if (!num_entries) goto cleanup;
|
|
new_cache_file = open_cache_file(self->cache_dir);
|
|
if (new_cache_file < 0) {
|
|
perror("Failed to open second file for defrag of disk cache");
|
|
goto cleanup;
|
|
}
|
|
defrag_entries = calloc(num_entries, sizeof(DefragEntry));
|
|
if (!defrag_entries) goto cleanup;
|
|
size_t total_data_size = 0, num_entries_to_defrag = 0;
|
|
CacheEntry *tmp, *s;
|
|
HASH_ITER(hh, self->entries, s, tmp) {
|
|
if (s->pos_in_cache_file > -1 && s->data_sz) {
|
|
total_data_size += s->data_sz;
|
|
DefragEntry *e = defrag_entries + num_entries_to_defrag++;
|
|
e->hash_keylen = s->hash_keylen;
|
|
e->old_offset = s->pos_in_cache_file;
|
|
e->data_sz = s->data_sz;
|
|
if (s->hash_key) memcpy(e->hash_key, s->hash_key, s->hash_keylen);
|
|
num_entries_to_defrag++;
|
|
}
|
|
}
|
|
if (ftruncate(new_cache_file, total_data_size) != 0) {
|
|
perror("Failed to allocate space for new disk cache file during defrag");
|
|
goto cleanup;
|
|
}
|
|
lseek(new_cache_file, 0, SEEK_SET);
|
|
#ifndef HAS_SENDFILE
|
|
buf = malloc(bufsz);
|
|
if (!buf) goto cleanup;
|
|
#endif
|
|
|
|
mutex(unlock); lock_released = true;
|
|
|
|
off_t current_pos = 0;
|
|
for (size_t i = 0; i < num_entries_to_defrag; i++) {
|
|
DefragEntry *e = defrag_entries + i;
|
|
if (!copy_between_files(self->cache_file_fd, new_cache_file, e->old_offset, e->data_sz, buf, bufsz)) {
|
|
perror("Failed to copy data to new disk cache file during defrag");
|
|
goto cleanup;
|
|
}
|
|
e->new_offset = current_pos;
|
|
current_pos += e->data_sz;
|
|
}
|
|
ok = true;
|
|
|
|
cleanup:
|
|
if (lock_released) mutex(lock);
|
|
if (ok) {
|
|
safe_close(self->cache_file_fd, __FILE__, __LINE__);
|
|
self->cache_file_fd = new_cache_file; new_cache_file = -1;
|
|
for (size_t i = 0; i < num_entries_to_defrag; i++) {
|
|
DefragEntry *e = defrag_entries + i;
|
|
s = NULL;
|
|
HASH_FIND(hh, self->entries, e->hash_key, e->hash_keylen, s);
|
|
if (s) s->pos_in_cache_file = e->new_offset;
|
|
}
|
|
}
|
|
if (defrag_entries) free(defrag_entries);
|
|
if (buf) free(buf);
|
|
if (new_cache_file > -1) safe_close(new_cache_file, __FILE__, __LINE__);
|
|
}
|
|
|
|
static int
|
|
cmp_pos_in_cache_file(void *a_, void *b_) {
|
|
CacheEntry *a = a_, *b = b_;
|
|
return a->pos_in_cache_file - b->pos_in_cache_file;
|
|
}
|
|
|
|
static inline void
|
|
find_hole(DiskCache *self) {
|
|
off_t required_size = self->currently_writing.data_sz, prev = -100;
|
|
HASH_SORT(self->entries, cmp_pos_in_cache_file);
|
|
CacheEntry *s, *tmp;
|
|
HASH_ITER(hh, self->entries, s, tmp) {
|
|
if (s->pos_in_cache_file >= 0 && s->data_sz > 0) {
|
|
if (prev >= 0 && s->pos_in_cache_file - prev >= required_size) {
|
|
self->currently_writing.pos_in_cache_file = prev;
|
|
return;
|
|
}
|
|
prev = s->pos_in_cache_file + s->data_sz;
|
|
}
|
|
}
|
|
}
|
|
|
|
static inline bool
|
|
find_cache_entry_to_write(DiskCache *self) {
|
|
CacheEntry *tmp, *s;
|
|
off_t size_on_disk = size_of_cache_file(self);
|
|
if (self->total_size && size_on_disk > 0 && (size_t)size_on_disk > self->total_size * 2) defrag(self);
|
|
HASH_ITER(hh, self->entries, s, tmp) {
|
|
if (!s->written_to_disk) {
|
|
if (s->data) {
|
|
if (self->currently_writing.data) free(self->currently_writing.data);
|
|
self->currently_writing.data = s->data;
|
|
s->data = NULL;
|
|
self->currently_writing.data_sz = s->data_sz;
|
|
self->currently_writing.pos_in_cache_file = -1;
|
|
xor_data(s->encryption_key, sizeof(s->encryption_key), self->currently_writing.data, s->data_sz);
|
|
self->currently_writing.hash_keylen = MIN(s->hash_keylen, MAX_KEY_SIZE);
|
|
memcpy(self->currently_writing.hash_key, s->hash_key, self->currently_writing.hash_keylen);
|
|
find_hole(self);
|
|
return true;
|
|
}
|
|
s->written_to_disk = true;
|
|
s->pos_in_cache_file = 0;
|
|
s->data_sz = 0;
|
|
}
|
|
}
|
|
return false;
|
|
}
|
|
|
|
static inline bool
|
|
write_dirty_entry(DiskCache *self) {
|
|
size_t left = self->currently_writing.data_sz;
|
|
uint8_t *p = self->currently_writing.data;
|
|
if (self->currently_writing.pos_in_cache_file < 0) {
|
|
self->currently_writing.pos_in_cache_file = size_of_cache_file(self);
|
|
if (self->currently_writing.pos_in_cache_file < 0) {
|
|
perror("Failed to seek in disk cache file");
|
|
return false;
|
|
}
|
|
}
|
|
off_t offset = self->currently_writing.pos_in_cache_file;
|
|
while (left > 0) {
|
|
ssize_t n = pwrite(self->cache_file_fd, p, left, offset);
|
|
if (n < 0) {
|
|
if (errno == EINTR || errno == EAGAIN) continue;
|
|
perror("Failed to write to disk-cache file");
|
|
self->currently_writing.pos_in_cache_file = -1;
|
|
return false;
|
|
}
|
|
if (n == 0) {
|
|
fprintf(stderr, "Failed to write to disk-cache file with zero return\n");
|
|
self->currently_writing.pos_in_cache_file = -1;
|
|
return false;
|
|
}
|
|
left -= n;
|
|
p += n;
|
|
offset += n;
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static inline void
|
|
retire_currently_writing(DiskCache *self) {
|
|
CacheEntry *s = NULL;
|
|
HASH_FIND(hh, self->entries, self->currently_writing.hash_key, self->currently_writing.hash_keylen, s);
|
|
if (s) {
|
|
s->written_to_disk = true;
|
|
s->pos_in_cache_file = self->currently_writing.pos_in_cache_file;
|
|
}
|
|
free(self->currently_writing.data);
|
|
self->currently_writing.data = NULL;
|
|
self->currently_writing.data_sz = 0;
|
|
}
|
|
|
|
static void*
|
|
write_loop(void *data) {
|
|
DiskCache *self = (DiskCache*)data;
|
|
set_thread_name("DiskCacheWrite");
|
|
struct pollfd fds[1] = {0};
|
|
fds[0].fd = self->loop_data.wakeup_read_fd;
|
|
fds[0].events = POLLIN;
|
|
bool found_dirty_entry = false;
|
|
|
|
while (!self->shutting_down) {
|
|
mutex(lock);
|
|
found_dirty_entry = find_cache_entry_to_write(self);
|
|
size_t count = HASH_COUNT(self->entries);
|
|
mutex(unlock);
|
|
if (found_dirty_entry) {
|
|
write_dirty_entry(self);
|
|
mutex(lock);
|
|
retire_currently_writing(self);
|
|
mutex(unlock);
|
|
continue;
|
|
} else if (!count) {
|
|
mutex(lock);
|
|
if (ftruncate(self->cache_file_fd, 0) == 0) lseek(self->cache_file_fd, 0, SEEK_END);
|
|
mutex(unlock);
|
|
}
|
|
|
|
if (poll(fds, 1, -1) > 0 && fds[0].revents & POLLIN) {
|
|
drain_fd(fds[0].fd); // wakeup
|
|
}
|
|
}
|
|
return 0;
|
|
}
|
|
// }}}
|
|
|
|
static bool
|
|
ensure_state(DiskCache *self) {
|
|
int ret;
|
|
if (self->fully_initialized) return true;
|
|
if (!self->loop_data_inited) {
|
|
if (!init_loop_data(&self->loop_data)) { PyErr_SetFromErrno(PyExc_OSError); return false; }
|
|
self->loop_data_inited = true;
|
|
}
|
|
if (!self->currently_writing.hash_key) {
|
|
self->currently_writing.hash_key = malloc(MAX_KEY_SIZE);
|
|
if (!self->currently_writing.hash_key) { PyErr_NoMemory(); return false; }
|
|
}
|
|
|
|
if (!self->lock_inited) {
|
|
if ((ret = pthread_mutex_init(&self->lock, NULL)) != 0) {
|
|
PyErr_Format(PyExc_OSError, "Failed to create disk cache lock mutex: %s", strerror(ret));
|
|
return false;
|
|
}
|
|
self->lock_inited = true;
|
|
}
|
|
|
|
if (!self->thread_started) {
|
|
if ((ret = pthread_create(&self->write_thread, NULL, write_loop, self)) != 0) {
|
|
PyErr_Format(PyExc_OSError, "Failed to start disk cache write thread with error: %s", strerror(ret));
|
|
return false;
|
|
}
|
|
self->thread_started = true;
|
|
}
|
|
|
|
if (!self->cache_dir) {
|
|
PyObject *kc = NULL, *cache_dir = NULL;
|
|
kc = PyImport_ImportModule("kitty.constants");
|
|
if (kc) {
|
|
cache_dir = PyObject_CallMethod(kc, "cache_dir", NULL);
|
|
if (cache_dir) {
|
|
self->cache_dir = strdup(PyUnicode_AsUTF8(cache_dir));
|
|
if (!self->cache_dir) PyErr_NoMemory();
|
|
}
|
|
}
|
|
Py_CLEAR(kc); Py_CLEAR(cache_dir);
|
|
if (PyErr_Occurred()) return false;
|
|
}
|
|
|
|
if (self->cache_file_fd < 0) {
|
|
self->cache_file_fd = open_cache_file(self->cache_dir);
|
|
if (self->cache_file_fd < 0) {
|
|
PyErr_SetFromErrnoWithFilename(PyExc_OSError, self->cache_dir);
|
|
return false;
|
|
}
|
|
}
|
|
|
|
self->fully_initialized = true;
|
|
return true;
|
|
}
|
|
|
|
static void
|
|
wakeup_write_loop(DiskCache *self) {
|
|
if (self->thread_started) wakeup_loop(&self->loop_data, false, "disk_cache_write_loop");
|
|
}
|
|
|
|
static void
|
|
dealloc(DiskCache* self) {
|
|
self->shutting_down = true;
|
|
if (self->thread_started) {
|
|
wakeup_write_loop(self);
|
|
pthread_join(self->write_thread, NULL);
|
|
self->thread_started = false;
|
|
}
|
|
if (self->currently_writing.hash_key) {
|
|
free(self->currently_writing.hash_key); self->currently_writing.hash_key = NULL;
|
|
}
|
|
if (self->lock_inited) {
|
|
pthread_mutex_destroy(&self->lock);
|
|
self->lock_inited = false;
|
|
}
|
|
if (self->loop_data_inited) {
|
|
free_loop_data(&self->loop_data);
|
|
self->loop_data_inited = false;
|
|
}
|
|
if (self->entries) {
|
|
CacheEntry *tmp, *s;
|
|
HASH_ITER(hh, self->entries, s, tmp) {
|
|
HASH_DEL(self->entries, s);
|
|
free_cache_entry(s); s = NULL;
|
|
}
|
|
self->entries = NULL;
|
|
}
|
|
if (self->cache_file_fd > -1) {
|
|
safe_close(self->cache_file_fd, __FILE__, __LINE__);
|
|
self->cache_file_fd = -1;
|
|
}
|
|
if (self->currently_writing.data) free(self->currently_writing.data);
|
|
Py_TYPE(self)->tp_free((PyObject*)self);
|
|
}
|
|
|
|
static inline CacheEntry*
|
|
create_cache_entry(const void *key, const size_t key_sz) {
|
|
CacheEntry *s = calloc(1, sizeof(CacheEntry));
|
|
if (!s) return (CacheEntry*)PyErr_NoMemory();
|
|
if (!secure_random_bytes(s->encryption_key, sizeof(s->encryption_key))) { free(s); PyErr_SetFromErrno(PyExc_OSError); return NULL; }
|
|
s->hash_key = malloc(key_sz);
|
|
if (!s->hash_key) { free(s); PyErr_NoMemory(); return NULL; }
|
|
s->hash_keylen = key_sz;
|
|
memcpy(s->hash_key, key, key_sz);
|
|
s->pos_in_cache_file = -2;
|
|
return s;
|
|
}
|
|
|
|
bool
|
|
add_to_disk_cache(PyObject *self_, const void *key, size_t key_sz, const void *data, size_t data_sz) {
|
|
DiskCache *self = (DiskCache*)self_;
|
|
if (!ensure_state(self)) return false;
|
|
if (key_sz > MAX_KEY_SIZE) { PyErr_SetString(PyExc_KeyError, "cache key is too long"); return false; }
|
|
CacheEntry *s = NULL;
|
|
uint8_t *copied_data = malloc(data_sz);
|
|
if (!copied_data) { PyErr_NoMemory(); return false; }
|
|
memcpy(copied_data, data, data_sz);
|
|
|
|
mutex(lock);
|
|
HASH_FIND(hh, self->entries, key, key_sz, s);
|
|
if (s == NULL) {
|
|
if (!(s = create_cache_entry(key, key_sz))) goto end;
|
|
HASH_ADD_KEYPTR(hh, self->entries, s->hash_key, s->hash_keylen, s);
|
|
} else {
|
|
s->written_to_disk = false;
|
|
if (s->data) free(s->data);
|
|
if (data_sz <= self->total_size) self->total_size -= data_sz;
|
|
else self->total_size = 0;
|
|
}
|
|
s->data = copied_data; s->data_sz = data_sz; copied_data = NULL;
|
|
self->total_size += s->data_sz;
|
|
end:
|
|
mutex(unlock);
|
|
|
|
if (copied_data) free(copied_data);
|
|
if (PyErr_Occurred()) return false;
|
|
wakeup_write_loop(self);
|
|
return true;
|
|
}
|
|
|
|
bool
|
|
remove_from_disk_cache(PyObject *self_, const void *key, size_t key_sz) {
|
|
DiskCache *self = (DiskCache*)self_;
|
|
if (!ensure_state(self)) return false;
|
|
if (key_sz > MAX_KEY_SIZE) { PyErr_SetString(PyExc_KeyError, "cache key is too long"); return false; }
|
|
CacheEntry *s = NULL;
|
|
bool removed = false;
|
|
|
|
mutex(lock);
|
|
HASH_FIND(hh, self->entries, key, key_sz, s);
|
|
if (s) {
|
|
removed = true;
|
|
HASH_DEL(self->entries, s);
|
|
self->total_size = (self->total_size > s->data_sz) ? self->total_size - s->data_sz : 0;
|
|
free_cache_entry(s);
|
|
}
|
|
mutex(unlock);
|
|
wakeup_write_loop(self);
|
|
return removed;
|
|
}
|
|
|
|
void
|
|
clear_disk_cache(PyObject *self_) {
|
|
DiskCache *self = (DiskCache*)self_;
|
|
CacheEntry *s, *tmp;
|
|
mutex(lock);
|
|
HASH_ITER(hh, self->entries, s, tmp) {
|
|
HASH_DEL(self->entries, s);
|
|
free_cache_entry(s);
|
|
}
|
|
mutex(unlock);
|
|
wakeup_write_loop(self);
|
|
}
|
|
|
|
static void
|
|
read_from_cache_file(const DiskCache *self, off_t pos, size_t sz, void *dest) {
|
|
uint8_t *p = dest;
|
|
while (sz) {
|
|
ssize_t n = pread(self->cache_file_fd, p, sz, pos);
|
|
if (n > 0) {
|
|
sz -= n;
|
|
p += n;
|
|
pos += n;
|
|
continue;
|
|
}
|
|
if (n < 0) {
|
|
if (errno == EINTR || errno == EAGAIN) continue;
|
|
PyErr_SetFromErrnoWithFilename(PyExc_OSError, self->cache_dir);
|
|
break;
|
|
}
|
|
if (n == 0) {
|
|
PyErr_SetString(PyExc_OSError, "Disk cache file truncated");
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
static void
|
|
read_from_cache_entry(const DiskCache *self, const CacheEntry *s, void *dest) {
|
|
size_t sz = s->data_sz;
|
|
off_t pos = s->pos_in_cache_file;
|
|
if (pos < 0) {
|
|
PyErr_SetString(PyExc_OSError, "Cache entry was not written, could not read from it");
|
|
return;
|
|
}
|
|
read_from_cache_file(self, pos, sz, dest);
|
|
}
|
|
|
|
void*
|
|
read_from_disk_cache(PyObject *self_, const void *key, size_t key_sz, void*(allocator)(void*, size_t), void* allocator_data) {
|
|
DiskCache *self = (DiskCache*)self_;
|
|
void *data = NULL;
|
|
if (!ensure_state(self)) return data;
|
|
if (key_sz > MAX_KEY_SIZE) { PyErr_SetString(PyExc_KeyError, "cache key is too long"); return data; }
|
|
|
|
mutex(lock);
|
|
CacheEntry *s = NULL;
|
|
HASH_FIND(hh, self->entries, key, key_sz, s);
|
|
if (!s) { PyErr_SetString(PyExc_KeyError, "No cached entry with specified key found"); goto end; }
|
|
data = allocator(allocator_data, s->data_sz);
|
|
if (!data) { PyErr_NoMemory(); goto end; }
|
|
|
|
if (s->data) { memcpy(data, s->data, s->data_sz); }
|
|
else if (self->currently_writing.data && self->currently_writing.hash_key && self->currently_writing.hash_keylen == key_sz && memcmp(self->currently_writing.hash_key, key, key_sz) == 0) {
|
|
memcpy(data, self->currently_writing.data, s->data_sz);
|
|
xor_data(s->encryption_key, sizeof(s->encryption_key), data, s->data_sz);
|
|
}
|
|
else {
|
|
read_from_cache_entry(self, s, data);
|
|
xor_data(s->encryption_key, sizeof(s->encryption_key), data, s->data_sz);
|
|
}
|
|
end:
|
|
mutex(unlock);
|
|
return data;
|
|
}
|
|
|
|
bool
|
|
disk_cache_wait_for_write(PyObject *self_, monotonic_t timeout) {
|
|
DiskCache *self = (DiskCache*)self_;
|
|
monotonic_t end_at = monotonic() + timeout;
|
|
while (!timeout || monotonic() <= end_at) {
|
|
bool pending = false;
|
|
mutex(lock);
|
|
CacheEntry *s, *tmp;
|
|
HASH_ITER(hh, self->entries, s, tmp) {
|
|
if (!s->written_to_disk) {
|
|
pending = true;
|
|
break;
|
|
}
|
|
}
|
|
mutex(unlock);
|
|
if (!pending) return true;
|
|
wakeup_write_loop(self);
|
|
usleep(10 * 1000);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
#define PYWRAP(name) static PyObject* py##name(DiskCache *self, PyObject *args)
|
|
#define PA(fmt, ...) if (!PyArg_ParseTuple(args, fmt, __VA_ARGS__)) return NULL;
|
|
PYWRAP(ensure_state) {
|
|
(void)args;
|
|
ensure_state(self);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
PYWRAP(xor_data) {
|
|
(void) self;
|
|
const char *key, *data;
|
|
Py_ssize_t keylen, data_sz;
|
|
PA("y#y#", &key, &keylen, &data, &data_sz);
|
|
PyObject *ans = PyBytes_FromStringAndSize(NULL, data_sz);
|
|
if (ans == NULL) return NULL;
|
|
void *dest = PyBytes_AS_STRING(ans);
|
|
memcpy(dest, data, data_sz);
|
|
xor_data((const uint8_t*)key, keylen, dest, data_sz);
|
|
return ans;
|
|
}
|
|
|
|
PYWRAP(read_from_cache_file) {
|
|
Py_ssize_t pos = 0, sz = -1;
|
|
PA("|nn", &pos, &sz);
|
|
if (sz < 0) sz = size_of_cache_file(self);
|
|
PyObject *ans = PyBytes_FromStringAndSize(NULL, sz);
|
|
if (ans) {
|
|
read_from_cache_file(self, pos, sz, PyBytes_AS_STRING(ans));
|
|
}
|
|
return ans;
|
|
}
|
|
|
|
static PyObject*
|
|
wait_for_write(PyObject *self, PyObject *args) {
|
|
double timeout = 0;
|
|
PA("|d", &timeout);
|
|
if (disk_cache_wait_for_write(self, s_double_to_monotonic_t(timeout))) Py_RETURN_TRUE;
|
|
Py_RETURN_FALSE;
|
|
}
|
|
|
|
static PyObject*
|
|
size_on_disk(PyObject *self, PyObject *args UNUSED) {
|
|
unsigned long long ans = disk_cache_size_on_disk(self);
|
|
return PyLong_FromUnsignedLongLong(ans);
|
|
}
|
|
|
|
static PyObject*
|
|
clear(PyObject *self, PyObject *args UNUSED) {
|
|
clear_disk_cache(self);
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
|
|
static PyObject*
|
|
add(PyObject *self, PyObject *args) {
|
|
const char *key, *data;
|
|
Py_ssize_t keylen, datalen;
|
|
PA("y#y#", &key, &keylen, &data, &datalen);
|
|
if (!add_to_disk_cache(self, key, keylen, data, datalen)) return NULL;
|
|
Py_RETURN_NONE;
|
|
}
|
|
|
|
static PyObject*
|
|
pyremove(PyObject *self, PyObject *args) {
|
|
const char *key;
|
|
Py_ssize_t keylen;
|
|
PA("y#", &key, &keylen);
|
|
bool removed = remove_from_disk_cache(self, key, keylen);
|
|
if (PyErr_Occurred()) return NULL;
|
|
if (removed) Py_RETURN_TRUE;
|
|
Py_RETURN_FALSE;
|
|
}
|
|
|
|
typedef struct {
|
|
PyObject *bytes;
|
|
} BytesWrapper;
|
|
|
|
static void*
|
|
bytes_alloc(void *x, size_t sz) {
|
|
BytesWrapper *w = x;
|
|
w->bytes = PyBytes_FromStringAndSize(NULL, sz);
|
|
if (!w->bytes) return NULL;
|
|
return PyBytes_AS_STRING(w->bytes);
|
|
}
|
|
|
|
static PyObject*
|
|
get(PyObject *self, PyObject *args) {
|
|
const char *key;
|
|
Py_ssize_t keylen;
|
|
PA("y#", &key, &keylen);
|
|
BytesWrapper w = {0};
|
|
read_from_disk_cache(self, key, keylen, bytes_alloc, &w);
|
|
if (PyErr_Occurred()) { Py_CLEAR(w.bytes); return NULL; }
|
|
return w.bytes;
|
|
}
|
|
|
|
|
|
#define MW(name, arg_type) {#name, (PyCFunction)py##name, arg_type, NULL}
|
|
static PyMethodDef methods[] = {
|
|
MW(ensure_state, METH_NOARGS),
|
|
MW(read_from_cache_file, METH_VARARGS),
|
|
{"add", add, METH_VARARGS, NULL},
|
|
{"remove", pyremove, METH_VARARGS, NULL},
|
|
{"get", get, METH_VARARGS, NULL},
|
|
{"wait_for_write", wait_for_write, METH_VARARGS, NULL},
|
|
{"size_on_disk", size_on_disk, METH_NOARGS, NULL},
|
|
{"clear", clear, METH_NOARGS, NULL},
|
|
|
|
{NULL} /* Sentinel */
|
|
};
|
|
|
|
static PyMemberDef members[] = {
|
|
{"total_size", T_ULONGLONG, offsetof(DiskCache, total_size), READONLY, "total_size"},
|
|
{NULL},
|
|
};
|
|
|
|
|
|
PyTypeObject DiskCache_Type = {
|
|
PyVarObject_HEAD_INIT(NULL, 0)
|
|
.tp_name = "fast_data_types.DiskCache",
|
|
.tp_basicsize = sizeof(DiskCache),
|
|
.tp_dealloc = (destructor)dealloc,
|
|
.tp_flags = Py_TPFLAGS_DEFAULT,
|
|
.tp_doc = "A disk based secure cache",
|
|
.tp_methods = methods,
|
|
.tp_members = members,
|
|
.tp_new = new,
|
|
};
|
|
|
|
static PyMethodDef module_methods[] = {
|
|
MW(xor_data, METH_VARARGS),
|
|
{NULL, NULL, 0, NULL} /* Sentinel */
|
|
};
|
|
|
|
INIT_TYPE(DiskCache)
|
|
PyObject* create_disk_cache(void) { return new(&DiskCache_Type, NULL, NULL); }
|