berkeley db now support custom comparison function

include/sdskv-common.h

@@ -17,8 +17,18 @@ typedef enum sdskv_db_type_t
 typedef uint64_t sdskv_database_id_t;
 #define SDSKV_DATABASE_ID_INVALID 0
 
-#define SDSKV_SUCCESS 0
-#define SDSKV_ERROR  -1
+#define SDSKV_SUCCESS          0 /* Success */
+#define SDSKV_ERR_ALLOCATION  -1 /* Error allocating something */
+#define SDSKV_ERR_INVALID_ARG -2 /* An argument is invalid */
+#define SDSKV_ERR_MERCURY     -3 /* An error happened calling a Mercury function */
+#define SDSKV_ERR_DB_CREATE   -4 /* Could not create database */
+#define SDSKV_ERR_DB_NAME     -5 /* Invalid database name */
+#define SDSKV_ERR_UNKNOWN_DB  -6 /* Database refered to by id is not known to provider */
+#define SDSKV_ERR_UNKNOWN_PR  -7 /* Mplex id could not be matched with a provider */
+#define SDSKV_ERR_PUT         -8 /* Could not put into the database */
+#define SDSKV_ERR_UNKNOWN_KEY -9 /* Key requested does not exist */
+#define SDSKV_ERR_SIZE        -10 /* Client did not allocate enough for the requested data */
+#define SDSKV_ERR_ERASE       -11 /* Could not erase the given key */
 
 #if defined(__cplusplus)
 }

src/datastore/berkeleydb_datastore.cc

@@ -3,6 +3,7 @@
 #include "berkeleydb_datastore.h"
 #include "kv-config.h"
 #include <chrono>
+#include <cstring>
 #include <iostream>
 #include <boost/filesystem.hpp>
 
@@ -21,7 +22,8 @@ BerkeleyDBDataStore::BerkeleyDBDataStore(Duplicates duplicates, bool eraseOnGet,
 };
   
 BerkeleyDBDataStore::~BerkeleyDBDataStore() {
-  delete _dbm;
+//  delete _dbm;
+  delete _wrapper;
   delete _dbenv;
 };
 
