From 56d7eef96baaf8365b258217b866f37159abf552 Mon Sep 17 00:00:00 2001
From: Oliver Reiche <oliver.reiche@huawei.com>
Date: Thu, 22 Aug 2024 17:07:29 +0200
Subject: CasClient: Retry batch update for missing response

The remote execution protocol is a bit unclear about how to
deal with blob updates for which we got no response. While
some clients consider a blob update failed only if a failed
response is received, we are going extra defensive here and
also consider missing responses to be a failed blob update.
Issue a retry for the missing blobs.
---
 .../remote/bazel/bazel_cas_client.cpp              | 24 ++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
index 552fab3c..aa6b7c9f 100644
--- a/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
+++ b/src/buildtool/execution_api/remote/bazel/bazel_cas_client.cpp
@@ -489,6 +489,7 @@ auto BazelCasClient::FindMissingBlobs(std::string const& instance_name,
     return result;
 }
 
+// NOLINTNEXTLINE(misc-no-recursion)
 auto BazelCasClient::BatchUpdateBlobs(
     std::string const& instance_name,
     std::vector<gsl::not_null<BazelBlob const*>>::const_iterator const& begin,
@@ -569,6 +570,29 @@ auto BazelCasClient::BatchUpdateBlobs(
         return oss.str();
     });
 
+    auto missing = std::distance(begin, end) - result.size();
+    if (not result.empty() and missing > 0) {
+        // The remote execution protocol is a bit unclear about how to deal with
+        // blob updates for which we got no response. While some clients
+        // consider a blob update failed only if a failed response is received,
+        // we are going extra defensive here and also consider missing responses
+        // to be a failed blob update. Issue a retry for the missing blobs.
+        logger_.Emit(LogLevel::Trace, "Retrying with missing blobs");
+        auto received_digests =
+            std::unordered_set<bazel_re::Digest>{result.begin(), result.end()};
+        auto missing_blobs = std::vector<gsl::not_null<BazelBlob const*>>{};
+        missing_blobs.reserve(missing);
+        std::for_each(
+            begin, end, [&received_digests, &missing_blobs](auto const& blob) {
+                if (not received_digests.contains(blob->digest)) {
+                    missing_blobs.emplace_back(blob);
+                }
+            });
+        return result.size() + BatchUpdateBlobs(instance_name,
+                                                missing_blobs.begin(),
+                                                missing_blobs.end());
+    }
+
     return result.size();
 }
 
-- 
cgit v1.2.3