database/database.c

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