@@ -88,11 +90,14 @@ void BerkeleyDBDataStore::createDatabase(std::string db_name) {
   }
   
   if (status == 0) {
-    _dbm = new Db(_dbenv, DB_CXX_NO_EXCEPTIONS);
+    _wrapper = new DbWrapper(_dbenv, DB_CXX_NO_EXCEPTIONS);
+    _dbm = &(_wrapper->_db);
 
     if (_duplicates == Duplicates::ALLOW) {
       _dbm->set_flags(DB_DUP); // Allow duplicate keys
     }
+
+    _dbm->set_bt_compare(&(BerkeleyDBDataStore::compkeys));
   
     uint32_t flags = DB_CREATE | DB_AUTO_COMMIT | DB_THREAD; // Allow database creation
     if (_in_memory) {
@@ -125,10 +130,7 @@ void BerkeleyDBDataStore::createDatabase(std::string db_name) {
 };
 
 void BerkeleyDBDataStore::set_comparison_function(comparator_fn less) {
-    if(less) {
-        fprintf(stderr, "Error: BerkeleyDBDataStore's comparison function cannot be changed\n");
-        exit(-1);
-    }
+    _wrapper->_less = less;
 }
 
 bool BerkeleyDBDataStore::put(const ds_bulk_t &key, const ds_bulk_t &data) {
@@ -176,9 +178,10 @@ bool BerkeleyDBDataStore::get(const ds_bulk_t &key, ds_bulk_t &data) {
   data.clear();
 
   Dbt db_key((void*)&(key[0]), uint32_t(key.size()));
+  db_key.set_ulen(uint32_t(key.size()));
   Dbt db_data;
   db_key.set_flags(DB_DBT_USERMEM);
-  db_data.set_flags(DB_DBT_REALLOC);
+  db_data.set_flags(DB_DBT_MALLOC);
   status = _dbm->get(NULL, &db_key, &db_data, 0);
 
   if (status != DB_NOTFOUND && status != DB_KEYEMPTY && db_data.get_size() > 0) {
@@ -322,3 +325,18 @@ std::vector<std::pair<ds_bulk_t,ds_bulk_t>> BerkeleyDBDataStore::vlist_keyvals(c
     cursorp->close();
     return result;
 }
+
+int BerkeleyDBDataStore::compkeys(Db *db, const Dbt *dbt1, const Dbt *dbt2, size_t *locp) {
+    DbWrapper* _wrapper = (DbWrapper*)(((char*)db) - offsetof(BerkeleyDBDataStore::DbWrapper, _db));
+    if(_wrapper->_less) {
+        return (_wrapper->_less)(dbt1->get_data(), dbt1->get_size(), 
+                dbt2->get_data(), dbt2->get_size());
+    } else {
+        size_t s = dbt1->get_size() > dbt2->get_size() ? dbt2->get_size() : dbt1->get_size();
+        int c = std::memcmp(dbt1->get_data(), dbt2->get_data(), s);
+        if(c != 0) return c;
+        if(dbt1->get_size() < dbt2->get_size()) return -1;
+        if(dbt1->get_size() > dbt2->get_size()) return 1;
+        return 0;
+    }
+}

src/datastore/berkeleydb_datastore.h

@@ -10,22 +10,37 @@
 
 // may want to implement some caching for persistent stores like BerkeleyDB
 class BerkeleyDBDataStore : public AbstractDataStore {
-public:
-  BerkeleyDBDataStore();
-  BerkeleyDBDataStore(Duplicates duplicates, bool eraseOnGet, bool debug);
-  virtual ~BerkeleyDBDataStore();
-  virtual void createDatabase(std::string db_name);
-  virtual bool put(const ds_bulk_t &key, const ds_bulk_t &data);
-  virtual bool get(const ds_bulk_t &key, ds_bulk_t &data);
-  virtual bool get(const ds_bulk_t &key, std::vector<ds_bulk_t> &data);
-  virtual bool erase(const ds_bulk_t &key);
-  virtual void set_in_memory(bool enable); // enable/disable in-memory mode
-  virtual void set_comparison_function(comparator_fn less);
-protected:
-  virtual std::vector<ds_bulk_t> vlist_keys(const ds_bulk_t &start, size_t count, const ds_bulk_t &prefix);
-  virtual std::vector<std::pair<ds_bulk_t,ds_bulk_t>> vlist_keyvals(const ds_bulk_t &start_key, size_t count, const ds_bulk_t &);
-  DbEnv *_dbenv = NULL;
-  Db *_dbm = NULL;
+
+    private:
+        struct DbWrapper {
+            Db _db;
+            AbstractDataStore::comparator_fn _less;
+
+            template<typename ... T>
+            DbWrapper(T&&... args) :
+            _db(std::forward<T>(args)...), _less(nullptr) {}
+        };
+
+        static int compkeys(Db *db, const Dbt *dbt1, const Dbt *dbt2, size_t *locp);
+
+    public:
+        BerkeleyDBDataStore();
+        BerkeleyDBDataStore(Duplicates duplicates, bool eraseOnGet, bool debug);
+        virtual ~BerkeleyDBDataStore();
+        virtual void createDatabase(std::string db_name);
+        virtual bool put(const ds_bulk_t &key, const ds_bulk_t &data);
+        virtual bool get(const ds_bulk_t &key, ds_bulk_t &data);
+        virtual bool get(const ds_bulk_t &key, std::vector<ds_bulk_t> &data);
+        virtual bool erase(const ds_bulk_t &key);
+        virtual void set_in_memory(bool enable); // enable/disable in-memory mode
+        virtual void set_comparison_function(comparator_fn less);
+    protected:
+        virtual std::vector<ds_bulk_t> vlist_keys(const ds_bulk_t &start, size_t count, const ds_bulk_t &prefix);
+        virtual std::vector<std::pair<ds_bulk_t,ds_bulk_t>> vlist_keyvals(const ds_bulk_t &start_key, size_t count, const ds_bulk_t &);
+        DbEnv *_dbenv = nullptr;
+        Db *_dbm = nullptr;
+        DbWrapper* _wrapper = nullptr;
+
 };
 
 #endif // bdb_datastore_h