Skip to content

Commit

Permalink
replace callback driven wait with new Wait() method (#526)
Browse files Browse the repository at this point in the history
  • Loading branch information
hliu18 authored Mar 7, 2024
1 parent 61846c0 commit 2795f85
Showing 1 changed file with 4 additions and 31 deletions.
35 changes: 4 additions & 31 deletions src/pq_flash_index.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -1201,37 +1201,10 @@ int PQFlashIndex<T, LabelT>::load_from_separate_paths(uint32_t num_threads, cons
}

#ifdef USE_BING_INFRA
bool getNextCompletedRequest(const IOContext &ctx, size_t size, int &completedIndex)
bool getNextCompletedRequest(std::shared_ptr<AlignedFileReader> &reader, IOContext &ctx, int &completedIndex)
{
bool waitsRemaining = false;
long completeCount = ctx.m_completeCount;
do
{
for (int i = 0; i < size; i++)
{
auto ithStatus = (*ctx.m_pRequestsStatus)[i];
if (ithStatus == IOContext::Status::READ_SUCCESS)
{
completedIndex = i;
return true;
}
else if (ithStatus == IOContext::Status::READ_WAIT)
{
waitsRemaining = true;
}
}

// if we didn't find one in READ_SUCCESS, wait for one to complete.
if (waitsRemaining)
{
WaitOnAddress(&ctx.m_completeCount, &completeCount, sizeof(completeCount), 100);
// this assumes the knowledge of the reader behavior (implicit
// contract). need better factoring?
}
} while (waitsRemaining);

completedIndex = -1;
return false;
reader->wait(ctx, completedIndex);
return completedIndex != -1;
}
#endif

Expand Down Expand Up @@ -1540,7 +1513,7 @@ void PQFlashIndex<T, LabelT>::cached_beam_search(const T *query1, const uint64_t
long requestCount = static_cast<long>(frontier_read_reqs.size());
// If we issued read requests and if a read is complete or there are
// reads in wait state, then enter the while loop.
while (requestCount > 0 && getNextCompletedRequest(ctx, requestCount, completedIndex))
while (requestCount > 0 && getNextCompletedRequest(reader, ctx, completedIndex))
{
assert(completedIndex >= 0);
auto &frontier_nhood = frontier_nhoods[completedIndex];
Expand Down

0 comments on commit 2795f85

Please sign in to comment.