Bound import memory with staged streaming

2 files changed, 127 insertions, 101 deletions

diff --git a/src/storage/store.c b/src/storage/store.c
index e7ade28..d49f176 100644
--- a/src/storage/store.c
+++ b/src/storage/store.c
@@ -42,6 +42,17 @@ struct tiny_store
   char last_error[160];
 };
 
+struct tiny_import_transaction
+{
+  tiny_sys_file_t lock;
+  tiny_sys_file_t staging;
+  char *staging_path;
+  char *log_path;
+  char *snapshot_path;
+  uint64_t sequence;
+  bool committed;
+};
+
 static void
 set_error(tiny_store_t *store, const char *message)
 {
@@ -2012,115 +2023,129 @@ tiny_store_health(const tiny_store_t *store, tiny_store_health_t *health)
   snprintf(health->last_error, sizeof(health->last_error), "%s", store->last_error);
 }
 
+void
+tiny_store_import_abort(tiny_import_transaction_t *transaction)
+{
+  if (!transaction)
+    {
+      return;
+    }
+  tiny_sys_file_close(&transaction->staging);
+  tiny_sys_file_close(&transaction->lock);
+  if (!transaction->committed && transaction->staging_path)
+    {
+      (void)tiny_sys_file_remove(transaction->staging_path);
+    }
+  free(transaction->snapshot_path);
+  free(transaction->log_path);
+  free(transaction->staging_path);
+  free(transaction);
+}
+
 tiny_store_status_t
-tiny_store_import(tiny_store_t *store,
-                  const tiny_import_record_t *records,
-                  size_t record_count,
-                  tiny_import_policy_t policy,
-                  size_t *imported_count,
-                  size_t *rejected_count)
+tiny_store_import_begin(tiny_import_transaction_t **output, const char *directory)
 {
-  tiny_store_entry_t *saved_entries;
-  size_t saved_count, saved_capacity, index;
-  tiny_rebuild_report_t report;
-  tiny_store_status_t status = TINY_STORE_OK;
-  if (!store || !records || !imported_count || !rejected_count)
+  tiny_import_transaction_t *transaction;
+  char *lock_path;
+  unsigned char generation[16];
+  tiny_sys_status_t status;
+  if (!output || !directory || !directory[0])
     {
       return TINY_STORE_IO_ERROR;
     }
-  *imported_count = 0;
-  *rejected_count = 0;
-  saved_entries = store->entries;
-  saved_count = store->entry_count;
-  saved_capacity = store->entry_capacity;
-  store->entries = NULL;
-  store->entry_count = 0;
-  store->entry_capacity = 0;
-  for (index = 0; index < record_count && status == TINY_STORE_OK; ++index)
+  *output = NULL;
+  transaction = calloc(1, sizeof(*transaction));
+  if (!transaction)
     {
-      const tiny_import_record_t *record = &records[index];
-      tiny_store_entry_t *existing = find_entry(store, record->message.message_id);
-      if (existing != NULL && policy == TINY_IMPORT_REJECT && !record->deleted
-          && !existing->deleted)
-        {
-          (*rejected_count)++;
-          continue;
-        }
-      if (existing != NULL)
-        {
-          tiny_message_t replacement;
-          tiny_message_init(&replacement);
-          if (tiny_message_copy(&replacement, &record->message) != 0)
-            {
-              status = TINY_STORE_NO_MEMORY;
-              break;
-            }
-          tiny_message_free(&existing->message);
-          existing->deleted = record->deleted;
-          existing->sequence = (uint64_t)record->message.timestamp_ms;
-          existing->message = replacement;
-          (*imported_count)++;
-        }
-      else
-        {
-          tiny_message_t copied;
-          tiny_message_init(&copied);
-          if (!ensure_entries(store) || tiny_message_copy(&copied, &record->message) != 0)
-            {
-              tiny_message_free(&copied);
-              status = TINY_STORE_NO_MEMORY;
-              break;
-            }
-          memset(
-              &store->entries[store->entry_count], 0, sizeof(store->entries[store->entry_count]));
-          store->entries[store->entry_count].deleted = record->deleted;
-          store->entries[store->entry_count].sequence = (uint64_t)record->message.timestamp_ms;
-          store->entries[store->entry_count].message = copied;
-          store->entry_count++;
-          (*imported_count)++;
-        }
+      return TINY_STORE_NO_MEMORY;
     }
-  if (status == TINY_STORE_OK)
+  tiny_sys_file_init(&transaction->lock);
+  tiny_sys_file_init(&transaction->staging);
+  transaction->staging_path = tiny_sys_path_join(directory, "active.import.tmp");
+  transaction->log_path = tiny_sys_path_join(directory, "active.log");
+  transaction->snapshot_path = tiny_sys_path_join(directory, "snapshot.bin");
+  lock_path = tiny_sys_path_join(directory, "store.lock");
+  if (!transaction->staging_path || !transaction->log_path || !transaction->snapshot_path
+      || !lock_path)
     {
-      store->update_records = 0;
-      store->dead_bytes = 0;
-      store->snapshot_valid = false;
-      store->recovered_tail = false;
-      status = tiny_store_rebuild(store, &report);
-      if (status != TINY_STORE_OK)
-        {
-          for (index = 0; index < store->entry_count; ++index)
-            {
-              tiny_message_free(&store->entries[index].message);
-            }
-          free(store->entries);
-          store->entries = saved_entries;
-          store->entry_count = saved_count;
-          store->entry_capacity = saved_capacity;
-        }
-      else
-        {
-          for (index = 0; index < saved_count; ++index)
-            {
-              tiny_message_free(&saved_entries[index].message);
-            }
-          free(saved_entries);
-        }
+      free(lock_path);
+      tiny_store_import_abort(transaction);
+      return TINY_STORE_NO_MEMORY;
     }
-  else
+  if (tiny_sys_directory_create(directory) != TINY_SYS_OK)
     {
-      for (index = 0; index < store->entry_count; ++index)
-        {
-          tiny_message_free(&store->entries[index].message);
-        }
-      free(store->entries);
-      store->entries = saved_entries;
-      store->entry_count = saved_count;
-      store->entry_capacity = saved_capacity;
+      free(lock_path);
+      tiny_store_import_abort(transaction);
+      return TINY_STORE_IO_ERROR;
+    }
+  status = tiny_sys_file_open(&transaction->lock, lock_path, TINY_SYS_FILE_LOCK);
+  free(lock_path);
+  if (status != TINY_SYS_OK)
+    {
+      tiny_store_import_abort(transaction);
+      return status == TINY_SYS_BUSY ? TINY_STORE_BUSY : TINY_STORE_IO_ERROR;
+    }
+  if (tiny_sys_file_open(&transaction->staging, transaction->staging_path, TINY_SYS_FILE_WRITE_NEW)
+          != TINY_SYS_OK
+      || !tiny_sys_random_bytes(generation, sizeof(generation))
+      || append_record_to(
+             &transaction->staging, RECORD_MARKER, 0, generation, sizeof(generation), false)
+             != TINY_STORE_OK)
+    {
+      tiny_store_import_abort(transaction);
+      return TINY_STORE_IO_ERROR;
+    }
+  *output = transaction;
+  return TINY_STORE_OK;
+}
+
+tiny_store_status_t
+tiny_store_import_append(tiny_import_transaction_t *transaction, const tiny_import_record_t *record)
+{
+  unsigned char *payload = NULL;
+  uint32_t length = 0;
+  tiny_store_status_t status;
+  if (!transaction || !record || transaction->committed)
+    {
+      return TINY_STORE_IO_ERROR;
+    }
+  transaction->sequence++;
+  if (record->deleted)
+    {
+      return append_record_to(&transaction->staging,
+                              RECORD_DEL,
+                              transaction->sequence,
+                              (const unsigned char *)record->message.message_id,
+                              TINY_MESSAGE_ID_LENGTH,
+                              false);
     }
+  if (!serialize_put(&record->message, &payload, &length))
+    {
+      return TINY_STORE_NO_MEMORY;
+    }
+  status = append_record_to(
+      &transaction->staging, RECORD_PUT, transaction->sequence, payload, length, false);
+  free(payload);
   return status;
 }
 
