Group the port picking functions

2 files changed, 34 insertions, 6 deletions

diff --git a/test/core/util/port.c b/test/core/util/port.c
index 4df7572f27..b1fc722858 100644
--- a/test/core/util/port.c
+++ b/test/core/util/port.c
@@ -105,8 +105,27 @@ static void grpc_recycle_unused_port_impl(int port) {
   GPR_ASSERT(free_chosen_port(port));
 }
 
-int (*grpc_pick_unused_port)(void) = grpc_pick_unused_port_impl;
-int (*grpc_pick_unused_port_or_die)(void) = grpc_pick_unused_port_or_die_impl;
-void (*grpc_recycle_unused_port)(int port) = grpc_recycle_unused_port_impl;
+static grpc_pick_port_functions g_pick_port_functions = {
+    grpc_pick_unused_port_impl, grpc_pick_unused_port_or_die_impl,
+    grpc_recycle_unused_port_impl};
+
+int grpc_pick_unused_port(void) {
+  return g_pick_port_functions.pick_unused_port_fn();
+}
+
+int grpc_pick_unused_port_or_die(void) {
+  return g_pick_port_functions.pick_unused_port_or_die_fn();
+}
+
+void grpc_recycle_unused_port(int port) {
+  g_pick_port_functions.recycle_unused_port_fn(port);
+}
+
+void grpc_set_pick_port_functions(grpc_pick_port_functions functions) {
+  GPR_ASSERT(functions.pick_unused_port_fn != NULL);
+  GPR_ASSERT(functions.pick_unused_port_or_die_fn != NULL);
+  GPR_ASSERT(functions.recycle_unused_port_fn != NULL);
+  g_pick_port_functions = functions;
+}
 
 #endif /* GRPC_TEST_PICK_PORT */
diff --git a/test/core/util/port.h b/test/core/util/port.h
index 0e15f397a5..602099dea6 100644
--- a/test/core/util/port.h
+++ b/test/core/util/port.h
@@ -23,18 +23,27 @@
 extern "C" {
 #endif
 
+typedef struct grpc_pick_port_functions {
+  int (*pick_unused_port_fn)(void);
+  int (*pick_unused_port_or_die_fn)(void);
+  void (*recycle_unused_port_fn)(int port);
+} grpc_pick_port_functions;
+
 /* pick a port number that is currently unused by either tcp or udp. return
    0 on failure. */
-extern int (*grpc_pick_unused_port)(void);
+int grpc_pick_unused_port(void);
 /* pick a port number that is currently unused by either tcp or udp. abort
    on failure. */
-extern int (*grpc_pick_unused_port_or_die)(void);
+int grpc_pick_unused_port_or_die(void);
 
 /* Return a port which was previously returned by grpc_pick_unused_port().
  * Implementations of grpc_pick_unused_port() backed by a portserver may limit
  * the total number of ports available; this lets a binary return its allocated
  * ports back to the server if it is going to allocate a large number. */
-extern void (*grpc_recycle_unused_port)(int port);
+void grpc_recycle_unused_port(int port);
+
+/** Request the family of pick_port functions in \a functions be used. */
+void grpc_set_pick_port_functions(grpc_pick_port_functions functions);
 
 #ifdef __cplusplus
 }