Do not lose increments

Indexing threw counts away in two ways, both without a word.

rocksdb was asked to write with low_pri and no_slowdown, which means "cancel
this write when compaction is behind" and "do not wait, report it instead", and
the reported status was discarded in both inc() overloads. Measured with 3
million increments against small write buffers: 2,253,066 arrived, 746,934 were
gone, a quarter of them. The writer no longer deprioritises itself and waits
when it has to, and merge_one() repeats a cancelled write in any case. A
cancelled write was not applied, so repeating it cannot count twice.

The second one is the end of a run. The database has no write ahead log, and
nothing ever closed it, so everything still in a write buffer was lost when the
indexer ended: of 2 million increments in a process that just exited, none were
readable afterwards. close_collocatordb() writes them and closes the database,
and the destructor does the same.

While testing this, two more:

- get_collocators() and get_collocation_scores() finish a collocate when the
  next one starts, so the last one was never returned. Every word was missing
  a collocate, in the API and in the JSON.
- the array both of them return was allocated with c.size() + sizeof c[0]
  instead of c.size() * sizeof c[0], far too small for more than one entry.
  Callers reading beyond the first entry read whatever was behind it. It is
  terminated with an empty entry now, so that they can tell where it ends.

The settings are aimed at what indexing does, a long stream of merge operands
for the same keys: bigger write buffers, several of them combined before they
are written, which collapses the operands early, and compaction that is allowed
to keep up. All of it can be set from the environment, to be able to tune a run
of several days without recompiling. It also explains why more indexing threads
stop helping: rocksdb inserts merge operands one writer at a time.

The test that found all of this is part of the suite now. The expectations of
the writing test were corrected, they described the missing collocate.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Change-Id: I283b073b7f3fd36aa3ce960978e82105b1e44355
diff --git a/CMakeLists.txt b/CMakeLists.txt
index 771e90e..912a49b 100644
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -1,5 +1,5 @@
 cmake_minimum_required(VERSION 3.9)
-project(collocatordb VERSION 1.5.0 DESCRIPTION "Storing and retrieving collocation counts based on RocksDB")
+project(collocatordb VERSION 1.6.0 DESCRIPTION "Storing and retrieving collocation counts based on RocksDB")
 set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -fno-rtti")
 set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} -g")
 
diff --git a/README.md b/README.md
index 00711f4..749ca63 100644
--- a/README.md
+++ b/README.md
@@ -97,6 +97,7 @@
 
 COLLOCATORDB *open_collocatordb(const char *s);
 COLLOCATORDB *open_collocatordb_for_write(const char *s);
+void close_collocatordb(COLLOCATORDB *db);
 void inc_collocator(COLLOCATORDB *db, uint64_t w1, uint64_t w2, int8_t dist);
 void dump_collocators(COLLOCATORDB *db, uint32_t w1, uint32_t w2, int8_t dist);
 COLLOCATOR *get_collocators(COLLOCATORDB *db, uint32_t w1);
