Merge pull request #12046 from y-zeng/port

Make port picking functions overridable

2 files changed, 37 insertions, 3 deletions

diff --git a/test/core/util/port.c b/test/core/util/port.c
index f430c543bd..b1fc722858 100644
--- a/test/core/util/port.c
+++ b/test/core/util/port.c
@@ -79,7 +79,7 @@ static void chose_port(int port) {
   chosen_ports[num_chosen_ports - 1] = port;
 }
 
-int grpc_pick_unused_port(void) {
+static int grpc_pick_unused_port_impl(void) {
   int port = grpc_pick_port_using_server();
   if (port != 0) {
     chose_port(port);
@@ -88,7 +88,7 @@ int grpc_pick_unused_port(void) {
   return port;
 }
 
-int grpc_pick_unused_port_or_die(void) {
+static int grpc_pick_unused_port_or_die_impl(void) {
   int port = grpc_pick_unused_port();
   if (port == 0) {
     fprintf(stderr,
@@ -101,6 +101,31 @@ int grpc_pick_unused_port_or_die(void) {
   return port;
 }
 
-void grpc_recycle_unused_port(int port) { GPR_ASSERT(free_chosen_port(port)); }
+static void grpc_recycle_unused_port_impl(int port) {
+  GPR_ASSERT(free_chosen_port(port));
+}
+
+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 154e8f830c..602099dea6 100644
--- a/test/core/util/port.h
+++ b/test/core/util/port.h
@@ -23,6 +23,12 @@
 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. */
 int grpc_pick_unused_port(void);
@@ -36,6 +42,9 @@ int grpc_pick_unused_port_or_die(void);
  * ports back to the server if it is going to allocate a large number. */
 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
 }
 #endif