1e33811c1d29f13f638ff603fa2481f32f3c5a29
[blerg.git] / 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 }