fix: numberMatched returns null in heavy search requests

[YuriZmytrakov]

Description:

In heavy search load, the execute_search function returns numberMatched = null in small fraction of request because the search task does not wait for the count task to complete before returning the results.

This issue occurs under heavy load in approximately 5–10% of requests. It was tested using a script that executed a butch of requests. To resolve this bug, the search task must wait for the count task to complete before returning results. This can be achieved by implementing asyncio.TaskGroup, which ensures that all tasks finish execution. After introducing asyncio.TaskGroup, the numberMatched = null issue no longer occurs.

PR Checklist:

• Code is formatted and linted (run pre-commit run --all-files)
• Tests pass (run make test)
• Documentation has been updated to reflect changes, if applicable
• Changes are added to the changelog

[ZSzCF]

asyncio. tasks moved to TaskGroup just

[jonhealy1]

@YuriZmytrakov Thank you for looking into this and putting this PR together!

While I definitely understand the desire to guarantee the numberMatched field is populated, I think we need to balance this with the performance implications of forcing the API to wait for slow count queries. In large Elasticsearch/OpenSearch clusters, count queries can be significantly slower than standard searches.

Additionally, we cannot use asyncio.TaskGroup here for a few critical technical reasons:

1. Exception Handling: In the original code, if the count task fails, we safely catch it when calling count_task.result(), log the error, and still return the search hits. If a task inside a TaskGroup fails, it raises an ExceptionGroup when the context manager exits. This bypasses our downstream error handling and will crash the API with a 500 error instead of failing gracefully.
2. Aggressive Cancellation: If the count query fails, the TaskGroup will immediately cancel the search task, destroying the user's search request entirely.

Proposed Solution: A Bounded Timeout Toggle
Instead of forcing all deployments to wait infinitely (which risks hanging the API worker threads if the database stalls), the safest middle ground may be to introduce a configurable timeout variable (e.g., COUNT_TIMEOUT).

If we set the default to a fast 0.5 seconds, the API gives the count task a reasonable grace period to finish, but strictly aborts the wait if the database is struggling. Deployments that need guaranteed counts can increase this timeout, while those that prioritize raw speed can set it to 0.

Could we update this PR to implement this bounded approach? The logic would revert the TaskGroup back to create_task and look something like this:

# 1. Create the tasks independently
search_task = asyncio.create_task(
    self.client.search(
        index=index_param,
        # ... (keep existing search args)
    )
)

count_task = asyncio.create_task(
    self.client.count(
        index=index_param,
        ignore_unavailable=ignore_unavailable,
        body=count_query,
    )
)

# 2. Await the search task as the absolute priority
try:
    es_response = await search_task
except exceptions.NotFoundError:
    raise NotFoundError(f"Collections '{collection_ids}' do not exist")

# 3. Explicitly wait for count with a strict boundary
import os
count_timeout = float(os.getenv("COUNT_TIMEOUT", 0.5))

if count_timeout > 0 and not count_task.done():
    await asyncio.wait([count_task], timeout=count_timeout)

# 4. Safely extract the count if it finished
if count_task.done():
    try:
        matched = count_task.result().get("count")
    except Exception as e:
        logger.error(f"Count task failed: {e}")

This ensures we don't break our error handling, protects the API's overall response times, but provides a safe, bounded way to vastly reduce numberMatched=null responses. Let me know what you think of this approach.

[YuriZmytrakov]

@YuriZmytrakov Thank you for looking into this and putting this PR together!

While I definitely understand the desire to guarantee the numberMatched field is populated, I think we need to balance this with the performance implications of forcing the API to wait for slow count queries. In large Elasticsearch/OpenSearch clusters, count queries can be significantly slower than standard searches.

Additionally, we cannot use asyncio.TaskGroup here for a few critical technical reasons:

1. Exception Handling: In the original code, if the count task fails, we safely catch it when calling count_task.result(), log the error, and still return the search hits. If a task inside a TaskGroup fails, it raises an ExceptionGroup when the context manager exits. This bypasses our downstream error handling and will crash the API with a 500 error instead of failing gracefully.
2. Aggressive Cancellation: If the count query fails, the TaskGroup will immediately cancel the search task, destroying the user's search request entirely.

Proposed Solution: A Bounded Timeout Toggle Instead of forcing all deployments to wait infinitely (which risks hanging the API worker threads if the database stalls), the safest middle ground may be to introduce a configurable timeout variable (e.g., COUNT_TIMEOUT).

If we set the default to a fast 0.5 seconds, the API gives the count task a reasonable grace period to finish, but strictly aborts the wait if the database is struggling. Deployments that need guaranteed counts can increase this timeout, while those that prioritize raw speed can set it to 0.

Could we update this PR to implement this bounded approach? The logic would revert the TaskGroup back to create_task and look something like this:

# 1. Create the tasks independently
search_task = asyncio.create_task(
    self.client.search(
        index=index_param,
        # ... (keep existing search args)
    )
)

count_task = asyncio.create_task(
    self.client.count(
        index=index_param,
        ignore_unavailable=ignore_unavailable,
        body=count_query,
    )
)

# 2. Await the search task as the absolute priority
try:
    es_response = await search_task
except exceptions.NotFoundError:
    raise NotFoundError(f"Collections '{collection_ids}' do not exist")

# 3. Explicitly wait for count with a strict boundary
import os
count_timeout = float(os.getenv("COUNT_TIMEOUT", 0.5))

if count_timeout > 0 and not count_task.done():
    await asyncio.wait([count_task], timeout=count_timeout)

# 4. Safely extract the count if it finished
if count_task.done():
    try:
        matched = count_task.result().get("count")
    except Exception as e:
        logger.error(f"Count task failed: {e}")

This ensures we don't break our error handling, protects the API's overall response times, but provides a safe, bounded way to vastly reduce numberMatched=null responses. Let me know what you think of this approach.

Thank you, @jonhealy1 This is a very good suggestion. Although it does not fully resolve the issue of the count task, it significantly reduces the number of null values returned in /search requests. I tested this solution and observed a dramatic reduction in nulls. Adding this functionality will allow us to better calibrate the count timeout, preventing /search requests from slowing down while also minimizing null values in numberMatched. Additionally, I updated the documentation and added a unit test that simulates a delayed count task and verifies how it is handled by asyncio.wait.

[jonhealy1]

Looks great - thank you - a couple of minor things

[jonhealy1]

This should be logger.debug not print.

[jonhealy1]

Changelog entry needs to be updated.

[jonhealy1]

Thanks @YuriZmytrakov