81 lines
2.4 KiB
C++
81 lines
2.4 KiB
C++
/**
|
|
* 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.
|
|
*/
|
|
|
|
#pragma once
|
|
|
|
#include "Index.h"
|
|
#include "IndexBinary.h"
|
|
#include "WorkerThread.h"
|
|
#include <memory>
|
|
#include <vector>
|
|
|
|
namespace faiss {
|
|
|
|
/// A holder of indices in a collection of threads
|
|
/// The interface to this class itself is not thread safe
|
|
template <typename IndexT>
|
|
class ThreadedIndex : public IndexT {
|
|
public:
|
|
explicit ThreadedIndex(bool threaded);
|
|
explicit ThreadedIndex(int d, bool threaded);
|
|
|
|
~ThreadedIndex() override;
|
|
|
|
/// override an index that is managed by ourselves.
|
|
/// WARNING: once an index is added, it becomes unsafe to touch it from any
|
|
/// other thread than that on which is managing it, until we are shut
|
|
/// down. Use runOnIndex to perform work on it instead.
|
|
void addIndex(IndexT* index);
|
|
|
|
/// Remove an index that is managed by ourselves.
|
|
/// This will flush all pending work on that index, and then shut
|
|
/// down its managing thread, and will remove the index.
|
|
void removeIndex(IndexT* index);
|
|
|
|
/// Run a function on all indices, in the thread that the index is
|
|
/// managed in.
|
|
/// Function arguments are (index in collection, index pointer)
|
|
void runOnIndex(std::function<void(int, IndexT*)> f);
|
|
void runOnIndex(std::function<void(int, const IndexT*)> f) const;
|
|
|
|
/// faiss::Index API
|
|
/// All indices receive the same call
|
|
void reset() override;
|
|
|
|
/// Returns the number of sub-indices
|
|
int count() const { return indices_.size(); }
|
|
|
|
/// Returns the i-th sub-index
|
|
IndexT* at(int i) { return indices_[i].first; }
|
|
|
|
/// Returns the i-th sub-index (const version)
|
|
const IndexT* at(int i) const { return indices_[i].first; }
|
|
|
|
/// Whether or not we are responsible for deleting our contained indices
|
|
bool own_fields;
|
|
|
|
protected:
|
|
/// Called just after an index is added
|
|
virtual void onAfterAddIndex(IndexT* index);
|
|
|
|
/// Called just after an index is removed
|
|
virtual void onAfterRemoveIndex(IndexT* index);
|
|
|
|
protected:
|
|
static void waitAndHandleFutures(std::vector<std::future<bool>>& v);
|
|
|
|
/// Collection of Index instances, with their managing worker thread if any
|
|
std::vector<std::pair<IndexT*, std::unique_ptr<WorkerThread>>> indices_;
|
|
|
|
/// Is this index multi-threaded?
|
|
bool isThreaded_;
|
|
};
|
|
|
|
} // namespace
|
|
|
|
#include "ThreadedIndex-inl.h"
|