+tiny_store_status_t
+tiny_store_import_commit(tiny_import_transaction_t *transaction)
+{
+  if (!transaction || transaction->committed || !tiny_sys_file_flush(&transaction->staging))
+    {
+      return TINY_STORE_IO_ERROR;
+    }
+  tiny_sys_file_close(&transaction->staging);
+  if (!tiny_sys_path_replace(transaction->staging_path, transaction->log_path))
+    {
+      return TINY_STORE_IO_ERROR;
+    }
+  transaction->committed = true;
+  (void)tiny_sys_file_remove(transaction->snapshot_path);
+  return TINY_STORE_OK;
+}
+
 size_t
 tiny_store_entry_count(const tiny_store_t *store)
 {
diff --git a/src/storage/store.h b/src/storage/store.h
index 9efb02a..78c8cbe 100644
--- a/src/storage/store.h
+++ b/src/storage/store.h
@@ -26,6 +26,7 @@ typedef struct tiny_store_entry
 } tiny_store_entry_t;
 
 typedef struct tiny_store tiny_store_t;
+typedef struct tiny_import_transaction tiny_import_transaction_t;
 
 typedef struct tiny_rebuild_report
 {
@@ -94,12 +95,12 @@ typedef struct tiny_import_record
 } tiny_import_record_t;
 
 tiny_store_status_t tiny_store_rebuild(tiny_store_t *store, tiny_rebuild_report_t *report);
-tiny_store_status_t tiny_store_import(tiny_store_t *store,
-                                      const tiny_import_record_t *records,
-                                      size_t record_count,
-                                      tiny_import_policy_t policy,
-                                      size_t *imported_count,
-                                      size_t *rejected_count);
+tiny_store_status_t tiny_store_import_begin(tiny_import_transaction_t **transaction,
+                                            const char *directory);
+tiny_store_status_t tiny_store_import_append(tiny_import_transaction_t *transaction,
+                                             const tiny_import_record_t *record);
+tiny_store_status_t tiny_store_import_commit(tiny_import_transaction_t *transaction);
+void tiny_store_import_abort(tiny_import_transaction_t *transaction);
 void tiny_store_health(const tiny_store_t *store, tiny_store_health_t *health);
 
 size_t tiny_store_entry_count(const tiny_store_t *store);