From b501c2d13df80ada3574022553628b4a8255e7b6 Mon Sep 17 00:00:00 2001
From: Suryansh Gupta <suryansh.gupta2000@gmail.com>
Date: Wed, 11 Dec 2024 18:56:43 +0530
Subject: [PATCH] Convert benchmark_reads to multithreaded

---
 apps/benchmark_reads.cpp            | 65 ++++++++++++++++++++++-------
 src/windows_aligned_file_reader.cpp |  2 +-
 2 files changed, 50 insertions(+), 17 deletions(-)

diff --git a/apps/benchmark_reads.cpp b/apps/benchmark_reads.cpp
index a66218375..8a3a5f227 100644
--- a/apps/benchmark_reads.cpp
+++ b/apps/benchmark_reads.cpp
@@ -3,27 +3,22 @@
 #include "windows_aligned_file_reader.h"
 #include "aligned_file_reader.h"
 #include "utils.h"
+#include "timer.h"
+#include <omp.h>
+#include <concurrent_queue.h>
 
 using namespace std;
 using namespace diskann;
 
-#define IS_ALIGNED(X, Y) ((uint64_t)(X) % (uint64_t)(Y) == 0)
-
 #define SECTOR_LEN 4096
 
-void do_reads()
+void do_reads(WindowsAlignedFileReader *reader, char *buf)
 {
-    string file_name = "C:\\DiskANN\\Data\\turning_100k\\index_disk.index";
-    auto reader = new WindowsAlignedFileReader();
-    reader->open(file_name.c_str());
     auto ctx = reader->get_ctx();
 
     std::vector<AlignedRead> read_reqs;
     int num_sectors = 100;
 
-    char *buf = nullptr;
-    alloc_aligned((void **)&buf, num_sectors * SECTOR_LEN, SECTOR_LEN);
-
     // create read requests
     for (size_t i = 0; i < num_sectors; ++i)
     {
@@ -36,18 +31,56 @@ void do_reads()
             read_reqs.push_back(read);
     }
 
-    auto s = std::chrono::high_resolution_clock::now();
+    reader->read(read_reqs, ctx, false);
+}
+
+void do_multiple_reads_with_threads(int thread_count)
+{
+    string file_name = "C:\\DiskANN\\Data\\turning_100k\\index_disk.index";
+    auto reader = new WindowsAlignedFileReader();
+    reader->open(file_name.c_str());
+    int num_sectors = 100;
+
+    ConcurrentQueue<char*> buffer_pool;
+
+    omp_set_num_threads(thread_count);
+
+#pragma omp parallel for num_threads((int)thread_count)
+    for (int i = 0; i < thread_count; i++)
+    {
+        char *buf = nullptr;
+        alloc_aligned((void **)&buf, num_sectors * SECTOR_LEN, SECTOR_LEN);
+        buffer_pool.push(buf);
+        reader->register_thread();
+    }
+
+    Timer timer;
+#pragma omp parallel for schedule(dynamic, 1)
     for (int i = 0; i < 10000; i++)
     {
-        reader->read(read_reqs, ctx, false);
+        char *buf = buffer_pool.pop();
+        do_reads(reader, buf);
+        buffer_pool.push(buf);
     }
-    auto e = std::chrono::high_resolution_clock::now();
-    std::chrono::duration<double> diff = e - s;
-    cout << "Time taken to read: " << diff.count() << endl;
+    cout << "Time taken to read in microseconds: " << timer.elapsed() << endl;
+
+    reader->close();
 }
 
-int main()
+int main(int argc, char *argv[])
 {
+    int val = 1;
+    if (argc >= 2)
+    {
+        std::istringstream iss( argv[1] );
+
+        if (iss >> val)
+        {
+            cout<<"Got cmd argument"<<endl;
+        }
+    }
+    cout<<"Using "<<val<<" threads."<<endl;
+
     cout << "Hello World" << endl;
-    do_reads();
+    do_multiple_reads_with_threads(val);
 }
\ No newline at end of file
diff --git a/src/windows_aligned_file_reader.cpp b/src/windows_aligned_file_reader.cpp
index 3650b928a..1b77d6373 100644
--- a/src/windows_aligned_file_reader.cpp
+++ b/src/windows_aligned_file_reader.cpp
@@ -18,7 +18,7 @@ void WindowsAlignedFileReader::open(const std::string &fname)
     m_filename = fname;
 #endif
 
-    this->register_thread();
+    // this->register_thread();
 }
 
 void WindowsAlignedFileReader::close()