faiss/impl/ThreadedIndex-inl.h

/**
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

#include <faiss/impl/FaissAssert.h>
#include <exception>
#include <iostream>

namespace faiss {

template <typename IndexT>
ThreadedIndex<IndexT>::ThreadedIndex(bool threaded)
        // 0 is default dimension
        : ThreadedIndex(0, threaded) {}

template <typename IndexT>
ThreadedIndex<IndexT>::ThreadedIndex(int d, bool threaded)
        : IndexT(d), own_fields(false), isThreaded_(threaded) {}

template <typename IndexT>
ThreadedIndex<IndexT>::~ThreadedIndex() {
    for (auto& p : indices_) {
        if (isThreaded_) {
            // should have worker thread
            FAISS_ASSERT((bool)p.second);

            // This will also flush all pending work
            p.second->stop();
            p.second->waitForThreadExit();
        } else {
            // should not have worker thread
            FAISS_ASSERT(!(bool)p.second);
        }

        if (own_fields) {
            delete p.first;
        }
    }
}

template <typename IndexT>
void ThreadedIndex<IndexT>::addIndex(IndexT* index) {
    // We inherit the dimension from the first index added to us if we don't
    // have a set dimension
    if (indices_.empty() && this->d == 0) {
        this->d = index->d;
    }

    // The new index must match our set dimension
    FAISS_THROW_IF_NOT_FMT(
            this->d == index->d,
            "addIndex: dimension mismatch for "
            "newly added index; expecting dim %d, "
            "new index has dim %d",
            this->d,
            index->d);

    if (!indices_.empty()) {
        auto& existing = indices_.front().first;

        FAISS_THROW_IF_NOT_MSG(
                index->metric_type == existing->metric_type,
                "addIndex: newly added index is "
                "of different metric type than old index");

        // Make sure this index is not duplicated
        for (auto& p : indices_) {
            FAISS_THROW_IF_NOT_MSG(
                    p.first != index,
                    "addIndex: attempting to add index "
                    "that is already in the collection");
        }
    }

    indices_.emplace_back(std::make_pair(
            index,
            std::unique_ptr<WorkerThread>(
                    isThreaded_ ? new WorkerThread : nullptr)));

    onAfterAddIndex(index);
}

template <typename IndexT>
void ThreadedIndex<IndexT>::removeIndex(IndexT* index) {
    for (auto it = indices_.begin(); it != indices_.end(); ++it) {
        if (it->first == index) {
            // This is our index; stop the worker thread before removing it,
            // to ensure that it has finished before function exit
            if (isThreaded_) {
                // should have worker thread
                FAISS_ASSERT((bool)it->second);
                it->second->stop();
                it->second->waitForThreadExit();
            } else {
                // should not have worker thread
                FAISS_ASSERT(!(bool)it->second);
            }

            indices_.erase(it);
            onAfterRemoveIndex(index);

            if (own_fields) {
                delete index;
            }

            return;
        }
    }

    // could not find our index
    FAISS_THROW_MSG("IndexReplicas::removeIndex: index not found");
}

template <typename IndexT>
void ThreadedIndex<IndexT>::runOnIndex(std::function<void(int, IndexT*)> f) {
    if (isThreaded_) {
        std::vector<std::future<bool>> v;

        for (int i = 0; i < this->indices_.size(); ++i) {
            auto& p = this->indices_[i];
            auto indexPtr = p.first;
            v.emplace_back(
                    p.second->add([f, i, indexPtr]() { f(i, indexPtr); }));
        }

        waitAndHandleFutures(v);
    } else {
        // Multiple exceptions may be thrown; gather them as we encounter them,
        // while letting everything else run to completion
        std::vector<std::pair<int, std::exception_ptr>> exceptions;

        for (int i = 0; i < this->indices_.size(); ++i) {
            auto& p = this->indices_[i];
            try {
                f(i, p.first);
            } catch (...) {
                exceptions.emplace_back(
                        std::make_pair(i, std::current_exception()));
            }
        }

        handleExceptions(exceptions);
    }
}

template <typename IndexT>
void ThreadedIndex<IndexT>::runOnIndex(
        std::function<void(int, const IndexT*)> f) const {
    const_cast<ThreadedIndex<IndexT>*>(this)->runOnIndex(
            [f](int i, IndexT* idx) { f(i, idx); });
}

template <typename IndexT>
void ThreadedIndex<IndexT>::reset() {
    runOnIndex([](int, IndexT* index) { index->reset(); });
    this->ntotal = 0;
    this->is_trained = false;
}

template <typename IndexT>
void ThreadedIndex<IndexT>::onAfterAddIndex(IndexT* index) {}

template <typename IndexT>
void ThreadedIndex<IndexT>::onAfterRemoveIndex(IndexT* index) {}

template <typename IndexT>
void ThreadedIndex<IndexT>::waitAndHandleFutures(
        std::vector<std::future<bool>>& v) {
    // Blocking wait for completion for all of the indices, capturing any
    // exceptions that are generated
    std::vector<std::pair<int, std::exception_ptr>> exceptions;

    for (int i = 0; i < v.size(); ++i) {
        auto& fut = v[i];

        try {
            fut.get();
        } catch (...) {
            exceptions.emplace_back(
                    std::make_pair(i, std::current_exception()));
        }
    }

    handleExceptions(exceptions);
}

} // namespace faiss