Bound import memory with staged streaming

1 files changed, 351 insertions, 85 deletions

diff --git a/src/cli/import.c b/src/cli/import.c
index f4688ab..ee4d1cf 100644
--- a/src/cli/import.c
+++ b/src/cli/import.c
@@ -3,6 +3,7 @@
 #include "domain/message.h"
 #include "mcp/json.h"
 #include "storage/store.h"
+#include "sys/system.h"
 
 #include "mjson.h"
 
@@ -11,55 +12,162 @@
 #include <string.h>
 #include <limits.h>
 
-#define IMPORT_MAX_LINE TINY_CLI_LINE_MAX
-#define IMPORT_MAX_RECORDS 100000U
-#define IMPORT_MAX_DECODED_BYTES (UINT64_C(4) * 1024U * 1024U * 1024U)
+#include "base/allocator.h"
 
-typedef struct import_record_list
+#ifndef TINY_IMPORT_MAX_LINE
+#define TINY_IMPORT_MAX_LINE TINY_CLI_LINE_MAX
+#endif
+#ifndef TINY_IMPORT_MAX_RECORDS
+#define TINY_IMPORT_MAX_RECORDS 100000U
+#endif
+#ifndef TINY_IMPORT_MAX_DECODED_BYTES
+#define TINY_IMPORT_MAX_DECODED_BYTES (UINT64_C(4) * 1024U * 1024U * 1024U)
+#endif
+#define IMPORT_READ_BUFFER 65536U
+
+typedef struct import_compact_entry
+{
+  char message_id[TINY_MESSAGE_ID_LENGTH + 1];
+  size_t selected_record;
+  bool deleted;
+} import_compact_entry_t;
+
+typedef struct import_compact_index
 {
-  tiny_import_record_t *items;
+  import_compact_entry_t *entries;
   size_t count;
   size_t capacity;
-} import_record_list_t;
+  size_t *slots;
+  size_t slot_capacity;
+} import_compact_index_t;
+
+typedef struct import_reader
+{
+  tiny_sys_file_t file;
+  unsigned char input[IMPORT_READ_BUFFER];
+  size_t input_offset;
+  size_t input_length;
+  char *line;
+  size_t line_length;
+  bool eof;
+} import_reader_t;
 
 static void
-import_record_list_free(import_record_list_t *list)
+import_compact_index_free(import_compact_index_t *index)
 {
-  size_t index;
-  if (!list)
+  if (!index)
     {
       return;
     }
-  for (index = 0; index < list->count; ++index)
+  free(index->slots);
+  free(index->entries);
+  memset(index, 0, sizeof(*index));
+}
+
+static uint64_t
+import_id_hash(const char *value)
+{
+  uint64_t hash = UINT64_C(1469598103934665603);
+  while (*value)
     {
-      tiny_message_free(&list->items[index].message);
+      hash ^= (unsigned char)*value++;
+      hash *= UINT64_C(1099511628211);
     }
-  free(list->items);
-  list->items = NULL;
-  list->count = 0;
-  list->capacity = 0;
+  return hash;
 }
 
 static bool
-import_record_list_push(import_record_list_t *list)
+import_compact_index_rehash(import_compact_index_t *index, size_t capacity)
 {
-  tiny_import_record_t *resized;
-  size_t capacity;
-  if (list->count < list->capacity)
+  size_t *slots;
+  size_t position;
+  if (capacity > SIZE_MAX / sizeof(*slots) || !(slots = malloc(capacity * sizeof(*slots))))
     {
-      return true;
+      return false;
     }
-  capacity = list->capacity == 0 ? 32 : list->capacity * 2;
-  resized = realloc(list->items, capacity * sizeof(*resized));
-  if (!resized)
+  for (position = 0; position < capacity; ++position)
     {
-      return false;
+      slots[position] = SIZE_MAX;
     }
-  list->items = resized;
-  list->capacity = capacity;
+  for (position = 0; position < index->count; ++position)
+    {
+      size_t slot = (size_t)import_id_hash(index->entries[position].message_id) & (capacity - 1U);
+      while (slots[slot] != SIZE_MAX)
+        {
+          slot = (slot + 1U) & (capacity - 1U);
+        }
+      slots[slot] = position;
+    }
+  free(index->slots);
+  index->slots = slots;
+  index->slot_capacity = capacity;
   return true;
 }
 
