4 files changed, 85 insertions, 9 deletions

diff --git a/src/common/threadpool.cpp b/src/common/threadpool.cpp
index 51e071577..fb238dca7 100644
--- a/src/common/threadpool.cpp
+++ b/src/common/threadpool.cpp
@@ -39,11 +39,11 @@ static __thread int depth = 0;
 
 namespace tools
 {
-threadpool::threadpool() : running(true), active(0) {
+threadpool::threadpool(unsigned int max_threads) : running(true), active(0) {
   boost::thread::attributes attrs;
   attrs.set_stack_size(THREAD_STACK_SIZE);
-  max = tools::get_max_concurrency();
-  size_t i = max;
+  max = max_threads ? max_threads : tools::get_max_concurrency();
+  unsigned int i = max;
   while(i--) {
     threads.push_back(boost::thread(attrs, boost::bind(&threadpool::run, this)));
   }
@@ -78,7 +78,7 @@ void threadpool::submit(waiter *obj, std::function<void()> f) {
   }
 }
 
-int threadpool::get_max_concurrency() {
+unsigned int threadpool::get_max_concurrency() const {
   return max;
 }
 
diff --git a/src/common/threadpool.h b/src/common/threadpool.h
index 34152541c..bf80a87f6 100644
--- a/src/common/threadpool.h
+++ b/src/common/threadpool.h
@@ -46,6 +46,9 @@ public:
     static threadpool instance;
     return instance;
   }
+  static threadpool *getNewForUnitTests(unsigned max_threads = 0) {
+    return new threadpool(max_threads);
+  }
 
   // The waiter lets the caller know when all of its
   // tasks are completed.
@@ -66,11 +69,12 @@ public:
   // task to finish.
   void submit(waiter *waiter, std::function<void()> f);
 
-  int get_max_concurrency();
+  unsigned int get_max_concurrency() const;
+
+  ~threadpool();
 
   private:
-    threadpool();
-    ~threadpool();
+    threadpool(unsigned int max_threads = 0);
     typedef struct entry {
       waiter *wo;
       std::function<void()> f;
@@ -79,8 +83,8 @@ public:
     boost::condition_variable has_work;
     boost::mutex mutex;
     std::vector<boost::thread> threads;
-    int active;
-    int max;
+    unsigned int active;
+    unsigned int max;
     bool running;
     void run();
 };
diff --git a/src/common/util.cpp b/src/common/util.cpp
index 7e77e19b1..008610117 100644
--- a/src/common/util.cpp
+++ b/src/common/util.cpp
@@ -34,6 +34,17 @@
 #include <gnu/libc-version.h>
 #endif
 
+#ifdef __GLIBC__
+#include <sys/types.h>
+#include <sys/stat.h>
+#include <ustat.h>
+#include <unistd.h>
+#include <dirent.h>
+#include <string.h>
+#include <ctype.h>
+#include <string>
+#endif
+
 #include "unbound.h"
 
 #include "include_base_utils.h"
@@ -632,6 +643,65 @@ std::string get_nix_version_display_string()
 #endif
   }
 
+  bool is_hdd(const char *path)
+  {
+#ifdef __GLIBC__
+    std::string device = "";
+    struct stat st, dst;
+    if (stat(path, &st) < 0)
+      return 0;
+
+    DIR *dir = opendir("/dev/block");
+    if (!dir)
+      return 0;
+    struct dirent *de;
+    while ((de = readdir(dir)))
+    {
+      if (strcmp(de->d_name, ".") && strcmp(de->d_name, ".."))
+      {
+        std::string dev_path = std::string("/dev/block/") + de->d_name;
+        char resolved[PATH_MAX];
+        if (realpath(dev_path.c_str(), resolved) && !strncmp(resolved, "/dev/", 5))
+        {
+          if (stat(resolved, &dst) == 0)
+          {
+            if (dst.st_rdev == st.st_dev)
+            {
+              // take out trailing digits (eg, sda1 -> sda)
+              char *ptr = resolved;
+              while (*ptr)
+                ++ptr;
+              while (ptr > resolved && isdigit(*--ptr))
+                *ptr = 0;
+              device = resolved + 5;
+              break;
+            }
+          }
+        }
+      }
+    }
+    closedir(dir);
+
+    if (device.empty())
+      return 0;
+
+    std::string sys_path = "/sys/block/" + device + "/queue/rotational";
+    FILE *f = fopen(sys_path.c_str(), "r");
+    if (!f)
+      return false;
+    char s[8];
+    char *ptr = fgets(s, sizeof(s), f);
+    fclose(f);
+    if (!ptr)
+      return 0;
+    s[sizeof(s) - 1] = 0;
+    int n = atoi(s); // returns 0 on parse error
+    return n == 1;
+#else
+    return 0;
+#endif
+  }
+
   namespace
   {
     boost::mutex max_concurrency_lock;
diff --git a/src/common/util.h b/src/common/util.h
index d3ba47a4f..7caf0e3c5 100644
--- a/src/common/util.h
+++ b/src/common/util.h
@@ -212,4 +212,6 @@ namespace tools
 
   bool sha256sum(const uint8_t *data, size_t len, crypto::hash &hash);
   bool sha256sum(const std::string &filename, crypto::hash &hash);
+
+  bool is_hdd(const char *path);
 }