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