Miscellaneous improvements

6 files changed, 33 insertions, 14 deletions

diff --git a/bin/nfcollect.c b/bin/nfcollect.c
index a36d177..3ffd3fa 100644
--- a/bin/nfcollect.c
+++ b/bin/nfcollect.c
@@ -24,6 +24,7 @@
 // SOFTWARE.
 
 #include "collect.h"
+#include "sql.h"
 #include "util.h"
 #include <dirent.h>
 #include <fcntl.h>
@@ -43,10 +44,11 @@ const char *help_text =
     "Options:\n"
     "  -c --compression=<algo>      compression algorithm to use (default: no "
     "compression)\n"
-    "  -d --storage_file=<filename>    sqlite database storage file\n"
+    "  -d --storage=<filename>         sqlite database storage file\n"
     "  -h --help                       print this help\n"
-    "  -g --nflog-group=<id>           the group id to collect\n"
+    "  -g --nflog_group=<id>           the group id to collect\n"
     "  -s --storage_size=<max DB size> maximum DB size in MiB\n"
+    "  -V --vacuum                     vacuum the database on startup\n"
     "  -v --version                    print version information\n"
     "\n";
 
@@ -63,18 +65,20 @@ int main(int argc, char *argv[]) {
     Global g;
     int nflog_group_id = -1;
     char *compression_flag = NULL, *storage = NULL;
+    bool do_vacuum = false;
 
     struct option longopts[] = {/* name, has_args, flag, val */
                                 {"nflog_group", required_argument, NULL, 'g'},
                                 {"storage", required_argument, NULL, 'd'},
                                 {"storage_size", required_argument, NULL, 's'},
                                 {"compression", optional_argument, NULL, 'z'},
+                                {"vacuum", optional_argument, NULL, 'V'},
                                 {"help", no_argument, NULL, 'h'},
                                 {"version", no_argument, NULL, 'v'},
                                 {0, 0, 0, 0}};
 
     int opt;
-    while ((opt = getopt_long(argc, argv, "c:g:d:s:hvp:", longopts, NULL)) !=
+    while ((opt = getopt_long(argc, argv, "c:g:d:s:hVvp:", longopts, NULL)) !=
            -1) {
         switch (opt) {
         case 'h':
@@ -128,7 +132,9 @@ int main(int argc, char *argv[]) {
 
     pthread_t worker;
     State *state;
-    INFO(PACKAGE ": storing in file '%s', capped by %d MiB", g.storage_file,
+    INFO(PACKAGE
+         ": storing in file '%s' (current size: %.2f MB), capped by %d MiB",
+         g.storage_file, (float)g.storage_consumed / 1024.0 / 1024.0,
          storage_size);
     INFO(PACKAGE ": workers started, entries per block = %d", g.max_nr_entries);
 
diff --git a/bin/nfextract.c b/bin/nfextract.c
index 0ca112e..038c889 100644
--- a/bin/nfextract.c
+++ b/bin/nfextract.c
@@ -59,9 +59,9 @@ const char *help_text =
     "  -d --storage=<dirname>     sqlite storage file\n"
     "  -h --help                  print this help\n"
     "  -v --version               print version information\n"
-    "  -s --since                 start showing entries on or newer than the "
+    "  -s --since=<date>          start showing entries on or newer than the "
     "specified date (format: " DATE_FORMAT_HUMAN ")\n"
-    "  -u --until                 stop showing entries on or older than the "
+    "  -u --until=<date>          stop showing entries on or older than the "
     "specified date (format: " DATE_FORMAT_HUMAN ")\n"
     "\n";
 
diff --git a/include/main.h b/include/main.h
index c2407e1..f64fe46 100644
--- a/include/main.h
+++ b/include/main.h
@@ -85,9 +85,9 @@
 #define unlikely(x) __builtin_expect((x), 0)
 
 #ifdef __GNUC__
-#define UNUSED_FUNCTION(x) __attribute__((__unused__)) UNUSED_ ## x
+#define UNUSED_FUNCTION(x) __attribute__((__unused__)) UNUSED_##x
 #else
-#define UNUSED_FUNCTION(x) UNUSED_ ## x
+#define UNUSED_FUNCTION(x) UNUSED_##x
 #endif
 
 enum CompressionType { COMPRESS_NONE, COMPRESS_LZ4, COMPRESS_ZSTD };
@@ -129,8 +129,8 @@ typedef struct _nfl_nl_t {
 typedef struct _Global {
     uint16_t nl_group_id;
 
-    uint32_t storage_budget;
-    uint32_t storage_consumed;
+    int64_t storage_budget;
+    int64_t storage_consumed;
     pthread_mutex_t storage_consumed_lock;
 
     uint32_t max_nr_entries;
diff --git a/include/util.h b/include/util.h
index f1fd965..3d92764 100644
--- a/include/util.h
+++ b/include/util.h
@@ -2,6 +2,7 @@
 #define UTIL_H
 #include "main.h"
 int check_basedir_exist(const char *storage);
+int check_file_size(const char *storage);
 int check_file_exist(const char *storage);
 enum CompressionType get_compression(const char *flag);
 
diff --git a/lib/sql.c b/lib/sql.c
index 99f97ef..1b4df89 100644
--- a/lib/sql.c
+++ b/lib/sql.c
@@ -232,7 +232,7 @@ int db_delete_oldest_bytes(sqlite3 *db, int64_t bytes) {
     size_t bufsize = 1024;
     char *buf = malloc(bufsize);
 
-    while (bytes >= 0) {
+    while (true) {
         rc = sqlite3_step(stmt);
         if (rc == SQLITE_DONE)
             break;
@@ -240,8 +240,12 @@ int db_delete_oldest_bytes(sqlite3 *db, int64_t bytes) {
         sqlite3_int64 index = sqlite3_column_int64(stmt, 2);
         int size = sqlite3_column_int(stmt, 0);
 
+        bytes -= size;
+        if (bytes <= 0)
+            break;
+
         char _buf[22];
-        sprintf(_buf, count ? "%lld" : ",%lld", index);
+        sprintf(_buf, count ? ",%lld" : "%lld", index);
         while (strlen(_buf) + strlen(buf) + 2 >= bufsize) {
             bufsize *= 2;
             char *__buf = malloc(bufsize);
diff --git a/lib/util.c b/lib/util.c
index e80954b..ef5238f 100644
--- a/lib/util.c
+++ b/lib/util.c
@@ -8,6 +8,12 @@ int check_file_exist(const char *storage) {
     return access(storage, F_OK) != -1;
 }
 
+int check_file_size(const char *storage) {
+    struct stat st;
+    stat(storage, &st);
+    return st.st_size;
+}
+
 int check_basedir_exist(const char *storage) {
     char *_storage = strdup(storage);
     char *basedir = dirname(_storage);
@@ -28,7 +34,9 @@ enum CompressionType get_compression(const char *flag) {
     } else if (!strcmp(flag, "lz4")) {
         return COMPRESS_LZ4;
     } else {
-        fprintf(stderr, "Unknown compression algorithm: %s\n", flag);
-        return 0;
+        FATAL("Unknown compression algorithm: %s\n", flag);
+        exit(1);
     }
+
+    return 0;
 }