database.c

   1 #include <stdint.h>
   2 #include <stdlib.h>
   3 #include <string.h>
   4 #include <stdio.h>
   5 #include <unistd.h>
   6 #include <sys/stat.h>
   7 #include <sys/types.h>
   8 #include <sys/mman.h>
   9 #include <sys/file.h>
  10 #include <fcntl.h>
  11 #include "database.h"
  12 #include "config.h"
  13
  14 #define RECORDS_PER_SEGMENT 65536
  15
  16 uint64_t blerg_get_record_count(struct blerg *blerg) {
  17         uint64_t count;
  18         flock(blerg->meta_fd, LOCK_SH);
  19         count = blerg->meta->sequence;
  20         flock(blerg->meta_fd, LOCK_UN);
  21         return count;
  22 }
  23
  24 // Returns last usable record
  25 uint64_t blerg_increment_record_count(struct blerg *blerg) {
  26         uint64_t count;
  27         flock(blerg->meta_fd, LOCK_EX);
  28         count = blerg->meta->sequence++;
  29         flock(blerg->meta_fd, LOCK_UN);
  30         return count;
  31 }
  32
  33 void blerg_segment_close(struct blerg *blerg) {
  34         if (blerg->data != NULL)
  35                 munmap((void *)blerg->data, blerg->data_size);
  36         if (blerg->data_fd != -1)
  37                 close(blerg->data_fd);
  38         if (blerg->index != NULL)
  39                 munmap((void *)blerg->index, RECORDS_PER_SEGMENT * sizeof(struct record));
  40         if (blerg->index_fd != -1)
  41                 close(blerg->index_fd);
  42 }
  43
  44 int blerg_segment_switch(struct blerg *blerg, int new_segment) {
  45         char filename[512];
  46         uint64_t max_sequence = blerg_get_record_count(blerg);
  47         struct stat st;
  48
  49         if (new_segment > max_sequence / RECORDS_PER_SEGMENT) {
  50                 fprintf(stderr, "Cannot switch to sequence beyond last record\n");
  51                 return 0;
  52         }
  53
  54         blerg_segment_close(blerg);
  55
  56         // Load and map the index
  57         snprintf(filename, 512, "%s/index%d", blerg->base_path, new_segment);
  58         blerg->index_fd = open(filename, O_RDWR | O_CREAT, 0600);
  59         if (blerg->index_fd == -1) {
  60                 perror("Could not open index");
  61                 goto open_failed_index_open;
  62         }
  63         flock(blerg->index_fd, LOCK_EX);
  64         fstat(blerg->index_fd, &st);
  65         if (st.st_size == 0) {
  66                 int i;
  67                 struct record r;
  68                 memset((void *)&r, 0, sizeof(struct record));
  69                 for (i = 0; i < RECORDS_PER_SEGMENT; i++) {
  70                         write(blerg->index_fd, &r, sizeof(struct record));
  71                 }
  72         }
  73         flock(blerg->index_fd, LOCK_UN);
  74
  75         blerg->index = (struct record *) mmap(NULL, RECORDS_PER_SEGMENT * sizeof(struct record), PROT_READ | PROT_WRITE, MAP_SHARED, blerg->index_fd, 0);
  76         if (blerg->index == MAP_FAILED) {
  77                 perror("Could not mmap index");
  78                 goto open_failed_index_mmap;
  79         }
  80
  81         // Load data file
  82         sprintf(filename, "%s/data%d", blerg->base_path, new_segment);
  83         blerg->data_fd = open(filename, O_RDWR | O_APPEND | O_CREAT, 0600);
  84         fstat(blerg->data_fd, &st);
  85         blerg->data_size = st.st_size;
  86         if (blerg->data_fd == -1) {
  87                 perror("Could not open data");
  88                 goto open_failed_data_open;
  89         }
  90
  91         if (blerg->data_size > 0) {
  92                 blerg->data = (char *) mmap(NULL, blerg->data_size, PROT_READ, MAP_SHARED, blerg->data_fd, 0);
  93                 if (blerg->data == MAP_FAILED) {
  94                         perror("Could not mmap data");
  95                         goto open_failed_data_mmap;
  96                 }
  97         }
  98
  99         return 1;
 100
 101 open_failed_data_mmap:
 102         close(blerg->data_fd);
 103 open_failed_data_open:
 104         munmap((void *)blerg->index, sizeof(65536 * sizeof(struct record)));
 105 open_failed_index_mmap:
 106         close(blerg->index_fd);
 107 open_failed_index_open:
 108         return 0;
 109 }
 110
 111 int blerg_exists(const char *name) {
 112         int namelen = strlen(name);
 113         char filename[512];
 114
 115         if (namelen > 32) {
 116                 perror("Name too long");
 117                 return 0;
 118         }
 119
 120         snprintf(filename, 512, "%s/%s", DATA_PATH, name);
 121         if (access(filename, F_OK) == -1)
 122                 return 0;
 123         else
 124                 return 1;
 125 }
 126
 127 struct blerg *blerg_open(const char *name) {
 128         int namelen = strlen(name);
 129         char filename[512];
 130         struct stat st;
 131         uint64_t sequence;
 132
 133         if (namelen > 32) {
 134                 perror("Name too long");
 135                 return NULL;
 136         }
 137         struct blerg *blerg = malloc(sizeof(struct blerg));
 138         if (!blerg) {
 139                 perror("Cannot allocate memory for blerg");
 140                 goto open_failed_blerg_malloc;
 141         }
 142         blerg->name = malloc(namelen + 1);
 143         memcpy(blerg->name, name, namelen + 1);
 144         blerg->meta_fd = blerg->index_fd = blerg->data_fd = -1;
 145         blerg->meta = NULL;
 146         blerg->index = NULL;
 147         blerg->data = NULL;
 148
 149         // Make the directory if it doesn't exist
 150         blerg->base_path = malloc(512);
 151         snprintf(blerg->base_path, 512, "%s/%s", DATA_PATH, name);
 152         if (access(blerg->base_path, F_OK) == -1)
 153                 mkdir(blerg->base_path, 0755);
 154
 155         // Open and map metadata
 156         snprintf(filename, 512, "%s/meta", blerg->base_path);
 157         blerg->meta_fd = open(filename, O_RDWR | O_CREAT, 0600);
 158         if (blerg->meta_fd == -1) {
 159                 perror("Could not open metadata");
 160                 goto open_failed_meta_open;
 161         }
 162         fstat(blerg->meta_fd, &st);
 163         if (st.st_size == 0) {
 164                 char *buf = (char *) malloc(sizeof(struct meta));
 165                 memset(buf, 0, sizeof(struct meta));
 166                 write(blerg->meta_fd, buf, sizeof(struct meta));
 167                 free(buf);
 168         }
 169         blerg->meta = (struct meta *) mmap(NULL, sizeof(struct meta), PROT_READ | PROT_WRITE, MAP_SHARED, blerg->meta_fd, 0);
 170         if (blerg->meta == MAP_FAILED) {
 171                 perror("Could not map metadata");
 172                 goto open_failed_meta_mmap;
 173         }
 174
 175         // Open and map index and data for the current segment
 176         blerg->current_segment = blerg_get_record_count(blerg) / RECORDS_PER_SEGMENT;
 177         if (!blerg_segment_switch(blerg, blerg->current_segment)) {
 178                 fprintf(stderr, "Could not switch segment\n");
 179                 goto open_failed_segment_switch;
 180         }
 181
 182         return blerg;
 183
 184 open_failed_segment_switch:
 185         munmap((void *)blerg->meta, sizeof(struct meta));
 186 open_failed_meta_mmap:
 187         close(blerg->meta_fd);
 188 open_failed_meta_open:
 189         free(blerg->name);
 190         free(blerg);
 191 open_failed_blerg_malloc:
 192         return NULL;
 193 }
 194
 195 int blerg_close(struct blerg *blerg) {
 196         blerg_segment_close(blerg);
 197         munmap((void *)blerg->meta, sizeof(struct meta));
 198         close(blerg->meta_fd);
 199         free(blerg->base_path);
 200         free(blerg->name);
 201         free(blerg);
 202         return 1;
 203 }
 204
 205 int blerg_store(struct blerg *blerg, const char *data, int len) {
 206         if (len > 65536) {
 207                 printf("len > 64K\n");
 208                 return -1;
 209         }
 210
 211         flock(blerg->index_fd, LOCK_EX);
 212         flock(blerg->data_fd, LOCK_EX);
 213
 214         uint64_t record = blerg_increment_record_count(blerg);
 215         if (record == -1) {
 216                 printf("Could not find free record\n");
 217                 return -1;
 218         }
 219         int segment = record / RECORDS_PER_SEGMENT;
 220         if (segment != blerg->current_segment)
 221                 blerg_segment_switch(blerg, segment);
 222         int seg_rec = record % RECORDS_PER_SEGMENT;
 223
 224         // Get the position for the new data
 225         FILE *datafile = fdopen(dup(blerg->data_fd), "a");
 226         fseek(datafile, 0, SEEK_END);
 227         int curpos = ftell(datafile);
 228         fclose(datafile);
 229
 230         int bytes = 0;
 231         do {
 232                 int n = write(blerg->data_fd, data + bytes, len);
 233                 if (n == -1) {
 234                         perror("Could not write data");
 235                         // Truncate anything we may have written
 236                         ftruncate(blerg->data_fd, curpos);
 237                         return -1;
 238                 }
 239                 bytes += n;
 240         } while (bytes < len);
 241         blerg->index[seg_rec].flags = 0x0001;
 242         blerg->index[seg_rec].offset = curpos;
 243         blerg->index[seg_rec].length = len;
 244
 245         tag_scan(blerg->name, data, len, record);
 246
 247         flock(blerg->data_fd, LOCK_UN);
 248         flock(blerg->index_fd, LOCK_UN);
 249
 250         return record;
 251 }
 252
 253 int blerg_fetch(struct blerg *blerg, int record, char **data, int *length) {
 254         if (record < 0) {
 255                 printf("Invalid record\n");
 256                 return 0;
 257         }
 258
 259         int segment = record / RECORDS_PER_SEGMENT;
 260         if (segment != blerg->current_segment)
 261                 blerg_segment_switch(blerg, segment);
 262         int seg_rec = record % RECORDS_PER_SEGMENT;
 263
 264         if ((blerg->index[seg_rec].flags & 0x1) == 0) {
 265                 printf("Invalid record\n");
 266                 return 0;
 267         }
 268
 269         int rec_offset = blerg->index[seg_rec].offset;
 270         int rec_length = blerg->index[seg_rec].length;
 271         if (rec_offset >= blerg->data_size) {
 272                 // We're accessing an out-of-bounds record in our mmap.
 273                 // Recheck size and remap.
 274                 struct stat st;
 275                 fstat(blerg->data_fd, &st);
 276                 blerg->data_size = st.st_size;
 277                 if (rec_offset > blerg->data_size) {
 278                         printf("Record offset outside of data!?");
 279                         return 0;
 280                 }
 281
 282                 munmap(blerg->data, blerg->data_size);
 283                 blerg->data = (char *) mmap(NULL, blerg->data_size, PROT_READ, MAP_SHARED, blerg->data_fd, 0);
 284                 if (blerg->data == MAP_FAILED) {
 285                         perror("Could not remap data");
 286                         return 0;
 287                 }
 288         }
 289
 290         *data = malloc(rec_length);
 291         if (*data == NULL) {
 292                 perror("Could not allocate string in fetch");
 293                 return 0;
 294         }
 295
 296         memcpy(*data, blerg->data + rec_offset, rec_length);
 297
 298         *length = rec_length;
 299
 300         return 1;
 301 }