diff --git a/test/core/security/oauth2_utils.cc b/test/core/security/oauth2_utils.cc index 5447e7ae87b..35fce7a6147 100644 --- a/test/core/security/oauth2_utils.cc +++ b/test/core/security/oauth2_utils.cc @@ -29,6 +29,7 @@ #include #include "src/core/lib/gprpp/crash.h" +#include "src/core/lib/gprpp/notification.h" #include "src/core/lib/iomgr/exec_ctx.h" #include "src/core/lib/promise/exec_ctx_wakeup_scheduler.h" #include "src/core/lib/promise/map.h" @@ -40,14 +41,14 @@ char* grpc_test_fetch_oauth2_token_with_credentials( grpc_call_credentials* creds) { grpc_core::ExecCtx exec_ctx; grpc_call_credentials::GetRequestMetadataArgs get_request_metadata_args; - + // TODO(hork): rm once GetRequestMetadata does not depend on pollsets. grpc_pollset* pollset = static_cast(gpr_zalloc(grpc_pollset_size())); gpr_mu* mu = nullptr; grpc_pollset_init(pollset, &mu); auto pops = grpc_polling_entity_create_from_pollset(pollset); bool is_done = false; - + grpc_core::Notification done; grpc_core::MemoryAllocator memory_allocator = grpc_core::MemoryAllocator(grpc_core::ResourceQuota::Default() ->memory_quota() @@ -69,7 +70,7 @@ char* grpc_test_fetch_oauth2_token_with_credentials( }); }, grpc_core::ExecCtxWakeupScheduler(), - [&is_done, &token, &initial_metadata](absl::Status result) { + [&is_done, &done, &token, &initial_metadata](absl::Status result) { is_done = true; if (!result.ok()) { gpr_log(GPR_ERROR, "Fetching token failed: %s", @@ -83,29 +84,31 @@ char* grpc_test_fetch_oauth2_token_with_credentials( .value_or("")) .c_str()); } + done.Notify(); }, arena.get(), &pops); - grpc_core::ExecCtx::Get()->Flush(); - gpr_mu_lock(mu); - while (!is_done) { - grpc_pollset_worker* worker = nullptr; - if (!GRPC_LOG_IF_ERROR( - "pollset_work", - grpc_pollset_work(grpc_polling_entity_pollset(&pops), &worker, - grpc_core::Timestamp::InfFuture()))) { - is_done = true; + if (grpc_core::IsEventEngineClientEnabled()) { + done.WaitForNotification(); + } else { + gpr_mu_lock(mu); + while (!is_done) { + grpc_pollset_worker* worker = nullptr; + if (!GRPC_LOG_IF_ERROR( + "pollset_work", + grpc_pollset_work(grpc_polling_entity_pollset(&pops), &worker, + grpc_core::Timestamp::InfFuture()))) { + is_done = true; + } } + gpr_mu_unlock(mu); } - gpr_mu_unlock(mu); - grpc_pollset_shutdown( grpc_polling_entity_pollset(&pops), GRPC_CLOSURE_CREATE([](void*, grpc_error_handle) {}, nullptr, nullptr)); grpc_core::ExecCtx::Get()->Flush(); grpc_pollset_destroy(grpc_polling_entity_pollset(&pops)); gpr_free(pollset); - return token; }