@@ -111,8 +112,46 @@
 uint64_t get_word_frequency(COLLOCATORDB *db, uint64_t w);
 ```
 
+## Indexing
+
+Building a collocation database is one long stream of increments for the same
+keys. The defaults are chosen for that, and can be adjusted from the
+environment, so that a run of several days does not have to be recompiled to be
+tuned:
+
+| variable | default | meaning |
+|---|---|---|
+| `COLLOCATORDB_WRITE_BUFFER_MB` | 256 | size of one write buffer |
+| `COLLOCATORDB_WRITE_BUFFERS` | 8 | how many of them |
+| `COLLOCATORDB_WRITE_BUFFERS_TO_MERGE` | 4 | how many are combined before they are written, which collapses the increments of the same key early |
+| `COLLOCATORDB_BACKGROUND_JOBS` | cores, at most 32 | threads for flushing and compaction |
+| `COLLOCATORDB_SUBCOMPACTIONS` | 4 | how far a single compaction is spread over threads |
+| `COLLOCATORDB_BLOCK_CACHE_MB` | 512 | block cache, only relevant for reading |
+
+Increments are inserted one writer at a time, rocksdb does not support
+concurrent memtable writes for merge operations, so more indexing threads stop
+helping at some point regardless of these settings.
+
+The database has no write ahead log, so an indexer has to call
+`close_collocatordb()` when it is done, otherwise everything that has not been
+flushed is lost.
+
 ## Changes
 
+* v1.6.0 (2026-08-02)
+  * fixed the loss of increments: rocksdb was told to cancel a write instead of
+    waiting when compaction is behind, and the result was not looked at. Under
+    write pressure a quarter of the increments were lost without any message.
+    A cancelled write is repeated now
+  * fixed the loss of everything that was not flushed yet when an indexer ends:
+    added `close_collocatordb()`, which writes it and closes the database
+  * fixed `get_collocators()` and `get_collocation_scores()` dropping the last
+    collocate of every word
+  * fixed the array returned by `get_collocators()` and
+    `get_collocation_scores()` being far too small, `+` instead of `*` in the
+    size. It is terminated with an empty entry now
+  * the settings for indexing can be adjusted from the environment, see above
+
 * v1.5.0 (2026-07-31)
   * fixed two memory leaks in the collocator lookup, a rocksdb iterator and the
     window sums, together 3.6 kB per call
diff --git a/src/collocatordb.cc b/src/collocatordb.cc
index 1e0a8ec..79d038d 100644
--- a/src/collocatordb.cc
+++ b/src/collocatordb.cc
@@ -9,6 +9,7 @@
 #include "rocksdb/table.h"
 #include "rocksdb/slice.h"
 #include <algorithm>
+#include <atomic>
 #include <cassert>
 #include <cmath>
 #include <cstdint>
@@ -299,6 +300,10 @@
 
 class CollocatorDB {
   WriteOptions merge_option_; // for merge
+  // to repeat a write that rocksdb cancelled, rather than lose the count
+  WriteOptions blocking_merge_option_;
+  std::atomic<uint64_t> stalled_writes_{0};
+  std::atomic<uint64_t> failed_writes_{0};
   char _one[sizeof(uint64_t)]{};
   Slice _one_slice;
   vector<VocabEntry> _vocab;
@@ -320,7 +325,7 @@
   std::shared_ptr<DB> OpenDbForRead(const char *dbname);
 
 public:
-  virtual ~CollocatorDB() = default;
+  virtual ~CollocatorDB();  // flushes what is still in memory, see close()
   void readVocab(const string& fname);
   string getWord(uint32_t w1);
 
@@ -397,14 +402,33 @@
     return value;
   }
 
+  /* A merge that does not lose the count. rocksdb cancels a write with
+     Status::Incomplete instead of waiting when it is asked not to slow down,
+     and then the increment is simply gone. Such a write was not applied, so
+     repeating it blocking cannot count twice. */
+  void merge_one(const Slice &key) {
+    Status s = db_->Merge(merge_option_, key, _one_slice);
+    if (s.ok())
+      return;
+    if (s.IsIncomplete()) {
+      ++stalled_writes_;
+      s = db_->Merge(blocking_merge_option_, key, _one_slice);
+      if (s.ok())
+        return;
+    }
+    if (failed_writes_++ == 0)
+      std::cerr << "collocatordb: cannot write, counts are lost: " << s.ToString()
+                << std::endl;
+  }
+
   virtual void inc(const std::string &key) {
-    db_->Merge(merge_option_, key, _one_slice);
+    merge_one(Slice(key));
   }
 
   void inc(const uint64_t key) {
     char encoded_key[sizeof(uint64_t)];
     EncodeFixed64(encoded_key, key);
-    db_->Merge(merge_option_, std::string(encoded_key, 8), _one_slice);
+    merge_one(Slice(encoded_key, sizeof(uint64_t)));
   }
 
   virtual void inc(uint32_t w1, uint32_t w2, uint8_t dist);
@@ -445,8 +469,32 @@
   }
 
   CollocatorIterator *SeekIterator(uint64_t w1, uint64_t w2, int8_t dist) const;
+
+  /* Writes what is still in memory and closes the database. Without this
+     everything that has not been flushed yet is lost when the process ends,
+     because the write ahead log is switched off for speed. */
+  void close() {
+    if (!db_)
+      return;
+    if (stalled_writes_ > 0)
+      std::cerr << "collocatordb: repeated " << stalled_writes_
+                << " writes that rocksdb had cancelled" << std::endl;
+    if (failed_writes_ > 0)
+      std::cerr << "collocatordb: " << failed_writes_
+                << " writes failed, the database is missing counts" << std::endl;
+    FlushOptions flush_options;
+    flush_options.wait = true;
+    Status s = db_->Flush(flush_options);
+    if (!s.ok() && !s.IsNotSupported())
+      std::cerr << "collocatordb: cannot write what is still in memory: "
+                << s.ToString() << std::endl;
+    db_.reset();
+  }
+
 };
 
+CollocatorDB::~CollocatorDB() { close(); }
+
 CollocatorDB::CollocatorDB(const char *db_name,
                                     bool read_only = false) {
   //		merge_option_.sync = true;
@@ -532,48 +580,80 @@
   return std::shared_ptr<DB>(ROCKSDB_DB_RELEASE(db));
 }
 
+  /* Reads a size from the environment, so that a long indexing run can be
+     tuned without recompiling. Returns the default when unset or unusable. */
+  static uint64_t env_size(const char *name, uint64_t fallback) {
+    const char *value = getenv(name);
+    if (value == nullptr)
+      return fallback;
+    char *end = nullptr;
+    unsigned long long parsed = strtoull(value, &end, 10);
+    if (end == value || parsed == 0) {
+      std::cerr << "collocatordb: ignoring " << name << "=" << value << std::endl;
+      return fallback;
+    }
+    return (uint64_t)parsed;
+  }
+
   std::shared_ptr<DB> CollocatorDB::OpenDb(const char *dbname) {
     ROCKSDB_DB_HANDLE db;
     Options options;
 
     int max_cores = static_cast<int>(std::thread::hardware_concurrency());
 
-    // options.env->SetBackgroundThreads(32, Env::Priority::HIGH); // Increase background threads for high priority
-    // options.env->SetBackgroundThreads(16, Env::Priority::LOW); // Increase background threads for low priority
     options.create_if_missing = true;
     options.merge_operator = std::make_shared<CountMergeOperator>();
-    //options.max_successive_merges = 0;
-    // options.IncreaseParallelism(max_cores); // Utilize all available cores
-    // options.OptimizeLevelStyleCompaction();
 
-    // Increase write buffer size and number of write buffers
-    // options.write_buffer_size = 512 * 1024 * 1024; // 512MB
-    // options.max_write_buffer_number = max_cores;
-    // options.min_write_buffer_number_to_merge = max_cores / 2;
+    /* Indexing a corpus is one long stream of merge operands for the same
+       keys. rocksdb only collapses them when memtables are merged and when
+       files are compacted, so the settings aim at doing that early and often -
+       every collapsed operand is one less to write, to read and to merge
+       again later.
 
-    // Enable concurrent memtable writes
-    options.allow_concurrent_memtable_write = true;
+       All of it can be overridden from the environment, to be able to tune a
+       run that takes days without recompiling. */
+    options.write_buffer_size = env_size("COLLOCATORDB_WRITE_BUFFER_MB", 256) << 20;
+    options.max_write_buffer_number = (int)env_size("COLLOCATORDB_WRITE_BUFFERS", 8);
+    /* collapses the operands of several memtables before they are written */
+    options.min_write_buffer_number_to_merge =
+        (int)env_size("COLLOCATORDB_WRITE_BUFFERS_TO_MERGE", 4);
+
+    /* Compaction has to keep up with the writer, otherwise the level 0 files
+       pile up and every read has to merge through all of them. */
+    options.max_background_jobs =
+        (int)env_size("COLLOCATORDB_BACKGROUND_JOBS",
+                      max_cores > 4 ? (max_cores < 32 ? max_cores : 32) : 4);
+    options.max_subcompactions = (int)env_size("COLLOCATORDB_SUBCOMPACTIONS", 4);
+    options.level0_file_num_compaction_trigger = 4;
+    options.level0_slowdown_writes_trigger = 20;
+    options.level0_stop_writes_trigger = 36;
+
+    /* Merge operands are inserted one writer at a time, rocksdb does not
+       support concurrent memtable writes for them, which is why more threads
+       in the indexer do not help beyond a certain point. */
+    options.allow_concurrent_memtable_write = false;
     options.enable_write_thread_adaptive_yield = true;
-    options.allow_mmap_writes = true;
     options.allow_mmap_reads = true;
 
-    // Optimize block cache size
     BlockBasedTableOptions table_options;
-    table_options.block_cache = NewLRUCache(8 * 1024 * 1024 * 1024L); // 8GB block cache
+    table_options.block_cache =
+        NewLRUCache(env_size("COLLOCATORDB_BLOCK_CACHE_MB", 512) << 20);
     options.table_factory.reset(NewBlockBasedTableFactory(table_options));
 
-    // Adjust compaction settings
-    options.level0_file_num_compaction_trigger = 100;
-    options.level0_slowdown_writes_trigger = 200;
-    options.level0_stop_writes_trigger = 400;
-    options.max_background_compactions = max_cores / 2;
-    options.max_background_flushes = max_cores / 4;
-    // options.disableWA
-    // Tune write options
-    merge_option_.low_pri = true; // Use low priority for compactions
-    merge_option_.disableWAL = true; // Disable Write-Ahead Logging for faster writes
-    merge_option_.sync = false; // Disable sync for faster writes
-    merge_option_.no_slowdown = true; // Disable write slowdown for faster writes
+    /* No write ahead log: an interrupted indexing run is repeated, and the log
+       would double the amount written. Everything that has not been flushed is
+       lost when the process dies, which is what close_collocatordb() is for. */
+    merge_option_.disableWAL = true;
+    merge_option_.sync = false;
+    /* Let rocksdb slow the writer down when compaction falls behind, instead
+       of cancelling the write. It used to be told to do neither, which threw
+       counts away. merge_one() repeats a cancelled write, whatever the
+       settings are. */
+    merge_option_.low_pri = false;
+    merge_option_.no_slowdown = false;
+    blocking_merge_option_ = merge_option_;
+    blocking_merge_option_.low_pri = false;
+    blocking_merge_option_.no_slowdown = false;
 
     Status s = DB::Open(options, dbname, &db);
     if (!s.ok()) {
@@ -751,6 +831,15 @@
     }
   }
 
+  /* A collocate is only finished when the next one begins, so the last one is
+     still in sumWindow when the iterator is through. It used to be dropped,
+     which cost every word one of its collocates. */
+  if (last_w2 != 0xffffffffffffffff && sum >= FREQUENCY_THRESHOLD) {
+    collocators.push_back({});
+    applyCAMeasures(w1, last_w2, sumWindow, sum, usedPositions,
+                    true_window_size, &(collocators[collocators.size() - 1]));
+  }
+
   // #pragma omp taskwait
   sort(collocators.begin(), collocators.end(), sortByLogDiceAF);
 
@@ -897,6 +986,8 @@
   return new CollocatorDB(dbname, false);
 }
 
+DLL_EXPORT void close_collocatordb(COLLOCATORS *db) { delete db; }
+
 DLL_EXPORT COLLOCATORS *open_collocatordb(char *dbname) {
   return new CollocatorDB(dbname, true);
 }
@@ -915,9 +1006,12 @@
   std::vector<Collocator> c = db->get_collocators(w1);
   if (c.empty())
     return nullptr;
-  uint64_t size = c.size() + sizeof c[0];
+  /* one entry more than there are collocators, terminated with w2 == 0 and
+     raw == 0, so that callers can tell where the array ends */
+  uint64_t size = (c.size() + 1) * sizeof c[0];
   auto *p = (COLLOCATORS *)malloc(size);
-  memcpy(p, c.data(), size);
+  memset(p, 0, size);
+  memcpy(p, c.data(), c.size() * sizeof c[0]);
   return p;
 }
 
@@ -926,9 +1020,12 @@
   std::vector<Collocator> c = db->get_collocation_scores(w1, w2);
   if (c.empty())
     return nullptr;
-  uint64_t size = c.size() + sizeof c[0];
+  /* one entry more than there are collocators, terminated with w2 == 0 and
+     raw == 0, so that callers can tell where the array ends */
+  uint64_t size = (c.size() + 1) * sizeof c[0];
   auto *p = (COLLOCATORS *)malloc(size);
-  memcpy(p, c.data(), size);
+  memset(p, 0, size);
+  memcpy(p, c.data(), c.size() * sizeof c[0]);
   return p;
 }
 
diff --git a/src/collocatordb.h b/src/collocatordb.h
index 74f2f0e..adf764d 100644
--- a/src/collocatordb.h
+++ b/src/collocatordb.h
@@ -120,6 +120,7 @@
 
 extern COLLOCATORDB *open_collocatordb(const char *s);
 extern COLLOCATORDB *open_collocatordb_for_write(const char *s);
+extern void close_collocatordb(COLLOCATORDB *db);
 extern void inc_collocator(COLLOCATORDB *db, uint64_t w1, uint64_t w2, int8_t dist);
 extern void dump_collocators(COLLOCATORDB *db, uint32_t w1, uint32_t w2, int8_t dist);
 extern COLLOCATOR *get_collocators(COLLOCATORDB *db, uint32_t w1);
diff --git a/tests/basic_test.c b/tests/basic_test.c
index 7f33b59..75eae5b 100644
--- a/tests/basic_test.c
+++ b/tests/basic_test.c
@@ -106,10 +106,77 @@
   inc_collocator(cdb, 1, 2, 4); size++;
   COLLOCATOR *c = get_collocators(cdb, 0);
   TEST_ASSERT(c != NULL);
-  TEST_CHECK(c[0].w2 == 1);
-  TEST_CHECK(c[0].raw == 2001);
-  TEST_CHECK(c[0].left_raw == 200);
-  TEST_CHECK(c[0].right_raw == 200);
+  /* Both collocates of word0 have to be there, in whatever order they are
+     sorted into: word1 with 2001 and word2 with 2000. The one that came last
+     used to be dropped, which is why this only asked for the first one. */
+  int n, seen1 = 0, seen2 = 0;
+  for (n = 0; c[n].raw > 0; n++) {
+    if (c[n].w2 == 1) {
+      seen1 = 1;
+      TEST_CHECK(c[n].raw == 2001);
+      TEST_CHECK(c[n].left_raw == 200);
+      TEST_CHECK(c[n].right_raw == 200);
+    } else if (c[n].w2 == 2) {
+      seen2 = 1;
+      TEST_CHECK(c[n].raw == 2000);
+    }
+  }
+  TEST_CHECK(n == 2);
+  TEST_MSG("expected 2 collocates, got %d", n);
+  TEST_CHECK(seen1 && seen2);
+
+  rmrf(rocksdbfn);
+}
+
+/* Every increment has to end up in the database, also when rocksdb would
+   rather cancel the write because compaction is behind, and also when they are
+   still in memory when the process that wrote them is gone. Both used to lose
+   counts silently: a quarter of them under write pressure, and everything that
+   had not been flushed when the indexer ended. */
+void test_no_increment_is_lost() {
+  char tmp_template[] = "/tmp/tmpfileXXXXXX";
+  int fd = mkstemp(tmp_template);
+  if (fd == -1) {
+    perror("mkstemp");
+    exit(EXIT_FAILURE);
+  }
+  close(fd);
+  char *tmp = strdup(tmp_template);
+  char rocksdbfn[1024], vocabfn[1024];
+  const long increments = 200000;
+  const int collocates = 100;
+  long i, total = 0;
+
+  snprintf(rocksdbfn, sizeof(rocksdbfn), "%s.rocksdb", tmp);
+  snprintf(vocabfn, sizeof(vocabfn), "%s.vocab", tmp);
+  FILE *h = fopen(vocabfn, "w");
+  for (i = 0; i <= collocates + 1; i++)
+    fprintf(h, "word%ld 1000\n", i);
+  fclose(h);
+
+  /* small write buffers, so that the writer runs into flushes and stalls */
+  setenv("COLLOCATORDB_WRITE_BUFFER_MB", "1", 1);
+  setenv("COLLOCATORDB_WRITE_BUFFERS", "2", 1);
+
+  COLLOCATORDB *cdb = open_collocatordb_for_write(rocksdbfn);
+  TEST_ASSERT(cdb != NULL);
+  for (i = 0; i < increments; i++)
+    inc_collocator(cdb, 1, 2 + (i % collocates), 1);   /* never w2 == w1 */
+  close_collocatordb(cdb);      /* has to write what is still in memory */
+
+  unsetenv("COLLOCATORDB_WRITE_BUFFER_MB");
+  unsetenv("COLLOCATORDB_WRITE_BUFFERS");
+
+  cdb = open_collocatordb(tmp);
+  TEST_ASSERT(cdb != NULL);
+  COLLOCATOR *c = get_collocators(cdb, 1);
+  TEST_ASSERT(c != NULL);
+  for (i = 0; i < collocates && c[i].raw > 0; i++)
+    total += c[i].raw;
+  TEST_CHECK(i == collocates);
+  TEST_MSG("only %ld of %d collocates are in the database", i, collocates);
+  TEST_CHECK(total == increments);
+  TEST_MSG("%ld of %ld increments survived", total, increments);
 
   rmrf(rocksdbfn);
 }
@@ -160,6 +227,7 @@
     { "collocation analysis", test_collocation_analysis },
     { "collocation analysis as json", test_collocation_analysis_as_json },
     { "writing", test_writing },
+    { "no increment is lost", test_no_increment_is_lost },
     { "version function", test_version_function },
     { "get word id", test_get_word_id },
     { "get corpus size", test_get_corpus_size},