Revert "Revert "Create and pass pollset_set to ALTS TSI handshaker""

16 files changed, 60 insertions, 15 deletions

diff --git a/src/core/ext/filters/client_channel/http_connect_handshaker.cc b/src/core/ext/filters/client_channel/http_connect_handshaker.cc
index 7ce8da8c00..bfabc68c66 100644
--- a/src/core/ext/filters/client_channel/http_connect_handshaker.cc
+++ b/src/core/ext/filters/client_channel/http_connect_handshaker.cc
@@ -351,6 +351,7 @@ static grpc_handshaker* grpc_http_connect_handshaker_create() {
 
 static void handshaker_factory_add_handshakers(
     grpc_handshaker_factory* factory, const grpc_channel_args* args,
+    grpc_pollset_set* interested_parties,
     grpc_handshake_manager* handshake_mgr) {
   grpc_handshake_manager_add(handshake_mgr,
                              grpc_http_connect_handshaker_create());
diff --git a/src/core/ext/transport/chttp2/client/chttp2_connector.cc b/src/core/ext/transport/chttp2/client/chttp2_connector.cc
index 0ac84032fd..5229304fa4 100644
--- a/src/core/ext/transport/chttp2/client/chttp2_connector.cc
+++ b/src/core/ext/transport/chttp2/client/chttp2_connector.cc
@@ -160,7 +160,7 @@ static void on_handshake_done(void* arg, grpc_error* error) {
 static void start_handshake_locked(chttp2_connector* c) {
   c->handshake_mgr = grpc_handshake_manager_create();
   grpc_handshakers_add(HANDSHAKER_CLIENT, c->args.channel_args,
-                       c->handshake_mgr);
+                       c->args.interested_parties, c->handshake_mgr);
   grpc_endpoint_add_to_pollset_set(c->endpoint, c->args.interested_parties);
   grpc_handshake_manager_do_handshake(
       c->handshake_mgr, c->args.interested_parties, c->endpoint,
diff --git a/src/core/ext/transport/chttp2/server/chttp2_server.cc b/src/core/ext/transport/chttp2/server/chttp2_server.cc
index 3f8a26ae32..b95baa9191 100644
--- a/src/core/ext/transport/chttp2/server/chttp2_server.cc
+++ b/src/core/ext/transport/chttp2/server/chttp2_server.cc
@@ -67,6 +67,7 @@ typedef struct {
   grpc_timer timer;
   grpc_closure on_timeout;
   grpc_closure on_receive_settings;
+  grpc_pollset_set* interested_parties;
 } server_connection_state;
 
 static void server_connection_state_unref(
@@ -76,6 +77,9 @@ static void server_connection_state_unref(
       GRPC_CHTTP2_UNREF_TRANSPORT(connection_state->transport,
                                   "receive settings timeout");
     }
+    grpc_pollset_set_del_pollset(connection_state->interested_parties,
+                                 connection_state->accepting_pollset);
+    grpc_pollset_set_destroy(connection_state->interested_parties);
     gpr_free(connection_state);
   }
 }
@@ -189,7 +193,11 @@ static void on_accept(void* arg, grpc_endpoint* tcp,
   connection_state->accepting_pollset = accepting_pollset;
   connection_state->acceptor = acceptor;
   connection_state->handshake_mgr = handshake_mgr;
+  connection_state->interested_parties = grpc_pollset_set_create();
+  grpc_pollset_set_add_pollset(connection_state->interested_parties,
+                               connection_state->accepting_pollset);
   grpc_handshakers_add(HANDSHAKER_SERVER, state->args,
+                       connection_state->interested_parties,
                        connection_state->handshake_mgr);
   const grpc_arg* timeout_arg =
       grpc_channel_args_find(state->args, GRPC_ARG_SERVER_HANDSHAKE_TIMEOUT_MS);
diff --git a/src/core/lib/channel/handshaker_factory.cc b/src/core/lib/channel/handshaker_factory.cc
index 4fd43635b6..8ade8fe4e2 100644
--- a/src/core/lib/channel/handshaker_factory.cc
+++ b/src/core/lib/channel/handshaker_factory.cc
@@ -24,11 +24,12 @@
 
 void grpc_handshaker_factory_add_handshakers(
     grpc_handshaker_factory* handshaker_factory, const grpc_channel_args* args,
+    grpc_pollset_set* interested_parties,
     grpc_handshake_manager* handshake_mgr) {
   if (handshaker_factory != nullptr) {
     GPR_ASSERT(handshaker_factory->vtable != nullptr);
-    handshaker_factory->vtable->add_handshakers(handshaker_factory, args,
-                                                handshake_mgr);
+    handshaker_factory->vtable->add_handshakers(
+        handshaker_factory, args, interested_parties, handshake_mgr);
   }
 }
 
diff --git a/src/core/lib/channel/handshaker_factory.h b/src/core/lib/channel/handshaker_factory.h
index 3e45fcf20e..e17a678179 100644
--- a/src/core/lib/channel/handshaker_factory.h
+++ b/src/core/lib/channel/handshaker_factory.h
@@ -32,6 +32,7 @@ typedef struct grpc_handshaker_factory grpc_handshaker_factory;
 typedef struct {
   void (*add_handshakers)(grpc_handshaker_factory* handshaker_factory,
                           const grpc_channel_args* args,
+                          grpc_pollset_set* interested_parties,
                           grpc_handshake_manager* handshake_mgr);
   void (*destroy)(grpc_handshaker_factory* handshaker_factory);
 } grpc_handshaker_factory_vtable;
@@ -42,6 +43,7 @@ struct grpc_handshaker_factory {
 
 void grpc_handshaker_factory_add_handshakers(
     grpc_handshaker_factory* handshaker_factory, const grpc_channel_args* args,
+    grpc_pollset_set* interested_parties,
     grpc_handshake_manager* handshake_mgr);
 
 void grpc_handshaker_factory_destroy(
diff --git a/src/core/lib/channel/handshaker_registry.cc b/src/core/lib/channel/handshaker_registry.cc
index eec3e1b352..fbafc43e79 100644
--- a/src/core/lib/channel/handshaker_registry.cc
+++ b/src/core/lib/channel/handshaker_registry.cc
@@ -51,9 +51,11 @@ static void grpc_handshaker_factory_list_register(
 
 static void grpc_handshaker_factory_list_add_handshakers(
     grpc_handshaker_factory_list* list, const grpc_channel_args* args,
+    grpc_pollset_set* interested_parties,
     grpc_handshake_manager* handshake_mgr) {
   for (size_t i = 0; i < list->num_factories; ++i) {
-    grpc_handshaker_factory_add_handshakers(list->list[i], args, handshake_mgr);
+    grpc_handshaker_factory_add_handshakers(list->list[i], args,
+                                            interested_parties, handshake_mgr);
   }
 }
 
@@ -91,7 +93,9 @@ void grpc_handshaker_factory_register(bool at_start,
 
 void grpc_handshakers_add(grpc_handshaker_type handshaker_type,
                           const grpc_channel_args* args,
+                          grpc_pollset_set* interested_parties,
                           grpc_handshake_manager* handshake_mgr) {
   grpc_handshaker_factory_list_add_handshakers(
-      &g_handshaker_factory_lists[handshaker_type], args, handshake_mgr);
+      &g_handshaker_factory_lists[handshaker_type], args, interested_parties,
+      handshake_mgr);
 }
diff --git a/src/core/lib/channel/handshaker_registry.h b/src/core/lib/channel/handshaker_registry.h
index 82ad9c5b9a..3dd4316de6 100644
--- a/src/core/lib/channel/handshaker_registry.h
+++ b/src/core/lib/channel/handshaker_registry.h
@@ -43,6 +43,7 @@ void grpc_handshaker_factory_register(bool at_start,
 
 void grpc_handshakers_add(grpc_handshaker_type handshaker_type,
                           const grpc_channel_args* args,
+                          grpc_pollset_set* interested_parties,
                           grpc_handshake_manager* handshake_mgr);
 
 #endif /* GRPC_CORE_LIB_CHANNEL_HANDSHAKER_REGISTRY_H */
diff --git a/src/core/lib/http/httpcli_security_connector.cc b/src/core/lib/http/httpcli_security_connector.cc
index 50078c37a1..98fb7d3937 100644
--- a/src/core/lib/http/httpcli_security_connector.cc
+++ b/src/core/lib/http/httpcli_security_connector.cc
@@ -189,7 +189,8 @@ static void ssl_handshake(void* arg, grpc_endpoint* tcp, const char* host,
   grpc_arg channel_arg = grpc_security_connector_to_arg(&sc->base);
   grpc_channel_args args = {1, &channel_arg};
   c->handshake_mgr = grpc_handshake_manager_create();
-  grpc_handshakers_add(HANDSHAKER_CLIENT, &args, c->handshake_mgr);
+  grpc_handshakers_add(HANDSHAKER_CLIENT, &args,
+                       nullptr /* interested_parties */, c->handshake_mgr);
   grpc_handshake_manager_do_handshake(
       c->handshake_mgr, nullptr /* interested_parties */, tcp,
       nullptr /* channel_args */, deadline, nullptr /* acceptor */,
diff --git a/src/core/lib/security/security_connector/alts_security_connector.cc b/src/core/lib/security/security_connector/alts_security_connector.cc
index 35a787871a..f157c93171 100644
--- a/src/core/lib/security/security_connector/alts_security_connector.cc
+++ b/src/core/lib/security/security_connector/alts_security_connector.cc
@@ -70,9 +70,9 @@ static void alts_channel_add_handshakers(
   auto c = reinterpret_cast<grpc_alts_channel_security_connector*>(sc);
   grpc_alts_credentials* creds =
       reinterpret_cast<grpc_alts_credentials*>(c->base.channel_creds);
-  GPR_ASSERT(alts_tsi_handshaker_create(creds->options, c->target_name,
-                                        creds->handshaker_service_url, true,
-                                        &handshaker) == TSI_OK);
+  GPR_ASSERT(alts_tsi_handshaker_create(
+                 creds->options, c->target_name, creds->handshaker_service_url,
+                 true, sc->base.interested_parties, &handshaker) == TSI_OK);
   grpc_handshake_manager_add(handshake_manager, grpc_security_handshaker_create(
                                                     handshaker, &sc->base));
 }
@@ -84,9 +84,9 @@ static void alts_server_add_handshakers(
   auto c = reinterpret_cast<grpc_alts_server_security_connector*>(sc);
   grpc_alts_server_credentials* creds =
       reinterpret_cast<grpc_alts_server_credentials*>(c->base.server_creds);
-  GPR_ASSERT(alts_tsi_handshaker_create(creds->options, nullptr,
-                                        creds->handshaker_service_url, false,
-                                        &handshaker) == TSI_OK);
+  GPR_ASSERT(alts_tsi_handshaker_create(
+                 creds->options, nullptr, creds->handshaker_service_url, false,
+                 sc->base.interested_parties, &handshaker) == TSI_OK);
   grpc_handshake_manager_add(handshake_manager, grpc_security_handshaker_create(
                                                     handshaker, &sc->base));
 }
diff --git a/src/core/lib/security/security_connector/security_connector.cc b/src/core/lib/security/security_connector/security_connector.cc
index 6246613e7b..f4490582a3 100644
--- a/src/core/lib/security/security_connector/security_connector.cc
+++ b/src/core/lib/security/security_connector/security_connector.cc
@@ -156,6 +156,13 @@ int grpc_security_connector_cmp(grpc_security_connector* sc,
   return sc->vtable->cmp(sc, other);
 }
 
+void grpc_security_connector_set_interested_parties(
+    grpc_security_connector* sc, grpc_pollset_set* interested_parties) {
+  if (sc != nullptr) {
+    sc->interested_parties = interested_parties;
+  }
+}
+
 int grpc_channel_security_connector_cmp(grpc_channel_security_connector* sc1,
                                         grpc_channel_security_connector* sc2) {
   GPR_ASSERT(sc1->channel_creds != nullptr);
diff --git a/src/core/lib/security/security_connector/security_connector.h b/src/core/lib/security/security_connector/security_connector.h
index 67a506b576..f0172594ad 100644
--- a/src/core/lib/security/security_connector/security_connector.h
+++ b/src/core/lib/security/security_connector/security_connector.h
@@ -63,6 +63,7 @@ struct grpc_security_connector {
   const grpc_security_connector_vtable* vtable;
   gpr_refcount refcount;
   const char* url_scheme;
+  grpc_pollset_set* interested_parties;
 };
 
 /* Refcounting. */
@@ -106,6 +107,10 @@ grpc_security_connector* grpc_security_connector_from_arg(const grpc_arg* arg);
 grpc_security_connector* grpc_security_connector_find_in_args(
     const grpc_channel_args* args);
 
+/* Util to set the interested_parties whose ownership is not transferred. */
+void grpc_security_connector_set_interested_parties(
+    grpc_security_connector* sc, grpc_pollset_set* interested_parties);
+
 /* --- channel_security_connector object. ---
 
     A channel security connector object represents a way to configure the
diff --git a/src/core/lib/security/transport/security_handshaker.cc b/src/core/lib/security/transport/security_handshaker.cc
index d76d582638..f3bdf573b3 100644
--- a/src/core/lib/security/transport/security_handshaker.cc
+++ b/src/core/lib/security/transport/security_handshaker.cc
@@ -475,20 +475,30 @@ static grpc_handshaker* fail_handshaker_create() {
 
 static void client_handshaker_factory_add_handshakers(
     grpc_handshaker_factory* handshaker_factory, const grpc_channel_args* args,
+    grpc_pollset_set* interested_parties,
     grpc_handshake_manager* handshake_mgr) {
   grpc_channel_security_connector* security_connector =
       reinterpret_cast<grpc_channel_security_connector*>(
           grpc_security_connector_find_in_args(args));
+  if (security_connector != nullptr) {
+    grpc_security_connector_set_interested_parties(&security_connector->base,
+                                                   interested_parties);
+  }
   grpc_channel_security_connector_add_handshakers(security_connector,
                                                   handshake_mgr);
 }
 
 static void server_handshaker_factory_add_handshakers(
     grpc_handshaker_factory* hf, const grpc_channel_args* args,
+    grpc_pollset_set* interested_parties,
     grpc_handshake_manager* handshake_mgr) {
   grpc_server_security_connector* security_connector =
       reinterpret_cast<grpc_server_security_connector*>(
           grpc_security_connector_find_in_args(args));
+  if (security_connector != nullptr) {
+    grpc_security_connector_set_interested_parties(&security_connector->base,
+                                                   interested_parties);
+  }
   grpc_server_security_connector_add_handshakers(security_connector,
                                                  handshake_mgr);
 }
diff --git a/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc b/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc
index 34608a3de1..dfdd659b87 100644
--- a/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc
+++ b/src/core/tsi/alts/handshaker/alts_tsi_handshaker.cc
@@ -347,7 +347,8 @@ static void init_shared_resources(const char* handshaker_service_url) {
 
 tsi_result alts_tsi_handshaker_create(
     const grpc_alts_credentials_options* options, const char* target_name,
-    const char* handshaker_service_url, bool is_client, tsi_handshaker** self) {
+    const char* handshaker_service_url, bool is_client,
+    grpc_pollset_set* interested_parties, tsi_handshaker** self) {
   if (handshaker_service_url == nullptr || self == nullptr ||
       options == nullptr || (is_client && target_name == nullptr)) {
     gpr_log(GPR_ERROR, "Invalid arguments to alts_tsi_handshaker_create()");
diff --git a/src/core/tsi/alts/handshaker/alts_tsi_handshaker.h b/src/core/tsi/alts/handshaker/alts_tsi_handshaker.h
index 227b30ce53..48ce69b1da 100644
--- a/src/core/tsi/alts/handshaker/alts_tsi_handshaker.h
+++ b/src/core/tsi/alts/handshaker/alts_tsi_handshaker.h
@@ -23,6 +23,7 @@
 
 #include <grpc/grpc.h>
 
+#include "src/core/lib/iomgr/pollset_set.h"
 #include "src/core/lib/security/credentials/alts/grpc_alts_credentials_options.h"
 #include "src/core/tsi/alts_transport_security.h"
 #include "src/core/tsi/transport_security.h"
@@ -51,6 +52,7 @@ typedef struct alts_tsi_handshaker alts_tsi_handshaker;
  *   "host:port".
  * - is_client: boolean value indicating if the handshaker is used at the client
  *   (is_client = true) or server (is_client = false) side.
+ * - interested_parties: set of pollsets interested in this connection.
  * - self: address of ALTS TSI handshaker instance to be returned from the
  *   method.
  *
@@ -58,7 +60,8 @@ typedef struct alts_tsi_handshaker alts_tsi_handshaker;
  */
 tsi_result alts_tsi_handshaker_create(
     const grpc_alts_credentials_options* options, const char* target_name,
-    const char* handshaker_service_url, bool is_client, tsi_handshaker** self);
+    const char* handshaker_service_url, bool is_client,
+    grpc_pollset_set* interested_parties, tsi_handshaker** self);
 
 /**
  * This method handles handshaker response returned from ALTS handshaker
diff --git a/test/core/handshake/readahead_handshaker_server_ssl.cc b/test/core/handshake/readahead_handshaker_server_ssl.cc
index 97e9c20ee4..14d96b5d89 100644
--- a/test/core/handshake/readahead_handshaker_server_ssl.cc
+++ b/test/core/handshake/readahead_handshaker_server_ssl.cc
@@ -75,6 +75,7 @@ static grpc_handshaker* readahead_handshaker_create() {
 
 static void readahead_handshaker_factory_add_handshakers(
     grpc_handshaker_factory* hf, const grpc_channel_args* args,
+    grpc_pollset_set* interested_parties,
     grpc_handshake_manager* handshake_mgr) {
   grpc_handshake_manager_add(handshake_mgr, readahead_handshaker_create());
 }
diff --git a/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc b/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
index 85a58114ba..e9eb7e175f 100644
--- a/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
+++ b/test/core/tsi/alts/handshaker/alts_tsi_handshaker_test.cc
@@ -421,7 +421,7 @@ static tsi_handshaker* create_test_handshaker(bool used_for_success_test,
       alts_mock_handshaker_client_create(used_for_success_test);
   grpc_alts_credentials_options* options =
       grpc_alts_credentials_client_options_create();
-  alts_tsi_handshaker_create(options, "target_name", "lame", is_client,
+  alts_tsi_handshaker_create(options, "target_name", "lame", is_client, nullptr,
                              &handshaker);
   alts_tsi_handshaker* alts_handshaker =
       reinterpret_cast<alts_tsi_handshaker*>(handshaker);