Catch MDB_MAP_FULL errors from mdb_txn_commit

diff --git a/src/caffe/util/db_lmdb.cpp b/src/caffe/util/db_lmdb.cpp
index df83a52..4567cd7 100644 (file)
--- a/src/caffe/util/db_lmdb.cpp
+++ b/src/caffe/util/db_lmdb.cpp
@@ -62,36 +62,42 @@ void LMDBTransaction::Commit() {
   MDB_CHECK(mdb_txn_begin(mdb_env_, NULL, 0, &mdb_txn));
   MDB_CHECK(mdb_dbi_open(mdb_txn, NULL, 0, &mdb_dbi));
 
-  bool out_of_memory = false;
   for (int i = 0; i < keys.size(); i++) {
     mdb_key.mv_size = keys[i].size();
     mdb_key.mv_data = const_cast<char*>(keys[i].data());
     mdb_data.mv_size = values[i].size();
     mdb_data.mv_data = const_cast<char*>(values[i].data());
 
+    // Add data to the transaction
     int put_rc = mdb_put(mdb_txn, mdb_dbi, &mdb_key, &mdb_data, 0);
     if (put_rc == MDB_MAP_FULL) {
-      out_of_memory = true;
-      break;
-    } else {
-      // Failed for some other reason
-      MDB_CHECK(put_rc);
+      // Out of memory - double the map size and retry
+      mdb_txn_abort(mdb_txn);
+      mdb_dbi_close(mdb_env_, mdb_dbi);
+      DoubleMapSize();
+      Commit();
+      return;
     }
+    // May have failed for some other reason
+    MDB_CHECK(put_rc);
   }
 
-  if (!out_of_memory) {
-    // Commit the transaction
-    MDB_CHECK(mdb_txn_commit(mdb_txn));
-    mdb_dbi_close(mdb_env_, mdb_dbi);
-    keys.clear();
-    values.clear();
-  } else {
-    // Double the map size and retry
-    mdb_txn_abort(mdb_txn);
+  // Commit the transaction
+  int commit_rc = mdb_txn_commit(mdb_txn);
+  if (commit_rc == MDB_MAP_FULL) {
+    // Out of memory - double the map size and retry
     mdb_dbi_close(mdb_env_, mdb_dbi);
     DoubleMapSize();
     Commit();
+    return;
   }
+  // May have failed for some other reason
+  MDB_CHECK(commit_rc);
+
+  // Cleanup after successful commit
+  mdb_dbi_close(mdb_env_, mdb_dbi);
+  keys.clear();
+  values.clear();
 }
 
 void LMDBTransaction::DoubleMapSize() {