Change db_connect to allow different arguments from different processes. (#142)

* Allow db_connect to take a variable number of arguments.

* Fix tests.

* Fixes.

* Formatting.

* Fixes.

* Simplifications.

* Fix typo.

src/photon/photon/photon_services.py

@@ -10,7 +10,7 @@ import time
 def random_name():
   return str(random.randint(0, 99999999))
 
-def start_local_scheduler(plasma_store_name, plasma_manager_name=None, plasma_address=None, redis_address=None, use_valgrind=False, use_profiler=False):
+def start_local_scheduler(plasma_store_name, plasma_manager_name=None, plasma_address=None, node_ip_address="127.0.0.1", redis_address=None, use_valgrind=False, use_profiler=False):
   """Start a local scheduler process.
 
   Args:
@@ -21,6 +21,8 @@ def start_local_scheduler(plasma_store_name, plasma_manager_name=None, plasma_ad
     plasma_address (str): The address of the plasma manager to connect to. This
       is only used by the global scheduler to figure out which plasma managers
       are connected to which local schedulers.
+    node_ip_address (str): The address of the node that this local scheduler is
+      running on.
     redis_address (str): The address of the Redis instance to connect to. If
       this is not provided, then the local scheduler will not connect to Redis.
     use_valgrind (bool): True if the local scheduler should be started inside of
@@ -38,7 +40,7 @@ def start_local_scheduler(plasma_store_name, plasma_manager_name=None, plasma_ad
     raise Exception("Cannot use valgrind and profiler at the same time.")
   local_scheduler_executable = os.path.join(os.path.dirname(os.path.abspath(__file__)), "../build/photon_scheduler")
   local_scheduler_name = "/tmp/scheduler{}".format(random_name())
-  command = [local_scheduler_executable, "-s", local_scheduler_name, "-p", plasma_store_name]
+  command = [local_scheduler_executable, "-s", local_scheduler_name, "-p", plasma_store_name, "-h", node_ip_address]
   if plasma_manager_name is not None:
     command += ["-m", plasma_manager_name]
   if redis_address is not None:

src/photon/photon_scheduler.c

@@ -25,6 +25,7 @@ UT_icd worker_icd = {sizeof(worker), NULL, NULL, NULL};
 UT_icd byte_icd = {sizeof(uint8_t), NULL, NULL, NULL};
 
 local_scheduler_state *init_local_scheduler(
+    const char *node_ip_address,
     event_loop *loop,
     const char *redis_addr,
     int redis_port,
@@ -39,8 +40,19 @@ local_scheduler_state *init_local_scheduler(
   utarray_new(state->workers, &worker_icd);
   /* Connect to Redis if a Redis address is provided. */
   if (redis_addr != NULL) {
-    state->db = db_connect_extended(redis_addr, redis_port, "photon", "", -1,
-                                    plasma_manager_address);
+    int num_args = 0;
+    const char **db_connect_args = NULL;
+    if (plasma_manager_address != NULL) {
+      num_args = 2;
+      db_connect_args = malloc(sizeof(char *) * num_args);
+      db_connect_args[0] = "aux_address";
+      db_connect_args[1] = plasma_manager_address;
+    }
+    state->db = db_connect(redis_addr, redis_port, "photon", node_ip_address,
+                           num_args, db_connect_args);
+    if (num_args != 0) {
+      free(db_connect_args);
+    };
     db_attach(state->db, loop, false);
   } else {
     state->db = NULL;
@@ -273,7 +285,8 @@ void handle_task_scheduled_callback(task *original_task, void *user_context) {
                         task_task_spec(original_task));
 }
 
-void start_server(const char *socket_name,
+void start_server(const char *node_ip_address,
+                  const char *socket_name,
                   const char *redis_addr,
                   int redis_port,
                   const char *plasma_store_socket_name,
@@ -283,7 +296,7 @@ void start_server(const char *socket_name,
   int fd = bind_ipc_sock(socket_name, true);
   event_loop *loop = event_loop_create();
   g_state =
-      init_local_scheduler(loop, redis_addr, redis_port,
+      init_local_scheduler(node_ip_address, loop, redis_addr, redis_port,
                            plasma_store_socket_name, plasma_manager_socket_name,
                            plasma_manager_address, global_scheduler_exists);
 
@@ -323,9 +336,11 @@ int main(int argc, char *argv[]) {
   char *plasma_manager_socket_name = NULL;
   /* Address for the plasma manager associated with this Photon instance. */
   char *plasma_manager_address = NULL;
+  /* The IP address of the node that this local scheduler is running on. */
+  char *node_ip_address = NULL;
   int c;
   bool global_scheduler_exists = true;
-  while ((c = getopt(argc, argv, "s:r:p:m:ga:")) != -1) {
+  while ((c = getopt(argc, argv, "s:r:p:m:ga:h:")) != -1) {
     switch (c) {
     case 's':
       scheduler_socket_name = optarg;
@@ -345,6 +360,9 @@ int main(int argc, char *argv[]) {
     case 'a':
       plasma_manager_address = optarg;
       break;
+    case 'h':
+      node_ip_address = optarg;
+      break;
     default:
       LOG_FATAL("unknown option %c", c);
     }
@@ -356,6 +374,9 @@ int main(int argc, char *argv[]) {
     LOG_FATAL(
         "please specify socket for connecting to Plasma store with -p switch");
   }
+  if (!node_ip_address) {
+    LOG_FATAL("please specify the node IP address with -p switch");
+  }
   if (!redis_addr_port) {
     /* Start the local scheduler without connecting to Redis. In this case, all
      * submitted tasks will be queued and scheduled locally. */
@@ -364,8 +385,9 @@ int main(int argc, char *argv[]) {
           "if a plasma manager socket name is provided with the -m switch, "
           "then a redis address must be provided with the -r switch");
     }
-    start_server(scheduler_socket_name, NULL, -1, plasma_store_socket_name,
-                 NULL, plasma_manager_address, global_scheduler_exists);
+    start_server(node_ip_address, scheduler_socket_name, NULL, -1,
+                 plasma_store_socket_name, NULL, plasma_manager_address,
+                 global_scheduler_exists);
   } else {
     /* Parse the Redis address into an IP address and a port. */
     char redis_addr[16] = {0};
@@ -382,9 +404,10 @@ int main(int argc, char *argv[]) {
           "please specify socket for connecting to Plasma manager with -m "
           "switch");
     }
-    start_server(scheduler_socket_name, &redis_addr[0], atoi(redis_port),
-                 plasma_store_socket_name, plasma_manager_socket_name,
-                 plasma_manager_address, global_scheduler_exists);
+    start_server(node_ip_address, scheduler_socket_name, &redis_addr[0],
+                 atoi(redis_port), plasma_store_socket_name,
+                 plasma_manager_socket_name, plasma_manager_address,
+                 global_scheduler_exists);
   }
 }
 #endif

src/photon/photon_scheduler.h

@@ -65,6 +65,7 @@ void reconstruct_object(local_scheduler_state *state, object_id object_id);
 /** The following methods are for testing purposes only. */
 #ifdef PHOTON_TEST
 local_scheduler_state *init_local_scheduler(
+    const char *node_ip_address,
     event_loop *loop,
     const char *redis_addr,
     int redis_port,

src/photon/test/photon_tests.c

@@ -58,7 +58,7 @@ photon_mock *init_photon_mock() {
       bind_ipc_sock_retry(photon_socket_name_format, &mock->photon_fd);
   CHECK(mock->plasma_fd >= 0 && mock->photon_fd >= 0);
   mock->photon_state = init_local_scheduler(
-      mock->loop, redis_addr, redis_port,
+      "127.0.0.1", mock->loop, redis_addr, redis_port,
       utstring_body(plasma_manager_socket_name),
       utstring_body(plasma_store_socket_name), NULL, false);
   /* Connect a Photon client. */
@@ -228,8 +228,14 @@ TEST object_reconstruction_suppression_test(void) {
     destroy_photon_mock(photon);
     exit(0);
   } else {
-    object_table_add(photon->photon_state->db, return_id, 1,
-                     (unsigned char *) NIL_DIGEST, (retry_info *) &photon_retry,
+    /* Connect a plasma manager client so we can call object_table_add. */
+    const char *db_connect_args[] = {"address", "127.0.0.1:12346"};
+    db_handle *db = db_connect("127.0.0.1", 6379, "plasma_manager", "127.0.0.1",
+                               2, db_connect_args);
+    db_attach(db, photon->loop, false);
+    /* Add the object to the object table. */
+    object_table_add(db, return_id, 1, (unsigned char *) NIL_DIGEST,
+                     (retry_info *) &photon_retry,
                      object_reconstruction_suppression_callback,
                      (void *) photon);
     /* Run the event loop. NOTE: OSX appears to require the parent process to
@@ -242,6 +248,7 @@ TEST object_reconstruction_suppression_test(void) {
     wait(NULL);
     ASSERT_EQ(num_tasks_in_queue(photon->photon_state->algorithm_state), 0);
     free_task_spec(object_reconstruction_suppression_spec);
+    db_disconnect(db);
     destroy_photon_mock(photon);
     PASS();
   }