libcruft-util/job/queue.cpp

/*
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 * Copyright 2018 Danny Robson <danny@nerdcruft.net>
 */

#include "queue.hpp"

#include "../parse.hpp"
#include "../raii.hpp"

using util::job::queue;


///////////////////////////////////////////////////////////////////////////////
static unsigned
default_parallelism (void)
{
    if (auto var = getenv ("JOB_THREADS")) {
        return util::parse<unsigned> (var);
    }

    return std::thread::hardware_concurrency ();
}


//-----------------------------------------------------------------------------
static unsigned
default_depth (void)
{
    if (auto var = getenv ("JOB_DEPTH")) {
        return util::parse<unsigned> (var);
    }

    return 1024;
}


///////////////////////////////////////////////////////////////////////////////
queue::queue ():
    queue (default_parallelism (), default_depth ())
{ ; }


//-----------------------------------------------------------------------------
queue::queue (unsigned thread_count, unsigned pool_size):
    m_tasks {
        {}, util::pool<task> {pool_size}, {}, {}
    },
    m_pending (0),
    m_threads (thread_count),
    m_doomed (0),
    m_reaper (&queue::reap, this)
{
    CHECK_GE (m_tasks.finishing.capacity (), thread_count);

    for (auto &t: m_threads)
        t = std::thread (&queue::loop, this);
}


//-----------------------------------------------------------------------------
queue::~queue ()
{
    m_stopping = true;

    // raise the semaphore enough times to resume all the worker threads
    for (size_t i = 0; i < m_threads.size (); ++i)
        m_pending.release ();

    // wait for everyone to tidy up. perhaps we'd like to use a timeout, but
    // if things deadlock then it's the users fault currently.
    for (auto &t: m_threads)
        t.join ();

    // wake the reaper up with enough tasks that it's guaranteed to resume.
    // it will bail early and drain the queue, so the validity of the increment
    // isn't super important.
    m_doomed.release (m_tasks.finishing.capacity ());
    m_reaper.join ();
}


///////////////////////////////////////////////////////////////////////////////
void
queue::loop ()
{
    while (true) {
        m_pending.acquire ();
        if (m_stopping)
            return;

        util::scoped_counter running_count (m_running);

        CHECK (!m_tasks.pending->empty ());

        auto todo = [this] () {
            auto obj = m_tasks.pending.acquire ();
            auto res = obj->front ();
            obj->pop_front ();
            return res;
        } ();

        util::scoped_function cleanup ([&, this] () {
            while (!m_tasks.finishing.push (todo))
                ;
            m_doomed.release ();
        });

        todo->function (*todo);
    }
}


//-----------------------------------------------------------------------------
void
queue::reap ()
{
    while (true) {
        // wait until we might have something to work with
        m_doomed.acquire ();
        if (m_stopping)
            break;

        // pop and notify as many doomed tasks as we can
        int count = 0;
        for (task *item; m_tasks.finishing.pop (item); ++count) {
            item->done.notify ();
            m_tasks.notified.push_back (item);
        }

        // destroy any tasks that have a zero reference count
        m_tasks.notified.erase (
            std::remove_if (
                m_tasks.notified.begin (),
                m_tasks.notified.end (),
                [&] (auto i) {
                    if (i->references.value () < 1)
                        return false;

                    m_tasks.store.destroy (i);
                    return true;
                }
            ),
            m_tasks.notified.end ()
        );

        // subtract the number of tasks we've dealt with (remembering we've
        // already claimed one in the process of waking).
        m_doomed.acquire (count - 1);
    }

    // pre-notify everyone. this lets any observers release the task before
    // we wait for them at destruction time. otherwise we tend to find
    // deadlocks amongst remaining tasks.
    std::vector<task*> doomed;
    for (auto &i: m_tasks.notified)
        m_tasks.store.destroy (i);
    for (auto &i: doomed)
        i->done.notify ();
    for (auto &i: doomed)
        m_tasks.store.destroy (i);
}