Add a lock mode and allow touching/locking individual idx entries in touchdb

Author: milot-mirdita

src/commons/Parameters.cpp

@@ -319,6 +319,8 @@ Parameters::Parameters():
         PARAM_TEMPERATURE(PARAM_TEMPERATURE_ID, "--temperature", "Temperature", "Temperature for forward-backward", typeid(float), (void *) &temperature, "^(0\\.[0-9]+|[1-9][0-9]*\\.?[0-9]*)$", MMseqsParameter::COMMAND_EXPERT),
         PARAM_BLOCKLEN(PARAM_BLOCKLEN_ID, "--blocklen", "Block length", "Block length for forward-backward", typeid(int), (void *) &blocklen, "^[1-9]{1}[0-9]*$", MMseqsParameter::COMMAND_EXPERT),
         PARAM_FWBW_BACKTRACE_MODE(PARAM_FWBW_BACKTRACE_MODE_ID, "--fwbw-backtrace-mode", "Backtrace mode", "Backtrace mode 0: no backtrace, 1: local", typeid(int), (void *) &fwbwBacktraceMode, "^[01]$", MMseqsParameter::COMMAND_EXPERT),
+        // touchdb
+        PARAM_TOUCH_LOCK(PARAM_TOUCH_LOCK_ID, "--touch-lock", "Touch lock", "Lock touched database or database entries into memory. Process will not exit until killed.", typeid(bool), (void *) &touchLock, "", MMseqsParameter::COMMAND_EXPERT),
         // for modules that should handle -h themselves
         PARAM_HELP(PARAM_HELP_ID, "-h", "Help", "Help", typeid(bool), (void *) &help, "", MMseqsParameter::COMMAND_HIDDEN),
         PARAM_HELP_LONG(PARAM_HELP_LONG_ID, "--help", "Help", "Help", typeid(bool), (void *) &help, "", MMseqsParameter::COMMAND_HIDDEN)
@@ -1493,6 +1495,8 @@ Parameters::Parameters():
     appenddbtoindex.push_back(&PARAM_V);
 
     // touchdb
+    touchdb.push_back(&PARAM_ID_LIST);
+    touchdb.push_back(&PARAM_TOUCH_LOCK);
     touchdb.push_back(&PARAM_THREADS);
     touchdb.push_back(&PARAM_V);
 
@@ -2707,6 +2711,10 @@ void Parameters::setDefaults() {
     temperature = 1;
     blocklen = 16;
     fwbwBacktraceMode = 1;
+
+    // touchdb
+    touchLock = false;
+
     // help
     help = 0;
 

src/commons/Parameters.h

@@ -740,6 +740,10 @@ class Parameters {
     float temperature;
     int blocklen;
     int fwbwBacktraceMode;
+
+    // touchdb
+    bool touchLock;
+
     // for modules that should handle -h themselves
     bool help;
 
@@ -1107,6 +1111,10 @@ class Parameters {
     PARAMETER(PARAM_TEMPERATURE)
     PARAMETER(PARAM_BLOCKLEN)
     PARAMETER(PARAM_FWBW_BACKTRACE_MODE)
+
+    // touchdb
+    PARAMETER(PARAM_TOUCH_LOCK)
+
     // for modules that should handle -h themselves
     PARAMETER(PARAM_HELP)
     PARAMETER(PARAM_HELP_LONG)

src/util/touchdb.cpp

@@ -3,19 +3,70 @@
 #include "PrefilteringIndexReader.h"
 #include "MemoryMapped.h"
 
+#include <sys/mman.h>
+
 int touchdb(int argc, const char **argv, const Command& command) {
     Parameters& par = Parameters::getInstance();
     par.parseParameters(argc, argv, command, true, 0, 0);
 
-    std::string db = par.db1;
+    const char* mlockError = "Consider granting the `CAP_IPC_LOCK` capability to this binary (e.g., `setcap cap_ipc_lock=+ep <program>`, or run as root) or raising the locked-memory limit so `mlock(2)` can succeed. Check your current limit with `ulimit -l` and increase it via `ulimit -l <KB>` for the shell, `LimitMEMLOCK=infinity` in your systemd unit, or persistent user limits in `/etc/security/limits.conf` (e.g., `youruser hard memlock unlimited`)\n";
 
+    std::string db = par.db1;
     std::string indexDB = PrefilteringIndexReader::searchForIndex(db);
     if (indexDB.empty() == false) {
         db = indexDB;
+        if (par.idList != "") {
+            std::string idx = db + ".index";
+            DBReader<unsigned int> reader(db.c_str(), idx.c_str(), 1, DBReader<unsigned int>::USE_DATA | DBReader<unsigned int>::USE_INDEX);
+            reader.open(DBReader<unsigned int>::NOSORT);
+
+            std::vector<std::string> ids = Util::split(par.idList, ",");
+            for (size_t i = 0; i < ids.size(); ++i) {
+                size_t id = reader.getId(Util::fast_atoi<unsigned int>(ids[i].c_str()));
+                if (id == UINT_MAX) {
+                    Debug(Debug::WARNING) << "Key " << ids[i] << " not found in database\n";
+                    continue;
+                }
+                size_t currDataOffset = reader.getOffset(id);
+                size_t nextDataOffset = reader.findNextOffsetid(id);
+                size_t dataSize = nextDataOffset - currDataOffset;
+                char* data = reader.getDataUncompressed(id);
+                Util::touchMemory(data, dataSize);
+                if (par.touchLock) {
+                    int res = mlock(data, dataSize);
+                    if (res != 0) {
+                        Debug(Debug::ERROR) << "Could not lock memory " << strerror(errno) << "\n";
+                        Debug(Debug::ERROR) << mlockError;
+                        reader.close();
+                        return EXIT_FAILURE;
+                    }
+                }
+            }
+
+            if (par.touchLock) {
+                Debug(Debug::INFO) << "Touched and locked " << ids.size() << " entries. Process will not exit until killed.\n";
+                pause();
+            } else {
+                Debug(Debug::INFO) << "Touched " << ids.size() << " entries\n";
+            }
+
+            reader.close();
+            return EXIT_SUCCESS;
+        }
     }
 
     MemoryMapped map(db, MemoryMapped::WholeFile, MemoryMapped::CacheHint::SequentialScan);
     Util::touchMemory(reinterpret_cast<const char*>(map.getData()), map.mappedSize());
+    if (par.touchLock) {
+        int res = mlock(map.getData(), map.mappedSize());
+        if (res != 0) {
+            Debug(Debug::ERROR) << "Could not lock memory " << strerror(errno) << "\n";
+            Debug(Debug::ERROR) << mlockError;
+            return EXIT_FAILURE;
+        }
+        Debug(Debug::INFO) << "Touched and locked database. Process will not exit until killed.\n";
+        pause();
+    }
 
     return EXIT_SUCCESS;
 }