database/database.c

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