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, sizeof(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 struct blerg *blerg_open(const char *name) {
 112         int namelen = strlen(name);
 113         char filename[512];
 114         struct stat st;
 115         uint64_t sequence;
 116
 117         if (namelen > 32) {
 118                 perror("Name too long");
 119                 return NULL;
 120         }
 121         struct blerg *blerg = malloc(sizeof(struct blerg));
 122         if (!blerg) {
 123                 perror("Cannot allocate memory for blerg");
 124                 goto open_failed_blerg_malloc;
 125         }
 126         blerg->meta_fd = blerg->index_fd = blerg->data_fd = -1;
 127         blerg->meta = NULL;
 128         blerg->index = NULL;
 129         blerg->data = NULL;
 130
 131         // Make the directory if it doesn't exist
 132         blerg->base_path = malloc(512);
 133         snprintf(blerg->base_path, 512, "%s/%s", DATA_PATH, name);
 134         if (access(blerg->base_path, F_OK) == -1)
 135                 mkdir(blerg->base_path, 0755);
 136
 137         // Open and map metadata
 138         snprintf(filename, 512, "%s/meta", blerg->base_path);
 139         blerg->meta_fd = open(filename, O_RDWR | O_CREAT, 0600);
 140         if (blerg->meta_fd == -1) {
 141                 perror("Could not open metadata");
 142                 goto open_failed_meta_open;
 143         }
 144         fstat(blerg->meta_fd, &st);
 145         if (st.st_size == 0) {
 146                 char *buf = (char *) malloc(sizeof(struct meta));
 147                 memset(buf, 0, sizeof(struct meta));
 148                 write(blerg->meta_fd, buf, sizeof(struct meta));
 149                 free(buf);
 150         }
 151         blerg->meta = (struct meta *) mmap(NULL, sizeof(struct meta), PROT_READ | PROT_WRITE, MAP_SHARED, blerg->meta_fd, 0);
 152         if (blerg->meta == MAP_FAILED) {
 153                 perror("Could not map metadata");
 154                 goto open_failed_meta_mmap;
 155         }
 156
 157         // Open and map index and data for the current segment
 158         blerg->current_segment = blerg_get_record_count(blerg) / RECORDS_PER_SEGMENT;
 159         if (!blerg_segment_switch(blerg, blerg->current_segment)) {
 160                 fprintf(stderr, "Could not switch segment\n");
 161                 goto open_failed_segment_switch;
 162         }
 163
 164         return blerg;
 165
 166 open_failed_segment_switch:
 167         munmap((void *)blerg->meta, sizeof(struct meta));
 168 open_failed_meta_mmap:
 169         close(blerg->meta_fd);
 170 open_failed_meta_open:
 171         free(blerg);
 172 open_failed_blerg_malloc:
 173         return NULL;
 174 }
 175
 176 int blerg_close(struct blerg *blerg) {
 177         blerg_segment_close(blerg);
 178         munmap((void *)blerg->meta, sizeof(struct meta));
 179         close(blerg->meta_fd);
 180         free(blerg->base_path);
 181         free(blerg);
 182         return 1;
 183 }
 184
 185 int blerg_store(struct blerg *blerg, const char *data, int len) {
 186         if (len > 65536) {
 187                 printf("len > 64K\n");
 188                 return -1;
 189         }
 190
 191         flock(blerg->index_fd, LOCK_EX);
 192         flock(blerg->data_fd, LOCK_EX);
 193
 194         uint64_t rec = blerg_increment_record_count(blerg);
 195         if (rec == -1) {
 196                 printf("Could not find free record\n");
 197                 return -1;
 198         }
 199         int segment = rec / RECORDS_PER_SEGMENT;
 200         if (segment != blerg->current_segment)
 201                 blerg_segment_switch(blerg, segment);
 202         rec = rec % RECORDS_PER_SEGMENT;
 203
 204         // Get the position for the new data
 205         FILE *datafile = fdopen(dup(blerg->data_fd), "a");
 206         fseek(datafile, 0, SEEK_END);
 207         int curpos = ftell(datafile);
 208         fclose(datafile);
 209
 210         int bytes = 0;
 211         do {
 212                 int n = write(blerg->data_fd, data + bytes, len);
 213                 if (n == -1) {
 214                         perror("Could not write data");
 215                         // Truncate anything we may have written
 216                         ftruncate(blerg->data_fd, curpos);
 217                         return -1;
 218                 }
 219                 bytes += n;
 220         } while (bytes < len);
 221         blerg->index[rec].flags = 0x0001;
 222         blerg->index[rec].offset = curpos;
 223         blerg->index[rec].length = len;
 224
 225         flock(blerg->data_fd, LOCK_UN);
 226         flock(blerg->index_fd, LOCK_UN);
 227
 228         return segment * RECORDS_PER_SEGMENT + rec;
 229 }
 230
 231 int blerg_fetch(struct blerg *blerg, int rec, char **data, int *length) {
 232         if (rec < 0) {
 233                 printf("Invalid record\n");
 234                 return 0;
 235         }
 236
 237         int segment = rec / RECORDS_PER_SEGMENT;
 238         if (segment != blerg->current_segment)
 239                 blerg_segment_switch(blerg, segment);
 240         rec = rec % RECORDS_PER_SEGMENT;
 241
 242         if ((blerg->index[rec].flags & 0x1) == 0) {
 243                 printf("Invalid record\n");
 244                 return 0;
 245         }
 246
 247         int rec_offset = blerg->index[rec].offset;
 248         int rec_length = blerg->index[rec].length;
 249         if (rec_offset >= blerg->data_size) {
 250                 // We're accessing an out-of-bounds record in our mmap.
 251                 // Recheck size and remap.
 252                 struct stat st;
 253                 fstat(blerg->data_fd, &st);
 254                 blerg->data_size = st.st_size;
 255                 if (rec_offset > blerg->data_size) {
 256                         printf("Record offset outside of data!?");
 257                         return 0;
 258                 }
 259
 260                 munmap(blerg->data, blerg->data_size);
 261                 blerg->data = (char *) mmap(NULL, blerg->data_size, PROT_READ, MAP_SHARED, blerg->data_fd, 0);
 262                 if (blerg->data == MAP_FAILED) {
 263                         perror("Could not remap data");
 264                         return 0;
 265                 }
 266         }
 267
 268         *data = malloc(rec_length);
 269         if (*data == NULL) {
 270                 perror("Could not allocate string in fetch");
 271                 return 0;
 272         }
 273
 274         memcpy(*data, blerg->data + rec_offset, rec_length);
 275
 276         *length = rec_length;
 277
 278         return 1;
 279 }