1 files changed, 94 insertions, 59 deletions

diff --git a/common/storage_lmdb.cpp b/common/storage_lmdb.cpp
index 08eea37..f7999d1 100644
--- a/common/storage_lmdb.cpp
+++ b/common/storage_lmdb.cpp
@@ -35,9 +35,6 @@
 #include <lmdb.h>
 #include "log.h"
 
-SINK_DEBUG_AREA("storage")
-// SINK_DEBUG_COMPONENT(d->storageRoot.toLatin1() + '/' + d->name.toLatin1())
-
 namespace Sink {
 namespace Storage {
 
@@ -169,6 +166,27 @@ public:
             if (const int rc = mdb_dbi_open(transaction, db.constData(), flags, &dbi)) {
                 //Create the db if it is not existing already
                 if (rc == MDB_NOTFOUND && !readOnly) {
+                    //Sanity check db name
+                    {
+                        auto parts = db.split('.');
+                        for (const auto &p : parts) {
+                            auto containsSpecialCharacter = [] (const QByteArray &p) {
+                                for (int i = 0; i < p.size(); i++) {
+                                    const auto c = p.at(i);
+                                    //Between 0 and z in the ascii table. Essentially ensures that the name is printable and doesn't contain special chars
+                                    if (c < 0x30 || c > 0x7A) {
+                                        return true;
+                                    }
+                                }
+                                return false;
+                            };
+                            if (p.isEmpty() || containsSpecialCharacter(p)) {
+                                SinkError() << "Tried to create a db with an invalid name. Hex:" << db.toHex() << " ASCII:" << db;
+                                Q_ASSERT(false);
+                                throw std::runtime_error("Fatal error while creating db.");
+                            }
+                        }
+                    }
                     if (const int rc = mdb_dbi_open(transaction, db.constData(), flags | MDB_CREATE, &dbi)) {
                         SinkWarning() << "Failed to create db " << QByteArray(mdb_strerror(rc));
                         Error error(name.toLatin1(), ErrorCodes::GenericError, "Error while creating database: " + QByteArray(mdb_strerror(rc)));
@@ -542,6 +560,7 @@ DataStore::Transaction::Transaction(Transaction &&other) : d(nullptr)
 DataStore::Transaction &DataStore::Transaction::operator=(DataStore::Transaction &&other)
 {
     if (&other != this) {
+        abort();
         delete d;
         d = other.d;
         other.d = nullptr;
@@ -639,11 +658,6 @@ static bool ensureCorrectDb(DataStore::NamedDatabase &database, const QByteArray
     return !openedTheWrongDatabase;
 }
 
-bool DataStore::Transaction::validateNamedDatabases()
-{
-    return true;
-}
-
 DataStore::NamedDatabase DataStore::Transaction::openDatabase(const QByteArray &db, const std::function<void(const DataStore::Error &error)> &errorHandler, bool allowDuplicates) const
 {
     if (!d) {
@@ -694,7 +708,7 @@ QList<QByteArray> DataStore::Transaction::getDatabaseNames() const
 class DataStore::Private
 {
 public:
-    Private(const QString &s, const QString &n, AccessMode m);
+    Private(const QString &s, const QString &n, AccessMode m, const DbLayout &layout = {});
     ~Private();
 
     QString storageRoot;
@@ -702,68 +716,85 @@ public:
 
     MDB_env *env;
     AccessMode mode;
+    Sink::Log::Context logCtx;
+
+    void initEnvironment(const QString &fullPath, const DbLayout &layout)
+    {
+        // Ensure the environment is only created once, and that we only have one environment per process
+        if (!(env = sEnvironments.value(fullPath))) {
+            QMutexLocker locker(&sMutex);
+            if (!(env = sEnvironments.value(fullPath))) {
+                int rc = 0;
+                if ((rc = mdb_env_create(&env))) {
+                    SinkWarningCtx(logCtx) << "mdb_env_create: " << rc << " " << mdb_strerror(rc);
+                    qCritical() << "mdb_env_create: " << rc << " " << mdb_strerror(rc);
+                } else {
+                    //Limit large enough to accomodate all our named dbs. This only starts to matter if the number gets large, otherwise it's just a bunch of extra entries in the main table.
+                    mdb_env_set_maxdbs(env, 50);
+                    const bool readOnly = (mode == ReadOnly);
+                    unsigned int flags = MDB_NOTLS;
+                    if (readOnly) {
+                        flags |= MDB_RDONLY;
+                    }
+                    if ((rc = mdb_env_open(env, fullPath.toStdString().data(), flags, 0664))) {
+                        SinkWarningCtx(logCtx) << "mdb_env_open: " << rc << ":" << mdb_strerror(rc);
+                        mdb_env_close(env);
+                        env = 0;
+                    } else {
+                        if (RUNNING_ON_VALGRIND) {
+                            // In order to run valgrind this size must be smaller than half your available RAM
+                            // https://github.com/BVLC/caffe/issues/2404
+                            mdb_env_set_mapsize(env, (size_t)10485760 * (size_t)1000); // 1MB * 1000
+                        } else {
+                            //This is the maximum size of the db (but will not be used directly), so we make it large enough that we hopefully never run into the limit.
+                            mdb_env_set_mapsize(env, (size_t)10485760 * (size_t)100000); // 1MB * 1000
+                        }
+                        Q_ASSERT(env);
+                        sEnvironments.insert(fullPath, env);
+                        //Open all available dbi's
+                        bool noLock = true;
+                        auto t = Transaction(new Transaction::Private(readOnly, nullptr, name, env, noLock));
+                        if (!layout.tables.isEmpty()) {
+
+                            //TODO upgrade db if the layout has changed:
+                            //* read existing layout
+                            //* if layout is not the same create new layout
+                        //If the db is read only, abort if the db is not yet existing.
+                        //If the db is not read-only but is not existing, ensure we have a layout and create all tables.
+
+                            for (auto it = layout.tables.constBegin(); it != layout.tables.constEnd(); it++) {
+                                bool allowDuplicates = it.value();
+                                t.openDatabase(it.key(), {}, allowDuplicates);
+                            }
+                        } else {
+                            for (const auto &db : t.getDatabaseNames()) {
+                                //Get dbi to store for future use.
+                                t.openDatabase(db);
+                            }
+                        }
+                        //To persist the dbis (this is also necessary for read-only transactions)
+                        t.commit();
+                    }
+                }
+            }
+        }
+    }
+
 };
 
-DataStore::Private::Private(const QString &s, const QString &n, AccessMode m) : storageRoot(s), name(n), env(0), mode(m)
+DataStore::Private::Private(const QString &s, const QString &n, AccessMode m, const DbLayout &layout) : storageRoot(s), name(n), env(0), mode(m), logCtx(n.toLatin1())
 {
+
     const QString fullPath(storageRoot + '/' + name);
     QFileInfo dirInfo(fullPath);
     if (!dirInfo.exists() && mode == ReadWrite) {
         QDir().mkpath(fullPath);
         dirInfo.refresh();
     }
-    Sink::Log::Context logCtx{n.toLatin1()};
     if (mode == ReadWrite && !dirInfo.permission(QFile::WriteOwner)) {
         qCritical() << fullPath << "does not have write permissions. Aborting";
     } else if (dirInfo.exists()) {
-        // Ensure the environment is only created once
-        QMutexLocker locker(&sMutex);
-
-        /*
-        * It seems we can only ever have one environment open in the process.
-        * Otherwise multi-threading breaks.
-        */
-        env = sEnvironments.value(fullPath);
-        if (!env) {
-            int rc = 0;
-            if ((rc = mdb_env_create(&env))) {
-                // TODO: handle error
-                SinkWarningCtx(logCtx) << "mdb_env_create: " << rc << " " << mdb_strerror(rc);
-            } else {
-                mdb_env_set_maxdbs(env, 50);
-                unsigned int flags = MDB_NOTLS;
-                if (mode == ReadOnly) {
-                    flags |= MDB_RDONLY;
-                }
-                if ((rc = mdb_env_open(env, fullPath.toStdString().data(), flags, 0664))) {
-                    SinkWarningCtx(logCtx) << "mdb_env_open: " << rc << ":" << mdb_strerror(rc);
-                    mdb_env_close(env);
-                    env = 0;
-                } else {
-                    if (RUNNING_ON_VALGRIND) {
-                        // In order to run valgrind this size must be smaller than half your available RAM
-                        // https://github.com/BVLC/caffe/issues/2404
-                        const size_t dbSize = (size_t)10485760 * (size_t)1000; // 1MB * 1000
-                        mdb_env_set_mapsize(env, dbSize);
-                    } else {
-                        // FIXME: dynamic resize
-                        const size_t dbSize = (size_t)10485760 * (size_t)8000; // 1MB * 8000
-                        mdb_env_set_mapsize(env, dbSize);
-                    }
-                    sEnvironments.insert(fullPath, env);
-                    //Open all available dbi's
-                    bool noLock = true;
-                    bool requestedRead = m == ReadOnly;
-                    auto t = Transaction(new Transaction::Private(requestedRead, nullptr, name, env, noLock));
-                    for (const auto &db : t.getDatabaseNames()) {
-                        //Get dbi to store for future use.
-                        t.openDatabase(db);
-                    }
-                    //To persist the dbis (this is also necessary for read-only transactions)
-                    t.commit();
-                }
-            }
-        }
+        initEnvironment(fullPath, layout);
     }
 }
 
@@ -777,6 +808,10 @@ DataStore::DataStore(const QString &storageRoot, const QString &name, AccessMode
 {
 }
 
+DataStore::DataStore(const QString &storageRoot, const DbLayout &dbLayout, AccessMode mode) : d(new Private(storageRoot, dbLayout.name, mode, dbLayout))
+{
+}
+
 DataStore::~DataStore()
 {
     delete d;