Initial commit

packages/leann-backend-diskann/third_party/DiskANN/include/concurrent_queue.h

@@ -0,0 +1,132 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT license.
+
+#pragma once
+#include <atomic>
+#include <chrono>
+#include <condition_variable>
+#include <mutex>
+#include <queue>
+#include <thread>
+#include <type_traits>
+#include <unordered_set>
+
+namespace diskann
+{
+
+template <typename T> class ConcurrentQueue
+{
+    typedef std::chrono::microseconds chrono_us_t;
+    typedef std::unique_lock<std::mutex> mutex_locker;
+
+    std::queue<T> q;
+    std::mutex mut;
+    std::mutex push_mut;
+    std::mutex pop_mut;
+    std::condition_variable push_cv;
+    std::condition_variable pop_cv;
+    T null_T;
+
+  public:
+    ConcurrentQueue()
+    {
+    }
+
+    ConcurrentQueue(T nullT)
+    {
+        this->null_T = nullT;
+    }
+
+    ~ConcurrentQueue()
+    {
+        this->push_cv.notify_all();
+        this->pop_cv.notify_all();
+    }
+
+    // queue stats
+    uint64_t size()
+    {
+        mutex_locker lk(this->mut);
+        uint64_t ret = q.size();
+        lk.unlock();
+        return ret;
+    }
+
+    bool empty()
+    {
+        return (this->size() == 0);
+    }
+
+    // PUSH BACK
+    void push(T &new_val)
+    {
+        mutex_locker lk(this->mut);
+        this->q.push(new_val);
+        lk.unlock();
+    }
+
+    template <class Iterator> void insert(Iterator iter_begin, Iterator iter_end)
+    {
+        mutex_locker lk(this->mut);
+        for (Iterator it = iter_begin; it != iter_end; it++)
+        {
+            this->q.push(*it);
+        }
+        lk.unlock();
+    }
+
+    // POP FRONT
+    T pop()
+    {
+        mutex_locker lk(this->mut);
+        if (this->q.empty())
+        {
+            lk.unlock();
+            return this->null_T;
+        }
+        else
+        {
+            T ret = this->q.front();
+            this->q.pop();
+            // diskann::cout << "thread_id: " << std::this_thread::get_id() <<
+            // ", ctx: "
+            // << ret.ctx << "\n";
+            lk.unlock();
+            return ret;
+        }
+    }
+
+    // register for notifications
+    void wait_for_push_notify(chrono_us_t wait_time = chrono_us_t{10})
+    {
+        mutex_locker lk(this->push_mut);
+        this->push_cv.wait_for(lk, wait_time);
+        lk.unlock();
+    }
+
+    void wait_for_pop_notify(chrono_us_t wait_time = chrono_us_t{10})
+    {
+        mutex_locker lk(this->pop_mut);
+        this->pop_cv.wait_for(lk, wait_time);
+        lk.unlock();
+    }
+
+    // just notify functions
+    void push_notify_one()
+    {
+        this->push_cv.notify_one();
+    }
+    void push_notify_all()
+    {
+        this->push_cv.notify_all();
+    }
+    void pop_notify_one()
+    {
+        this->pop_cv.notify_one();
+    }
+    void pop_notify_all()
+    {
+        this->pop_cv.notify_all();
+    }
+};
+} // namespace diskann