Skip to content

Commit

Permalink
Replace TSL's BlockingCounter with absl's.
Browse files Browse the repository at this point in the history
PiperOrigin-RevId: 705936067
  • Loading branch information
majnemer authored and Google-ML-Automation committed Dec 13, 2024
1 parent b119483 commit 3324edc
Show file tree
Hide file tree
Showing 2 changed files with 14 additions and 5 deletions.
2 changes: 2 additions & 0 deletions xla/tsl/platform/cloud/BUILD
Original file line number Diff line number Diff line change
Expand Up @@ -376,6 +376,8 @@ tsl_cc_test(
":now_seconds_env",
":ram_file_block_cache",
"//xla/tsl/lib/core:status_test_util",
"@com_google_absl//absl/synchronization",
"@com_google_absl//absl/time",
"@tsl//tsl/platform:blocking_counter",
"@tsl//tsl/platform:env",
"@tsl//tsl/platform:env_impl",
Expand Down
17 changes: 12 additions & 5 deletions xla/tsl/platform/cloud/ram_file_block_cache_test.cc
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,9 @@ limitations under the License.

#include <cstring>

#include "absl/synchronization/blocking_counter.h"
#include "absl/synchronization/notification.h"
#include "absl/time/time.h"
#include "xla/tsl/lib/core/status_test_util.h"
#include "xla/tsl/platform/cloud/now_seconds_env.h"
#include "tsl/platform/blocking_counter.h"
Expand Down Expand Up @@ -483,11 +486,15 @@ TEST(RamFileBlockCacheTest, ParallelReads) {
// concurrently (at which point it will respond with success to all callers),
// or 10 seconds have elapsed (at which point it will respond with an error).
const int callers = 4;
BlockingCounter counter(callers);
auto fetcher = [&counter](const string& filename, size_t offset, size_t n,
char* buffer, size_t* bytes_transferred) {
counter.DecrementCount();
if (!counter.WaitFor(std::chrono::seconds(10))) {
absl::BlockingCounter counter(callers);
absl::Notification notification;
auto fetcher = [&counter, &notification](
const string& filename, size_t offset, size_t n,
char* buffer, size_t* bytes_transferred) {
if (counter.DecrementCount()) {
notification.Notify();
}
if (!notification.WaitForNotificationWithTimeout(absl::Seconds(10))) {
// This avoids having the test time out, which is harder to debug.
return errors::FailedPrecondition("desired concurrency not reached");
}
Expand Down

0 comments on commit 3324edc

Please sign in to comment.