pass a pollset_set to ALTS TSI handshaker

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());

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,

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);

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,
+    handshaker_factory->vtable->add_handshakers(
-                                                handshake_mgr);
+        handshaker_factory, args, interested_parties, handshake_mgr);
   }
 }
 

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(

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);
 }

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 */

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 */,

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,
+  GPR_ASSERT(alts_tsi_handshaker_create(
-                                        creds->handshaker_service_url, true,
+                 creds->options, c->target_name, creds->handshaker_service_url,
-                                        &handshaker) == TSI_OK);
+                 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,
+  GPR_ASSERT(alts_tsi_handshaker_create(
-                                        creds->handshaker_service_url, false,
+                 creds->options, nullptr, creds->handshaker_service_url, false,
-                                        &handshaker) == TSI_OK);
+                 sc->base.interested_parties, &handshaker) == TSI_OK);
   grpc_handshake_manager_add(handshake_manager, grpc_security_handshaker_create(
                                                     handshaker, &sc->base));
 }

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);

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

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);
 }

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()");

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

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());
 }

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);