WIP

Author: kevinAlbs

src/libmongoc/doc/mongoc_server_description_service_id.rst

@@ -0,0 +1,6 @@
+:man_page: mongoc_server_description_service_id
+
+mongoc_server_description_service_id()
+======================================
+
+// LBTODO

src/libmongoc/src/mongoc/mongoc-cluster.c

@@ -2281,6 +2281,8 @@ _mongoc_cluster_stream_for_server (mongoc_cluster_t *cluster,
       _mongoc_bson_init_with_transient_txn_error (cs, reply);
    }
 
+   // LBTODO: if this is a load balanced topology and the server stream does not have a service id, disconnect and return an error.
+
    RETURN (server_stream);
 }
 

src/libmongoc/src/mongoc/mongoc-server-description.c

@@ -723,6 +723,7 @@ mongoc_server_description_handle_hello (mongoc_server_description_t *sd,
             sd, &incoming_topology_version);
          bson_destroy (&incoming_topology_version);
       }
+      // LBTODO: process serviceID
    }
 
    if (is_shard) {
@@ -1238,3 +1239,10 @@ mongoc_server_description_set_topology_version (mongoc_server_description_t *sd,
    bson_destroy (&sd->topology_version);
    bson_copy_to (tv, &sd->topology_version);
 }
+
+bool
+mongoc_server_description_service_id (const mongoc_server_description_t *description, bson_oid_t *oid) {
+   // LBTODO
+   memset (oid, 0, sizeof (bson_oid_t));
+   return false;
+}

src/libmongoc/src/mongoc/mongoc-server-description.h

@@ -66,6 +66,10 @@ MONGOC_EXPORT (int32_t)
 mongoc_server_description_compressor_id (
    const mongoc_server_description_t *description);
 
+// LBTODO
+MONGOC_EXPORT (bool)
+mongoc_server_description_service_id (const mongoc_server_description_t *description, bson_oid_t *oid);
+
 BSON_END_DECLS
 
 #endif

src/libmongoc/src/mongoc/mongoc-topology-scanner-private.h

@@ -83,6 +83,7 @@ typedef struct mongoc_topology_scanner_node {
    /* handshake_sd is a server description constructed from the response of the
     * initial handshake. It is bound to the lifetime of stream. */
    mongoc_server_description_t *handshake_sd;
+   bool loadbalanced;
 } mongoc_topology_scanner_node_t;
 
 typedef struct mongoc_topology_scanner {
@@ -241,6 +242,9 @@ void
 _mongoc_topology_scanner_set_server_api (mongoc_topology_scanner_t *ts,
                                          const mongoc_server_api_t *api);
 
+void
+_mongoc_topology_scanner_set_loadbalanced (mongoc_topology_scanner_t *ts, bool val);
+
 /* for testing. */
 mongoc_stream_t *
 _mongoc_topology_scanner_tcp_initiate (mongoc_async_cmd_t *acmd);

src/libmongoc/src/mongoc/mongoc-topology-scanner.c

@@ -260,6 +260,8 @@ _build_handshake_cmd (mongoc_topology_scanner_t *ts)
    bson_destroy (doc);
    bson_copy_to (ts->api ? &ts->hello_cmd : &ts->legacy_hello_cmd, doc);
 
+   // LBTODO: if ts->loadbalanced, append loadBalanced: true
+
    BSON_APPEND_DOCUMENT_BEGIN (doc, HANDSHAKE_FIELD, &subdoc);
    res = _mongoc_handshake_build_doc_with_application (&subdoc, ts->appname);
    bson_append_document_end (doc, &subdoc);
@@ -1398,3 +1400,10 @@ _mongoc_topology_scanner_set_server_api (mongoc_topology_scanner_t *ts,
    ts->api = mongoc_server_api_copy (api);
    _reset_hello (ts);
 }
+
+/* This must be called before the handshake command is constructed. Caller does not need to lock the topology->mutex. */
+void
+_mongoc_topology_scanner_set_loadbalanced (mongoc_topology_scanner_t *ts, bool val) {
+   BSON_ASSERT (bson_empty (&ts->handshake_cmd));
+   ts->loadbalanced = true;
+}

src/libmongoc/src/mongoc/mongoc-topology.c

@@ -434,6 +434,7 @@ mongoc_topology_new (const mongoc_uri_t *uri, bool single_threaded)
       if (topology->single_threaded) {
          _mongoc_topology_bypass_cooldown (topology);
       }
+      _mongoc_topology_scanner_set_loadbalanced (topology->scanner, true);
    } else if (service && !has_directconnection) {
       init_type = MONGOC_TOPOLOGY_UNKNOWN;
    } else if (has_directconnection) {

src/libmongoc/tests/test-mongoc-loadbalanced.c

@@ -20,6 +20,10 @@
 #include "test-libmongoc.h"
 #include "TestSuite.h"
 
+#include "mock_server/future-functions.h"
+#include "mock_server/mock-server.h"
+#include "mock_server/request.h"
+
 typedef struct {
    int server_changed_events;
    int server_opening_events;
@@ -433,6 +437,112 @@ test_loadbalanced_cooldown_is_bypassed_single (void *unused)
    free_and_assert_stats (stats);
 }
 
+/* Tests:
+ * - loadBalanced: true is added to the handshake
+ * - serviceID is set in the server description.
+ */
+#define LB_HELLO                                                               \
+   "{'ismaster': true, 'maxWireVersion': 13, 'msg': 'isdbgrid', 'serviceID', " \
+   "{'$oid': 'AAAAAAAAAAAAAAAAAAAAAAAA'}}"
+static void
+test_loadbalanced_handshake_sends_loadbalanced (void)
+{
+   mock_server_t *server;
+   mongoc_client_t *client;
+   mongoc_uri_t *uri;
+   request_t *request;
+   future_t *future;
+   bson_error_t error;
+   mongoc_server_description_t *monitor_sd;
+   mongoc_server_description_t *handshake_sd;
+   bson_oid_t expected;
+   bson_oid_t actual;
+
+   server = mock_server_new ();
+   mock_server_run (server);
+   uri = mongoc_uri_copy (mock_server_get_uri (server));
+   mongoc_uri_set_option_as_bool (uri, MONGOC_URI_LOADBALANCED, true);
+   client = mongoc_client_new_from_uri (uri);
+
+   future = future_client_command_simple (client,
+                                          "admin",
+                                          tmp_bson ("{'ping': 1}"),
+                                          NULL /* read prefs */,
+                                          NULL /* reply */,
+                                          &error);
+   request =
+      mock_server_receives_legacy_hello (server, "{'loadBalanced': true}");
+   mock_server_replies_simple (request, LB_HELLO);
+   request_destroy (request);
+
+   request = mock_server_receives_msg (server, 0, tmp_bson ("{'ping': 1}"));
+   mock_server_replies_ok_and_destroys (request);
+
+   ASSERT_OR_PRINT (future_get_bool (future), error);
+   future_destroy (future);
+
+   monitor_sd = mongoc_client_select_server (
+      client, true /* for writes */, NULL /* read prefs */, &error);
+   ASSERT_OR_PRINT (monitor_sd, error);
+   handshake_sd = mongoc_client_get_handshake_description (
+      client, 1, NULL /* opts */, &error);
+   ASSERT_OR_PRINT (handshake_sd, error);
+
+   bson_oid_init_from_string (&expected, "AAAAAAAAAAAAAAAAAAAAAAAA");
+   BSON_ASSERT (mongoc_server_description_service_id (handshake_sd, &actual));
+   ASSERT_CMPOID (&actual, &expected);
+
+   mongoc_server_description_destroy (handshake_sd);
+   mongoc_server_description_destroy (monitor_sd);
+   mongoc_uri_destroy (uri);
+   mongoc_client_destroy (client);
+   mock_server_destroy (server);
+}
+
+/* Tests that a connection is rejected if the handshake reply does not include a
+ * serviceID field. */
+#define NON_LB_HELLO \
+   "{'ismaster': true, 'maxWireVersion': 13, 'msg': 'isdbgrid'}"
+static void
+test_loadbalanced_handshake_rejects_non_loadbalanced (void)
+{
+   mock_server_t *server;
+   mongoc_client_t *client;
+   mongoc_uri_t *uri;
+   request_t *request;
+   future_t *future;
+   bson_error_t error;
+
+   server = mock_server_new ();
+   mock_server_run (server);
+   uri = mongoc_uri_copy (mock_server_get_uri (server));
+   mongoc_uri_set_option_as_bool (uri, MONGOC_URI_LOADBALANCED, true);
+   client = mongoc_client_new_from_uri (uri);
+
+   future = future_client_command_simple (client,
+                                          "admin",
+                                          tmp_bson ("{'ping': 1}"),
+                                          NULL /* read prefs */,
+                                          NULL /* reply */,
+                                          &error);
+   request =
+      mock_server_receives_legacy_hello (server, "{'loadBalanced': true}");
+   mock_server_replies_simple (request, NON_LB_HELLO);
+   request_destroy (request);
+   BSON_ASSERT (!future_get_bool (future));
+   future_destroy (future);
+
+   ASSERT_ERROR_CONTAINS (error,
+                          MONGOC_ERROR_CLIENT,
+                          123,
+                          "Driver attempted to initialize in load balancing "
+                          "mode, but the server does not support this mode");
+
+   mongoc_uri_destroy (uri);
+   mongoc_client_destroy (client);
+   mock_server_destroy (server);
+}
+
 static int
 skip_if_not_loadbalanced (void)
 {
@@ -495,4 +605,12 @@ test_loadbalanced_install (TestSuite *suite)
                       NULL /* ctx */,
                       skip_if_not_loadbalanced,
                       test_framework_skip_if_no_failpoint);
+
+   TestSuite_AddMockServerTest (suite,
+                                "/loadbalanced/handshake_sends_loadbalanced",
+                                test_loadbalanced_handshake_sends_loadbalanced);
+   TestSuite_AddMockServerTest (
+      suite,
+      "/loadbalanced/handshake_rejects_non_loadbalanced",
+      test_loadbalanced_handshake_rejects_non_loadbalanced);
 }