From 29c75e59d68db6212ed353fdf71d8d55e7b8bd62 Mon Sep 17 00:00:00 2001
From: "opensearch-trigger-bot[bot]"
 <98922864+opensearch-trigger-bot[bot]@users.noreply.github.com>
Date: Tue, 30 Apr 2024 14:30:33 -0700
Subject: [PATCH] fix number of total shards in partial search response (#543)
 (#565)

(cherry picked from commit 0402b6de3541f925c99ea29a0829ab70efc622a4)

Signed-off-by: panguixin <panguixin@bytedance.com>
Signed-off-by: Marc Handalian <marc.handalian@gmail.com>
Signed-off-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
---
 .../listener/AsynchronousSearchProgressListener.java            | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/main/java/org/opensearch/search/asynchronous/listener/AsynchronousSearchProgressListener.java b/src/main/java/org/opensearch/search/asynchronous/listener/AsynchronousSearchProgressListener.java
index adaad799..c62b9b86 100644
--- a/src/main/java/org/opensearch/search/asynchronous/listener/AsynchronousSearchProgressListener.java
+++ b/src/main/java/org/opensearch/search/asynchronous/listener/AsynchronousSearchProgressListener.java
@@ -68,7 +68,7 @@ public SearchResponse partialResponse() {
     protected void onListShards(List<SearchShard> shards, List<SearchShard> skippedShards, SearchResponse.Clusters clusters,
                                 boolean fetchPhase) {
         partialResultsHolder.hasFetchPhase.set(fetchPhase);
-        partialResultsHolder.totalShards.set(shards.size());
+        partialResultsHolder.totalShards.set(shards.size() + skippedShards.size());
         partialResultsHolder.skippedShards.set(skippedShards.size());
         partialResultsHolder.successfulShards.set(skippedShards.size());
         partialResultsHolder.clusters.set(clusters);