From 51d1d7ca1b3927a9cd8d4f8d38b0f8a97a972c9f Mon Sep 17 00:00:00 2001
From: Craig Tiller <ctiller@google.com>
Date: Wed, 5 Jun 2024 16:37:51 -0700
Subject: [PATCH] [thready-tsan] Make sure the DNS resolver can complete
 callbacks (#36825)

Closes #36825

COPYBARA_INTEGRATE_REVIEW=https://github.com/grpc/grpc/pull/36825 from ctiller:csc 53928bdcb17aad89b520df3a54fc06c4945fff4d
PiperOrigin-RevId: 640689749
---
 .../thready_event_engine.cc                   | 27 ++++++++++---------
 .../thready_event_engine.h                    |  7 ++---
 2 files changed, 18 insertions(+), 16 deletions(-)

diff --git a/src/core/lib/event_engine/thready_event_engine/thready_event_engine.cc b/src/core/lib/event_engine/thready_event_engine/thready_event_engine.cc
index 1d80104f663..f448108a390 100644
--- a/src/core/lib/event_engine/thready_event_engine/thready_event_engine.cc
+++ b/src/core/lib/event_engine/thready_event_engine/thready_event_engine.cc
@@ -85,7 +85,9 @@ bool ThreadyEventEngine::IsWorkerThread() {
 absl::StatusOr<std::unique_ptr<EventEngine::DNSResolver>>
 ThreadyEventEngine::GetDNSResolver(
     const DNSResolver::ResolverOptions& options) {
-  return std::make_unique<ThreadyDNSResolver>(*impl_->GetDNSResolver(options));
+  return std::make_unique<ThreadyDNSResolver>(
+      *impl_->GetDNSResolver(options),
+      std::static_pointer_cast<ThreadyEventEngine>(shared_from_this()));
 }
 
 void ThreadyEventEngine::Run(Closure* closure) {
@@ -116,10 +118,10 @@ void ThreadyEventEngine::ThreadyDNSResolver::LookupHostname(
     LookupHostnameCallback on_resolve, absl::string_view name,
     absl::string_view default_port) {
   return impl_->LookupHostname(
-      [this, on_resolve = std::move(on_resolve)](
+      [engine = engine_, on_resolve = std::move(on_resolve)](
           absl::StatusOr<std::vector<ResolvedAddress>> addresses) mutable {
-        engine_->Asynchronously([on_resolve = std::move(on_resolve),
-                                 addresses = std::move(addresses)]() mutable {
+        engine->Asynchronously([on_resolve = std::move(on_resolve),
+                                addresses = std::move(addresses)]() mutable {
           on_resolve(std::move(addresses));
         });
       },
@@ -129,13 +131,12 @@ void ThreadyEventEngine::ThreadyDNSResolver::LookupHostname(
 void ThreadyEventEngine::ThreadyDNSResolver::LookupSRV(
     LookupSRVCallback on_resolve, absl::string_view name) {
   return impl_->LookupSRV(
-      [this, on_resolve = std::move(on_resolve)](
+      [engine = engine_, on_resolve = std::move(on_resolve)](
           absl::StatusOr<std::vector<SRVRecord>> records) mutable {
-        return engine_->Asynchronously(
-            [on_resolve = std::move(on_resolve),
-             records = std::move(records)]() mutable {
-              on_resolve(std::move(records));
-            });
+        return engine->Asynchronously([on_resolve = std::move(on_resolve),
+                                       records = std::move(records)]() mutable {
+          on_resolve(std::move(records));
+        });
       },
       name);
 }
@@ -143,10 +144,10 @@ void ThreadyEventEngine::ThreadyDNSResolver::LookupSRV(
 void ThreadyEventEngine::ThreadyDNSResolver::LookupTXT(
     LookupTXTCallback on_resolve, absl::string_view name) {
   return impl_->LookupTXT(
-      [this, on_resolve = std::move(on_resolve)](
+      [engine = engine_, on_resolve = std::move(on_resolve)](
           absl::StatusOr<std::vector<std::string>> record) mutable {
-        return engine_->Asynchronously([on_resolve = std::move(on_resolve),
-                                        record = std::move(record)]() mutable {
+        return engine->Asynchronously([on_resolve = std::move(on_resolve),
+                                       record = std::move(record)]() mutable {
           on_resolve(std::move(record));
         });
       },
diff --git a/src/core/lib/event_engine/thready_event_engine/thready_event_engine.h b/src/core/lib/event_engine/thready_event_engine/thready_event_engine.h
index 6e4f17002f2..8b87d90a757 100644
--- a/src/core/lib/event_engine/thready_event_engine/thready_event_engine.h
+++ b/src/core/lib/event_engine/thready_event_engine/thready_event_engine.h
@@ -77,8 +77,9 @@ class ThreadyEventEngine final : public EventEngine {
  private:
   class ThreadyDNSResolver final : public DNSResolver {
    public:
-    explicit ThreadyDNSResolver(std::unique_ptr<DNSResolver> impl)
-        : impl_(std::move(impl)) {}
+    ThreadyDNSResolver(std::unique_ptr<DNSResolver> impl,
+                       std::shared_ptr<ThreadyEventEngine> engine)
+        : impl_(std::move(impl)), engine_(std::move(engine)) {}
     void LookupHostname(LookupHostnameCallback on_resolve,
                         absl::string_view name,
                         absl::string_view default_port) override;
@@ -89,7 +90,7 @@ class ThreadyEventEngine final : public EventEngine {
 
    private:
     std::unique_ptr<DNSResolver> impl_;
-    ThreadyEventEngine* engine_;
+    std::shared_ptr<ThreadyEventEngine> engine_;
   };
 
   void Asynchronously(absl::AnyInvocable<void()> fn);