loaders: factor out BackgroundProducer

Author: mlin

loaders/common.hpp

@@ -3,10 +3,14 @@
 #include "SQLiteCpp/SQLiteCpp.h"
 #include "genomicsqlite.h"
 #include "strlcpy.h"
+#include <atomic>
+#include <chrono>
 #include <iostream>
 #include <memory>
 #include <sqlite3.h>
 #include <string>
+#include <thread>
+#include <vector>
 
 #ifndef NDEBUG
 #define _DBG cerr << __FILE__ << ":" << __LINE__ << ": "
@@ -43,7 +47,7 @@ void split(string &s, char delim, Out result, uint64_t maxsplit = ULLONG_MAX) {
     split(&s[0], delim, result, maxsplit);
 }
 
-// because std::ostringstream is too slow :(
+// because ostringstream is too slow :(
 class OStringStream {
   public:
     OStringStream(size_t initial_capacity) : buf_size_(initial_capacity), cursor_(0) {
@@ -84,7 +88,7 @@ class OStringStream {
         Add(s);
         return *this;
     }
-    inline OStringStream &operator<<(const std::string &s) { return *this << s.c_str(); }
+    inline OStringStream &operator<<(const string &s) { return *this << s.c_str(); }
 
     inline const char *Get() const {
         assert(buf_[cursor_] == 0);
@@ -118,3 +122,117 @@ class OStringStream {
     unique_ptr<char[]> buf_;
     size_t buf_size_, cursor_;
 };
+
+// Producer-consumer pattern: background producer thread preprocessing "Items" to be consumed by
+// the main thread. Queues up to `ringsize` prepared items.
+template <class Item> class BackgroundProducer {
+  protected:
+    // Will be called on the background thread to populate the next item. Fills it out in-place
+    // from undefined initial state (avoids reallocating); returns true on success, false if the
+    // item stream is successfully complete, or throws an exception.
+    virtual bool Produce(Item &) = 0;
+
+  private:
+    vector<Item> ring_;
+    int R_;
+    unique_ptr<thread> worker_;
+    atomic<bool> stop_;
+    string errmsg_;
+    atomic<long long> p_, // produced
+        c_;               // if c_>0 then item (c_-1)%R is currently being consumed
+    chrono::duration<double> p_blocked_ = chrono::duration<double>::zero(),
+                             c_blocked_ = chrono::duration<double>::zero();
+    chrono::time_point<chrono::high_resolution_clock> t0_;
+
+    bool empty() { return p_ == c_; }
+
+    bool full() { return p_ - max(c_.load(), 1LL) == R_ - 1; }
+
+    void background_thread() {
+        t0_ = chrono::high_resolution_clock::now();
+        do {
+            assert(p_ >= c_);
+            bool ok = false;
+            try {
+                ok = Produce(ring_[p_ % R_]);
+            } catch (exception &exn) {
+                errmsg_ = exn.what();
+                if (errmsg_.empty()) {
+                    errmsg_ = "unknown error on producer thread";
+                }
+                stop_ = true;
+            }
+            if (ok) {
+                ++p_;
+                if (full()) {
+                    auto t_spin = chrono::high_resolution_clock::now();
+                    do {
+                        this_thread::sleep_for(chrono::milliseconds(1));
+                    } while (!stop_ && full());
+                    p_blocked_ += chrono::high_resolution_clock::now() - t_spin;
+                }
+            } else {
+                stop_ = true;
+            }
+        } while (!stop_);
+    }
+
+  public:
+    BackgroundProducer(int ringsize) : R_(ringsize) {
+        assert(R_ > 1);
+        stop_ = false;
+        p_ = 0;
+        c_ = 0;
+    }
+
+    virtual ~BackgroundProducer() { abort(); }
+
+    // advance to next item for consumption, return false when item stream has ended successfully,
+    // or throw an exception.
+    bool next() {
+        if (!worker_) {
+            while (ring_.size() < R_) {
+                ring_.push_back(Item());
+            }
+            worker_.reset(new thread([this]() { this->background_thread(); }));
+        }
+        if (empty()) {
+            auto t_start = chrono::high_resolution_clock::now();
+            while (!stop_ && empty()) {
+                this_thread::yield();
+            }
+            c_blocked_ += chrono::high_resolution_clock::now() - t_start;
+        }
+        if (stop_ && empty()) {
+            if (!errmsg_.empty()) {
+                throw runtime_error(errmsg_);
+            }
+            return false;
+        }
+        assert(c_ < p_);
+        c_++;
+        return true;
+    }
+
+    // get current item for consumption; defined only after next() returned true
+    Item &item() {
+        assert(c_ && errmsg_.empty());
+        return ring_[(c_ - 1) % R_];
+    }
+
+    void abort() {
+        stop_ = true;
+        if (worker_) {
+            worker_->join();
+        }
+    }
+
+    string log() {
+        chrono::duration<double> elapsed = chrono::high_resolution_clock::now() - t0_;
+        OStringStream ans;
+        ans << to_string(c_) << " item(s) processed in " << to_string(elapsed.count()) << "s"
+            << "; producer blocked for " << to_string(p_blocked_.count()) << "s"
+            << "; consumer blocked for " << to_string(c_blocked_.count()) << "s";
+        return string(ans.Get());
+    }
+};

loaders/vcf_into_sqlite.cc

@@ -7,18 +7,14 @@
  * - array values are stored as JSON
  */
 #include "common.hpp"
-#include <atomic>
-#include <chrono>
 #include <cmath>
 #include <functional>
 #include <getopt.h>
 #include <htslib/vcf.h>
 #include <map>
 #include <set>
 #include <sstream>
-#include <thread>
 #include <unistd.h>
-#include <vector>
 
 // unpack each bcf_hrec_t with the key type (e.g. INFO, FORMAT) into an easier-to-use map
 vector<map<string, string>> extract_hrecs(bcf_hdr_t *hdr, const char *key,
@@ -513,98 +509,33 @@ void insert_genotypes(bcf_hdr_t *hdr, bcf1_t *rec, vector<map<string, string>> &
 }
 
 // stream BCF records using background thread
-class BCFReader {
+class BCFReader : public BackgroundProducer<shared_ptr<bcf1_t>> {
+  public:
+    BCFReader(vcfFile *vcf, bcf_hdr_t *hdr, int ringsize)
+        : BackgroundProducer<shared_ptr<bcf1_t>>(ringsize), vcf_(vcf), hdr_(hdr) {}
+
+  private:
     vcfFile *vcf_;
     bcf_hdr_t *hdr_;
 
-    vector<unique_ptr<bcf1_t, void (*)(bcf1_t *)>> ring_;
-    unique_ptr<thread> worker_;
-    atomic<bool> stop_;
-    int err_ = 0, errcode_ = 0;
-    atomic<long long> p_, // produced
-        c_;               // if c_>0 then item (c_-1)%R is currently being consumed
-    chrono::duration<double> p_spin_ = chrono::duration<double>::zero(),
-                             c_spin_ = chrono::duration<double>::zero();
-    chrono::time_point<chrono::high_resolution_clock> t0_;
-
-    void background() {
-        t0_ = chrono::high_resolution_clock::now();
-        auto R = ring_.size();
-        do {
-            assert(p_ >= c_);
-            bcf1_t *rec = ring_[p_ % R].get();
-            int ret = bcf_read(vcf_, hdr_, rec);
-            if (ret != 0) {
-                if (ret != -1 || rec->errcode) {
-                    err_ = ret;
-                    errcode_ = rec->errcode;
-                }
-                stop_ = true;
-            } else {
-                ret = bcf_unpack(rec, BCF_UN_ALL);
-                if (ret != 0) {
-                    err_ = ret;
-                    stop_ = true;
-                } else {
-                    ++p_;
-                }
-            }
-            auto t_spin = chrono::high_resolution_clock::now();
-            for (int i = 0; !stop_ && p_ - max(c_.load(), 1LL) == R - 1; ++i) {
-                this_thread::sleep_for(chrono::milliseconds(1));
-            }
-            p_spin_ += chrono::high_resolution_clock::now() - t_spin;
-        } while (!stop_);
-    }
-
-  public:
-    BCFReader(vcfFile *vcf, bcf_hdr_t *hdr, int ringsize) : vcf_(vcf), hdr_(hdr) {
-        assert(ringsize > 1);
-        stop_ = false;
-        p_ = 0;
-        c_ = 0;
-        for (int i = 0; i < ringsize; i++) {
-            ring_.emplace_back(bcf_init(), &bcf_destroy);
+    bool Produce(shared_ptr<bcf1_t> &it) override {
+        if (!it) {
+            it = shared_ptr<bcf1_t>(bcf_init(), &bcf_destroy);
         }
-    }
-
-    bcf1_t *read() {
-        if (!worker_) {
-            worker_.reset(new thread([this]() { this->background(); }));
-        }
-        auto t_spin = chrono::high_resolution_clock::now();
-        while (!stop_ && p_ == c_)
-            this_thread::yield();
-        c_spin_ += chrono::high_resolution_clock::now() - t_spin;
-        if (stop_) {
-            if (err_ || errcode_) {
-                worker_->join();
-                ostringstream msg;
-                msg << "vcf_into_sqlite: failed reading VCF; bcf_read() -> " << err_
-                    << " bcf1_t::errcode = " << errcode_ << '\n';
-                throw runtime_error(msg.str());
-            }
-            if (c_ == p_) {
-                worker_->join();
-                return nullptr;
-            }
+        int ret = bcf_read(vcf_, hdr_, it.get());
+        if (ret == -1) {
+            return false;
+        } else if (ret != 0 || it->errcode) {
+            ostringstream msg;
+            msg << "VCF parser failed: bcf1_read() -> " << ret
+                << ", bcf1_t::errcode = " << it->errcode;
+            throw runtime_error(msg.str());
         }
-        assert(c_ < p_);
-        return ring_[c_++ % ring_.size()].get();
-    }
-
-    void cancel() {
-        stop_ = true;
-        if (worker_) {
-            worker_->join();
+        ret = bcf_unpack(it.get(), BCF_UN_ALL);
+        if (ret != 0) {
+            throw runtime_error("Corrupt VCF/BCF record; bcf_unpack() -> " + to_string(ret));
         }
-    }
-
-    void log() {
-        chrono::duration<double> elapsed = chrono::high_resolution_clock::now() - t0_;
-        cerr << c_ << " record(s) processed in " << elapsed.count() << "s"
-             << "; producer thread spun for " << p_spin_.count() << "s"
-             << "; consumer thread spun for " << c_spin_.count() << "s" << endl;
+        return true;
     }
 };
 
@@ -778,8 +709,9 @@ int main(int argc, char *argv[]) {
                         << (format_hrecs.empty() ? "..." : " & genotypes...") << endl;
 
         BCFReader reader(vcf.get(), hdr.get(), 64);
-        bcf1_t *rec;
-        while ((rec = reader.read())) {
+        while (reader.next()) {
+            bcf1_t *rec = reader.item().get();
+            assert(rec);
             try {
                 insert_variant(hdr.get(), rec, info_hrecs, *insert_variant_stmt);
                 if (!format_hrecs.empty()) {
@@ -791,11 +723,11 @@ int main(int argc, char *argv[]) {
                                      *insert_genotype_stmt);
                 }
             } catch (exception &exn) {
-                reader.cancel();
+                reader.abort();
                 throw exn;
             }
         }
-        progress && (reader.log(), true);
+        progress &&cerr << reader.log() << endl;
 
         // create GRI
         if (gri) {