multiprocess.h

#ifndef _theplu_yat_utility_multi_processor_
#define _theplu_yat_utility_multi_processor_

// $Id: multiprocess.h 4076 2021-08-25 01:35:16Z peter $

/*
  Copyright (C) 2021 Peter Johansson

  This file is part of the yat library, http://dev.thep.lu.se/yat

  The yat library is free software; you can redistribute it and/or
  modify it under the terms of the GNU General Public License as
  published by the Free Software Foundation; either version 3 of the
  License, or (at your option) any later version.

  The yat library is distributed in the hope that it will be useful,
  but WITHOUT ANY WARRANTY; without even the implied warranty of
  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
  General Public License for more details.

  You should have received a copy of the GNU General Public License
  along with yat. If not, see <http://www.gnu.org/licenses/>.
*/

#include "Queue.h"
#include "yat_assert.h"

#include <iterator>
#include <map>
#include <thread>
#include <vector>

namespace theplu {
namespace yat {
namespace utility {

  template<typename InputIterator, typename OutputIterator,
           class Function, class Compare>
  void multiprocess(InputIterator begin, InputIterator end,
                    OutputIterator out, size_t cap1, size_t n_threads,
                    size_t cap2, Function function, Compare compare)
  {
    typedef typename std::iterator_traits<InputIterator>::value_type T;

    class Reader
    {
    public:
      Reader(InputIterator first, InputIterator last,
             Queue<std::shared_ptr<T>>& queue)
        : first_(first), last_(last), queue_(queue)
      {}

      void operator()(void)
      {
        for (; first_!=last_; ++first_)
          queue_.push(std::make_shared<T>(*first_));
        // send a kill pill to workers
        queue_.push(std::shared_ptr<T>());
      }
    private:
      InputIterator first_;
      InputIterator last_;
      Queue<std::shared_ptr<T>>& queue_;
    };


    class Worker
    {
    public:
      Worker(Queue<std::shared_ptr<T>>& in, const Function& func,
             Queue<std::shared_ptr<T>>& out)
        : in_(in), out_(out), function_(func)
      {}


      void operator()(void)
      {
        std::shared_ptr<T> element;
        while (true) {
          in_.pop(element);
          if (element) {
            // Do the work and if function returns true, push worked
            // element into out_ queue.
            if (function_(*element))
              out_.push(element);
          }
          else {
            // When a Worker receive a kill pill, share it with other
            // workers by pushing it into in_, and push it to the
            // Writer to inform that this queue is completed.
            in_.push(element);
            out_.push(element);
            // After that we are done, and can return home.
            return;
          }
        }
      }


    private:
      Queue<std::shared_ptr<T>>& in_;
      Queue<std::shared_ptr<T>>& out_;
      Function function_;
      size_t cap_;
    };


    class PtrCompare
    {
    public:
      PtrCompare(const Compare& c)
        : compare_(c)
      {}

      bool operator()(const std::shared_ptr<T>& lhs,
                      const std::shared_ptr<T>& rhs) const
      {
        YAT_ASSERT(lhs.get());
        YAT_ASSERT(rhs.get());
        return compare_(*lhs, *rhs);
      }
    private:
      Compare compare_;
    };


    class Writer
    {
    public:
      Writer(OutputIterator out, std::vector<Queue<std::shared_ptr<T>>>& Qs,
             const Compare& compare)
        : out_(out), queue_(Qs), buffer_(PtrCompare(compare))
      {}


      void operator()(void)
      {
        for (size_t i=0; i<queue_.size(); ++i)
          read(i);

        while (!buffer_.empty()) {
          size_t idx = buffer_.begin()->second;
          *out_ = *buffer_.begin()->first;
          ++out_;
          buffer_.erase(buffer_.begin());
          // read from the same queue erased element came from
          read(idx);
        }
      }
    private:
      void read(size_t idx)
      {
        YAT_ASSERT(idx < queue_.size());
        std::shared_ptr<T> element;
        queue_[idx].pop(element);
        // If element was not "null" insert it into buffer. If element
        // was a "null", size of buffer is efectively decreasing by
        // one untile we reach zero-size.
        if (element)
          buffer_.insert(buffer_.end(), std::make_pair(element, idx));
      }

      OutputIterator out_;
      std::vector<Queue<std::shared_ptr<T>>>& queue_;
      std::multimap<std::shared_ptr<T>, size_t, PtrCompare> buffer_;
    };

    std::vector<std::thread> threads;
    // queue used to communicate between reader and workers
    Queue<std::shared_ptr<T>> queue;
    queue.capacity(cap1);

    YAT_ASSERT(n_threads);
    // queues used to communicate between workers and writer; each
    // worker has its own queue and it's the job of the writer to
    // merge them into a sorted output.
    Queue<std::shared_ptr<T>> queue2;
    queue2.capacity(cap2);
    std::vector<Queue<std::shared_ptr<T>>> Qs(n_threads, queue2);

    // launch threads
    threads.push_back(std::thread(Reader(begin, end, queue)));
    for (size_t i=0; i<n_threads; ++i)
      threads.push_back(std::thread(Worker(queue, function, Qs[i])));
    threads.push_back(std::thread(Writer(out, Qs, compare)));

    // wait for threads to complete
    for (size_t i=0; i<threads.size(); ++i)
      threads[i].join();
  }


  template<typename InputIterator, typename OutputIterator, class Function>
  void multiprocess(InputIterator begin, InputIterator end,
                    OutputIterator out, size_t cap1, size_t n_threads,
                    size_t cap2, Function function)
  {
    typedef typename std::iterator_traits<InputIterator>::value_type T;
    std::less<T> compare;
    multiprocess(begin, end, out, cap1, n_threads, cap2, function, compare);
  }

}}} // of namespace utility, yat, and theplu

#endif