+static import_compact_entry_t *
+import_compact_index_find(const import_compact_index_t *index, const char *message_id)
+{
+  size_t slot;
+  if (!index->slot_capacity)
+    {
+      return NULL;
+    }
+  slot = (size_t)import_id_hash(message_id) & (index->slot_capacity - 1U);
+  while (index->slots[slot] != SIZE_MAX)
+    {
+      import_compact_entry_t *entry = &index->entries[index->slots[slot]];
+      if (strcmp(entry->message_id, message_id) == 0)
+        {
+          return entry;
+        }
+      slot = (slot + 1U) & (index->slot_capacity - 1U);
+    }
+  return NULL;
+}
+
+static import_compact_entry_t *
+import_compact_index_add(import_compact_index_t *index, const char *message_id)
+{
+  import_compact_entry_t *resized;
+  size_t capacity;
+  if (index->count == index->capacity)
+    {
+      capacity = index->capacity ? index->capacity * 2U : 32U;
+      if (capacity < index->capacity || capacity > SIZE_MAX / sizeof(*resized)
+          || !(resized = realloc(index->entries, capacity * sizeof(*resized))))
+        {
+          return NULL;
+        }
+      index->entries = resized;
+      index->capacity = capacity;
+    }
+  if (!index->slot_capacity || (index->count + 1U) * 10U > index->slot_capacity * 7U)
+    {
+      capacity = index->slot_capacity ? index->slot_capacity * 2U : 64U;
+      if (capacity < index->slot_capacity || !import_compact_index_rehash(index, capacity))
+        {
+          return NULL;
+        }
+    }
+  memset(&index->entries[index->count], 0, sizeof(index->entries[index->count]));
+  memcpy(index->entries[index->count].message_id, message_id, TINY_MESSAGE_ID_LENGTH + 1U);
+  {
+    size_t slot = (size_t)import_id_hash(message_id) & (index->slot_capacity - 1U);
+    while (index->slots[slot] != SIZE_MAX)
+      {
+        slot = (slot + 1U) & (index->slot_capacity - 1U);
+      }
+    index->slots[slot] = index->count;
+  }
+  index->count++;
+  if (index->count > TINY_IMPORT_MAX_RECORDS)
+    {
+      index->count--;
+      return NULL;
+    }
+  return &index->entries[index->count - 1U];
+}
+
 static bool
 import_get_string(const char *json, int len, const char *path, char *out, size_t out_size)
 {
@@ -518,7 +626,8 @@ import_manifest_valid(const char *line, int len)
 {
   char type_name[64];
   uint64_t version = 0;
-  if (!import_get_string(line, len, "$.type", type_name, sizeof(type_name))
+  if (mjson_find(line, len, "$", NULL, NULL) != MJSON_TOK_OBJECT
+      || !import_get_string(line, len, "$.type", type_name, sizeof(type_name))
       || strcmp(type_name, "tiny-convmem-manifest") != 0
       || !import_get_u64(line, len, "$.format_version", &version))
     {
@@ -532,7 +641,8 @@ import_parse_del_record(const char *line, int len, tiny_import_record_t *record)
 {
   char type_name[32];
   tiny_message_init(&record->message);
-  if (!import_get_string(line, len, "$.type", type_name, sizeof(type_name))
+  if (mjson_find(line, len, "$", NULL, NULL) != MJSON_TOK_OBJECT
+      || !import_get_string(line, len, "$.type", type_name, sizeof(type_name))
       || strcmp(type_name, "DEL") != 0
       || !import_get_string(
           line, len, "$.message_id", record->message.message_id, sizeof(record->message.message_id))
@@ -553,7 +663,8 @@ import_parse_message_record(const char *line, int len, tiny_import_record_t *rec
   size_t content_length = 0;
   tiny_message_init(&record->message);
   record->deleted = false;
-  if (!import_get_string(
+  if (mjson_find(line, len, "$", NULL, NULL) != MJSON_TOK_OBJECT
+      || !import_get_string(
           line, len, "$.message_id", record->message.message_id, sizeof(record->message.message_id))
       || !import_get_string(line, len, "$.role", role_name, sizeof(role_name))
       || !import_get_allocated_string(
@@ -615,35 +726,129 @@ import_record_decoded_bytes(const tiny_import_record_t *record)
 }
 
 static bool
-import_parse_stream(FILE *file, import_record_list_t *records)
+import_reader_open(import_reader_t *reader, const char *path)
 {
-  char *line = malloc((size_t)IMPORT_MAX_LINE + 2U);
-  bool manifest_seen = false;
-  uint64_t decoded_total = 0;
-  if (!line)
+  memset(reader, 0, sizeof(*reader));
+  tiny_sys_file_init(&reader->file);
+  reader->line = malloc((size_t)TINY_IMPORT_MAX_LINE + 1U);
+  return reader->line && tiny_sys_file_open(&reader->file, path, TINY_SYS_FILE_READ) == TINY_SYS_OK;
+}
+
+static void
+import_reader_close(import_reader_t *reader)
+{
+  tiny_sys_file_close(&reader->file);
+  free(reader->line);
+  memset(reader, 0, sizeof(*reader));
+}
+
+static bool
+import_reader_rewind(import_reader_t *reader)
+{
+  reader->input_offset = 0;
+  reader->input_length = 0;
+  reader->line_length = 0;
+  reader->eof = false;
+  return tiny_sys_file_seek(&reader->file, 0, TINY_SYS_SEEK_BEGIN);
+}
+
+static int
+import_reader_next(import_reader_t *reader, const char **line, size_t *length)
+{
+  reader->line_length = 0;
+  while (true)
     {
-      return false;
+      unsigned char value;
+      if (reader->input_offset == reader->input_length)
+        {
+          uint32_t actual = 0;
+          if (reader->eof)
+            {
+              if (reader->line_length == 0)
+                {
+                  return 0;
+                }
+              break;
+            }
+          if (!tiny_sys_file_read(&reader->file, reader->input, sizeof(reader->input), &actual))
+            {
+              return -1;
+            }
+          reader->input_offset = 0;
+          reader->input_length = actual;
+          reader->eof = actual == 0;
+          continue;
+        }
+      value = reader->input[reader->input_offset++];
+      if (value == '\n')
+        {
+          break;
+        }
+      if (value == 0 || reader->line_length >= TINY_IMPORT_MAX_LINE)
+        {
+          return -1;
+        }
+      reader->line[reader->line_length++] = (char)value;
     }
-  while (fgets(line, (int)IMPORT_MAX_LINE + 2, file) != NULL)
+  if (reader->line_length > 0 && reader->line[reader->line_length - 1U] == '\r')
     {
-      size_t length = strlen(line);
-      tiny_import_record_t record;
-      uint64_t decoded;
-      if (length == 0 || length > IMPORT_MAX_LINE
-          || (length == IMPORT_MAX_LINE + 1U && line[length - 1] != '\n')
-          || (line[length - 1] != '\n' && !feof(file)))
+      reader->line_length--;
+    }
+  reader->line[reader->line_length] = '\0';
+  *line = reader->line;
+  *length = reader->line_length;
+  return 1;
+}
+
+static bool
+import_record_references_valid(const tiny_import_record_t *record,
+                               const import_compact_index_t *index)
+{
+  size_t position;
+  if (record->deleted)
+    {
+      return true;
+    }
+  for (position = 0; position < record->message.metadata.relation_count; ++position)
+    {
+      import_compact_entry_t *target = import_compact_index_find(
+          index, record->message.metadata.relations[position].message_id);
+      if (!target || target->deleted)
         {
-          free(line);
           return false;
         }
-      if (length > 0 && line[length - 1] == '\n')
-        {
-          line[--length] = '\0';
-        }
-      if (length > 0 && line[length - 1] == '\r')
+    }
+  for (position = 0; position < record->message.metadata.source_count; ++position)
+    {
+      import_compact_entry_t *target
+          = import_compact_index_find(index, record->message.metadata.source_message_ids[position]);
+      if (!target || target->deleted)
         {
-          line[--length] = '\0';
+          return false;
         }
+    }
+  return true;
+}
+
+static bool
+import_validate_pass(import_reader_t *reader,
+                     tiny_import_policy_t policy,
+                     import_compact_index_t *index,
+                     size_t *accepted,
+                     size_t *rejected)
+{
+  const char *line;
+  size_t length, record_number = 0;
+  uint64_t decoded_total = 0;
+  bool manifest_seen = false;
+  int next;
+  *accepted = 0;
+  *rejected = 0;
+  while ((next = import_reader_next(reader, &line, &length)) > 0)
+    {
+      tiny_import_record_t record;
+      import_compact_entry_t *entry;
+      uint64_t decoded;
       if (length == 0)
         {
           continue;
@@ -652,53 +857,111 @@ import_parse_stream(FILE *file, import_record_list_t *records)
         {
           if (!import_manifest_valid(line, (int)length))
             {
-              free(line);
               return false;
             }
           manifest_seen = true;
           continue;
         }
-      if (records->count >= IMPORT_MAX_RECORDS)
-        {
-          free(line);
-          return false;
-        }
-      if (!import_parse_del_record(line, (int)length, &record)
-          && !import_parse_message_record(line, (int)length, &record))
+      if (record_number >= TINY_IMPORT_MAX_RECORDS
+          || (!import_parse_del_record(line, (int)length, &record)
+              && !import_parse_message_record(line, (int)length, &record)))
         {
-          free(line);
           return false;
         }
       decoded = import_record_decoded_bytes(&record);
-      if (decoded > IMPORT_MAX_DECODED_BYTES - decoded_total || !import_record_list_push(records))
+      if (decoded > TINY_IMPORT_MAX_DECODED_BYTES - decoded_total)
         {
           tiny_message_free(&record.message);
-          free(line);
           return false;
         }
       decoded_total += decoded;
-      records->items[records->count++] = record;
+      entry = import_compact_index_find(index, record.message.message_id);
+      if (!entry)
+        {
+          entry = import_compact_index_add(index, record.message.message_id);
+          if (!entry)
+            {
+              tiny_message_free(&record.message);
+              return false;
+            }
+        }
+      else if (policy == TINY_IMPORT_REJECT && !record.deleted && !entry->deleted)
+        {
+          (*rejected)++;
+          tiny_message_free(&record.message);
+          record_number++;
+          continue;
+        }
+      entry->selected_record = record_number;
+      entry->deleted = record.deleted;
+      (*accepted)++;
+      tiny_message_free(&record.message);
+      record_number++;
     }
-  if (ferror(file))
+  return next == 0 && manifest_seen;
+}
+
+static tiny_store_status_t
+import_staging_pass(import_reader_t *reader,
+                    const import_compact_index_t *index,
+                    tiny_import_transaction_t *transaction)
+{
+  const char *line;
+  size_t length, record_number = 0;
+  bool manifest_seen = false;
+  int next;
+  while ((next = import_reader_next(reader, &line, &length)) > 0)
     {
-      free(line);
-      return false;
+      tiny_import_record_t record;
+      import_compact_entry_t *entry;
+      tiny_store_status_t status;
+      if (length == 0)
+        {
+          continue;
+        }
+      if (!manifest_seen)
+        {
+          manifest_seen = true;
+          continue;
+        }
+      if (!import_parse_del_record(line, (int)length, &record)
+          && !import_parse_message_record(line, (int)length, &record))
+        {
+          return TINY_STORE_CORRUPT;
+        }
+      entry = import_compact_index_find(index, record.message.message_id);
+      if (entry && entry->selected_record == record_number && !record.deleted)
+        {
+          if (!import_record_references_valid(&record, index))
+            {
+              tiny_message_free(&record.message);
+              return TINY_STORE_CORRUPT;
+            }
+          status = tiny_store_import_append(transaction, &record);
+          if (status != TINY_STORE_OK)
+            {
+              tiny_message_free(&record.message);
+              return status;
+            }
+        }
+      tiny_message_free(&record.message);
+      record_number++;
     }
-  free(line);
-  return manifest_seen;
+  return next == 0 ? TINY_STORE_OK : TINY_STORE_IO_ERROR;
 }
 
 int
-tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options)
+tiny_cli_import(const tiny_cli_options_t *options)
 {
-  import_record_list_t records = { NULL, 0, 0 };
+  import_compact_index_t index = { 0 };
+  import_reader_t reader;
+  tiny_import_transaction_t *transaction = NULL;
   size_t imported = 0, rejected = 0;
-  tiny_store_status_t status;
+  tiny_store_status_t status = TINY_STORE_CORRUPT;
   tiny_import_policy_t policy = TINY_IMPORT_OVERWRITE;
-  FILE *file = NULL;
   int result = 1;
 
-  if (!options || !options->import_path)
+  if (!options || !options->data_dir || !options->import_path)
     {
       fprintf(stderr, "import: path is required\n");
       return 1;
@@ -721,29 +984,30 @@ tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options)
         }
     }
 
-  if (fopen_s(&file, options->import_path, "rb") != 0 || !file)
+  if (!import_reader_open(&reader, options->import_path))
     {
       fprintf(stderr, "import: unable to open '%s'\n", options->import_path);
+      import_reader_close(&reader);
       return 1;
     }
-
-  if (!import_parse_stream(file, &records))
+  if (!import_validate_pass(&reader, policy, &index, &imported, &rejected)
+      || !import_reader_rewind(&reader))
     {
       fprintf(stderr, "import: invalid JSONL in '%s'\n", options->import_path);
-      fclose(file);
-      import_record_list_free(&records);
-      return 1;
     }
-  if (fclose(file) != 0)
+  else
     {
-      fprintf(stderr, "import: unable to close '%s'\n", options->import_path);
-      import_record_list_free(&records);
-      return 1;
+      status = tiny_store_import_begin(&transaction, options->data_dir);
+      if (status == TINY_STORE_OK)
+        {
+          status = import_staging_pass(&reader, &index, transaction);
+        }
+      if (status == TINY_STORE_OK)
+        {
+          status = tiny_store_import_commit(transaction);
+        }
     }
 
-  status = tiny_store_import(store, records.items, records.count, policy, &imported, &rejected);
-  import_record_list_free(&records);
-
   if (status == TINY_STORE_OK)
     {
       printf("{\"imported\":%zu,\"rejected\":%zu}\n", imported, rejected);
@@ -762,6 +1026,8 @@ tiny_cli_import(tiny_store_t *store, const tiny_cli_options_t *options)
     {
       fprintf(stderr, "import: failed\n");
     }
-
+  tiny_store_import_abort(transaction);
+  import_reader_close(&reader);
+  import_compact_index_free(&index);
   return result;
 }