Gracefully handle invalid futures in JoinFutures

This CL makes `JoinFutures` return an invalid future when it receives at least one invalid future. This allows `JoinFutures` to handle invalid futures gracefully rather than crashing inside it.

PiperOrigin-RevId: 936819255

Author: junwhanahn

third_party/xla/xla/tsl/concurrency/future.cc

@@ -68,6 +68,9 @@ Future<> JoinFutures(absl::Span<const Future<>> futures) {
       std::make_shared<JoinStateless>(futures.size(), std::move(promise));
 
   for (const Future<>& future : futures) {
+    if (!future.IsValid()) [[unlikely]] {
+      return {};
+    }
     future.OnReady(
         [join](const absl::Status& status) { join->OnReady(status); });
   }

third_party/xla/xla/tsl/concurrency/future.h

@@ -83,11 +83,15 @@ struct IsFuture<Future<T>> : std::true_type {};
 
 // Returns a `Future` that will be successful if all `futures` complete
 // successfully, or return a first encountered error.
+//
+// If any of the futures is invalid, the returned future will be invalid.
 Future<> JoinFutures(absl::Span<const Future<>> futures);
 
 // Returns a `Future` that will be successful if all `futures` complete
 // successfully, or return a first encountered error. Copies values from
 // completed futures into the result vector.
+//
+// If any of the futures is invalid, the returned future will be invalid.
 template <typename T, std::enable_if_t<!std::is_void_v<T>>* = nullptr>
 Future<std::vector<T>> JoinFutures(absl::Span<const Future<T>> futures);
 
@@ -96,6 +100,8 @@ Future<std::vector<T>> JoinFutures(absl::Span<const Future<T>> futures);
 // completed futures into the result vector and leaves `futures` in move-from
 // state (for copyable `T` it still incurs a copy overhead, see `OnReady`
 // documentation for details).
+//
+// If any of the futures is invalid, the returned future will be invalid.
 template <typename T, std::enable_if_t<!std::is_void_v<T>>* = nullptr>
 Future<std::vector<T>> JoinFutures(absl::Span<Future<T>> futures);
 
@@ -124,6 +130,8 @@ Future<std::vector<T>> JoinFutures(absl::Span<Future<T>> futures);
 // If custom result type for `JoinFutures` is not defined (is void by default),
 // then the result type will be inferred as `std::tuple`. Otherwise the result
 // value of type `R` will be constructed from expanded tuple values.
+//
+// If any of the futures is invalid, the returned future will be invalid.
 template <typename R = void, typename... Futures,
           std::enable_if_t<std::conjunction_v<
               internal::IsFuture<std::decay_t<Futures>>...>>* = nullptr>
@@ -1581,6 +1589,9 @@ Future<std::vector<T>> JoinFutures(absl::Span<const Future<T>> futures) {
                                                           std::move(promise));
 
   for (size_t index = 0; index < futures.size(); ++index) {
+    if (!futures[index].IsValid()) [[unlikely]] {
+      return {};
+    }
     futures[index].OnReady([index, join](absl::StatusOr<T> value) {
       join->OnReady(index, std::move(value));
     });
@@ -1602,6 +1613,9 @@ Future<std::vector<T>> JoinFutures(absl::Span<Future<T>> futures) {
                                                           std::move(promise));
 
   for (size_t index = 0; index < futures.size(); ++index) {
+    if (!futures[index].IsValid()) [[unlikely]] {
+      return {};
+    }
     std::move(futures[index]).OnReady([index, join](absl::StatusOr<T> value) {
       join->OnReady(index, std::move(value));
     });
@@ -1729,6 +1743,10 @@ auto JoinFutures(Futures&&... futures) {
     auto join = std::make_shared<internal::JoinStatic<PromiseResult, State>>(
         sizeof...(futures), std::move(promise));
 
+    if (((!futures.IsValid()) || ...)) [[unlikely]] {
+      return tsl::Future<PromiseResult>();
+    }
+
     using Is = std::make_index_sequence<sizeof...(Futures)>;
     join->OnReady(std::move(join), Is{}, std::forward<Futures>(futures)...);
 

third_party/xla/xla/tsl/concurrency/future_test.cc

@@ -1005,6 +1005,23 @@ TEST(FutureTest, JoinErrors) {
   EXPECT_EQ(join_two.Await(), absl::InternalError("error #0"));
 }
 
+TEST(FutureTest, JoinInvalid) {
+  auto [promise0, future0] = MakePromise();
+  auto [promise1, future1] = MakePromise();
+
+  std::vector<Future<>> futures0 = {future0, {}};
+  std::vector<Future<>> futures1 = {future0, {}, future1};
+
+  auto join_one = JoinFutures(futures0);
+  EXPECT_FALSE(join_one.IsValid());
+
+  auto join_two = JoinFutures(futures1);
+  EXPECT_FALSE(join_two.IsValid());
+
+  promise0.Set();
+  promise1.Set();
+}
+
 TEST(FutureTest, JoinCopyableFutures) {
   auto [promise0, future0] = MakePromise<int32_t>();
   auto [promise1, future1] = MakePromise<int32_t>();
@@ -1156,6 +1173,16 @@ TEST(FutureTest, JoinStaticallyError) {
   EXPECT_EQ(joined.Await().status(), absl::InternalError("error0"));
 }
 
+TEST(FutureTest, JoinStaticallyInvalid) {
+  auto [promise, future] = MakePromise<int32_t>();
+
+  Future<std::tuple<int32_t, int32_t>> joined =
+      JoinFutures(future, tsl::Future<int32_t>());
+  EXPECT_FALSE(joined.IsValid());
+
+  promise.Set(absl::InternalError("error0"));
+}
+
 TEST(FutureTest, JoinStaticallyToCustomType) {
   struct TwoInts {
     TwoInts(int32_t a, int32_t b) : a(a), b(b) {}