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