CDRIVER-3645 Close operation connections to removed servers (#1618)

* add regression test: check connection counts of servers removed by SRV polling

* prune clients on push

Track last known server IDs in the client pool. Before doing an expensive prune, do a cheaper check to see if server IDs changed.

* test direct topology changes

Servers may be removed from the topology in other scenarios aside from SRV polling (e.g. during SDAM)

* document expected use of pop/push

---------

Co-authored-by: Ezra Chung <[email protected]>

Author: kevinAlbs

src/libmongoc/doc/mongoc_client_pool_pop.rst

@@ -16,6 +16,10 @@ Retrieve a :symbol:`mongoc_client_t` from the client pool, or create one. The to
 
 The returned :symbol:`mongoc_client_t` must be returned to the pool with :symbol:`mongoc_client_pool_push()`.
 
+.. note::
+
+   Return a checked out :symbol:`mongoc_client_t` to the pool with :symbol:`mongoc_client_pool_push` quickly to encourage reuse of clients among threads.
+
 Parameters
 ----------
 

src/libmongoc/doc/mongoc_client_pool_try_pop.rst

@@ -14,6 +14,10 @@ Synopsis
 
 This function is identical to :symbol:`mongoc_client_pool_pop()` except it will return ``NULL`` instead of blocking for a client to become available.
 
+.. note::
+
+   Return a checked out :symbol:`mongoc_client_t` to the pool with :symbol:`mongoc_client_pool_push` quickly to encourage reuse of clients among threads.
+
 Parameters
 ----------
 

src/libmongoc/src/mongoc/mongoc-client-pool.c

@@ -17,6 +17,7 @@
 
 #include "mongoc.h"
 #include "mongoc-apm-private.h"
+#include "mongoc-array-private.h"
 #include "mongoc-counters-private.h"
 #include "mongoc-client-pool-private.h"
 #include "mongoc-client-pool.h"
@@ -52,6 +53,8 @@ struct _mongoc_client_pool_t {
    bool error_api_set;
    mongoc_server_api_t *api;
    bool client_initialized;
+   // `last_known_serverids` is a sorted array of uint32_t.
+   mongoc_array_t last_known_serverids;
 };
 
 
@@ -146,6 +149,7 @@ mongoc_client_pool_new_with_error (const mongoc_uri_t *uri, bson_error_t *error)
    }
 
    pool = (mongoc_client_pool_t *) bson_malloc0 (sizeof *pool);
+   _mongoc_array_init (&pool->last_known_serverids, sizeof (uint32_t));
    bson_mutex_init (&pool->mutex);
    mongoc_cond_init (&pool->cond);
    _mongoc_queue_init (&pool->queue);
@@ -229,6 +233,8 @@ mongoc_client_pool_destroy (mongoc_client_pool_t *pool)
    _mongoc_ssl_opts_cleanup (&pool->ssl_opts, true);
 #endif
 
+   _mongoc_array_destroy (&pool->last_known_serverids);
+
    bson_free (pool);
 
    mongoc_counter_client_pools_active_dec ();
@@ -357,6 +363,53 @@ mongoc_client_pool_try_pop (mongoc_client_pool_t *pool)
    RETURN (client);
 }
 
+typedef struct {
+   mongoc_array_t *known_server_ids;
+   mongoc_cluster_t *cluster;
+} prune_ctx;
+
+static int
+server_id_cmp (const void *a_, const void *b_)
+{
+   const uint32_t *const a = (const uint32_t *) a_;
+   const uint32_t *const b = (const uint32_t *) b_;
+
+   if (*a == *b) {
+      return 0;
+   }
+
+   return *a < *b ? -1 : 1;
+}
+
+// `maybe_prune` removes a `mongoc_cluster_node_t` if the node refers to a removed server.
+static bool
+maybe_prune (void *item, void *ctx_)
+{
+   mongoc_cluster_node_t *cn = (mongoc_cluster_node_t *) item;
+   prune_ctx *ctx = (prune_ctx *) ctx_;
+   // Get the server ID from the cluster node.
+   uint32_t server_id = cn->handshake_sd->id;
+
+   // Check if the cluster node's server ID references a removed server.
+   if (!bsearch (
+          &server_id, ctx->known_server_ids->data, ctx->known_server_ids->len, sizeof (uint32_t), server_id_cmp)) {
+      mongoc_cluster_disconnect_node (ctx->cluster, server_id);
+   }
+   return true;
+}
+
+// `prune_client` closes connections from `client` to servers not contained in `known_server_ids`.
+static void
+prune_client (mongoc_client_t *client, mongoc_array_t *known_server_ids)
+{
+   BSON_ASSERT_PARAM (client);
+   BSON_ASSERT_PARAM (known_server_ids);
+
+   mongoc_cluster_t *cluster = &client->cluster;
+   prune_ctx ctx = {.cluster = cluster, .known_server_ids = known_server_ids};
+   mongoc_set_for_each (cluster->nodes, maybe_prune, &ctx);
+}
+
 
 void
 mongoc_client_pool_push (mongoc_client_pool_t *pool, mongoc_client_t *client)
@@ -370,6 +423,48 @@ mongoc_client_pool_push (mongoc_client_pool_t *pool, mongoc_client_t *client)
    mongoc_cluster_reset_sockettimeoutms (&client->cluster);
 
    bson_mutex_lock (&pool->mutex);
+   // Check if `last_known_server_ids` needs update.
+   bool serverids_have_changed = false;
+   {
+      mongoc_array_t current_serverids;
+      _mongoc_array_init (&current_serverids, sizeof (uint32_t));
+
+      {
+         mc_shared_tpld td = mc_tpld_take_ref (pool->topology);
+         const mongoc_set_t *servers = mc_tpld_servers_const (td.ptr);
+         for (size_t i = 0; i < servers->items_len; i++) {
+            _mongoc_array_append_val (&current_serverids, servers->items[i].id);
+         }
+         mc_tpld_drop_ref (&td);
+      }
+
+      serverids_have_changed = (current_serverids.len != pool->last_known_serverids.len) ||
+                               memcmp (current_serverids.data,
+                                       pool->last_known_serverids.data,
+                                       current_serverids.len * current_serverids.element_size) != 0;
+
+      if (serverids_have_changed) {
+         _mongoc_array_destroy (&pool->last_known_serverids);
+         pool->last_known_serverids = current_serverids; // Ownership transfer.
+      } else {
+         _mongoc_array_destroy (&current_serverids);
+      }
+   }
+
+   // Check if pooled clients need to be pruned.
+   if (serverids_have_changed) {
+      // The set of last known server IDs has changed. Prune all clients in pool.
+      mongoc_queue_item_t *ptr = pool->queue.head;
+      while (ptr != NULL) {
+         prune_client ((mongoc_client_t *) ptr->data, &pool->last_known_serverids);
+         ptr = ptr->next;
+      }
+   }
+
+   // Always prune incoming client. The topology may have changed while client was checked out.
+   prune_client (client, &pool->last_known_serverids);
+
+   // Push client back into pool.
    _mongoc_queue_push_head (&pool->queue, client);
 
    if (pool->min_pool_size && _mongoc_queue_get_length (&pool->queue) > pool->min_pool_size) {

src/libmongoc/tests/TestSuite.c

@@ -1226,3 +1226,35 @@ bson_value_eq (const bson_value_t *a, const bson_value_t *b)
    bson_destroy (tmp_a);
    return ret;
 }
+
+int32_t
+get_current_connection_count (const char *host_and_port)
+{
+   char *uri_str = bson_strdup_printf ("mongodb://%s\n", host_and_port);
+   char *uri_str_with_auth = test_framework_add_user_password_from_env (uri_str);
+   mongoc_client_t *client = mongoc_client_new (uri_str_with_auth);
+   test_framework_set_ssl_opts (client);
+   bson_t *cmd = BCON_NEW ("serverStatus", BCON_INT32 (1));
+   bson_t reply;
+   bson_error_t error;
+   bool ok = mongoc_client_command_simple (client, "admin", cmd, NULL, &reply, &error);
+   if (!ok) {
+      printf ("serverStatus failed: %s\n", error.message);
+      abort ();
+   }
+   int32_t conns;
+   // Get `connections.current` from the reply.
+   {
+      bson_iter_t iter;
+      BSON_ASSERT (bson_iter_init_find (&iter, &reply, "connections"));
+      BSON_ASSERT (bson_iter_recurse (&iter, &iter));
+      BSON_ASSERT (bson_iter_find (&iter, "current"));
+      conns = bson_iter_int32 (&iter);
+   }
+   bson_destroy (&reply);
+   bson_destroy (cmd);
+   mongoc_client_destroy (client);
+   bson_free (uri_str_with_auth);
+   bson_free (uri_str);
+   return conns;
+}

src/libmongoc/tests/TestSuite.h

@@ -577,6 +577,44 @@ bson_value_eq (const bson_value_t *a, const bson_value_t *b);
    } else                                                                       \
       (void) 0
 
+
+// `get_current_connection_count` returns the server reported connection count.
+int32_t
+get_current_connection_count (const char *host_and_port);
+
+#define ASSERT_CONN_COUNT(host, expect)                           \
+   if (1) {                                                       \
+      int32_t _got = get_current_connection_count (host);         \
+      if (_got != expect) {                                       \
+         test_error ("Got unexpected connection count to %s:\n"   \
+                     "  Expected %" PRId32 ", got %" PRId32 "\n", \
+                     host,                                        \
+                     expect,                                      \
+                     _got);                                       \
+      }                                                           \
+   } else                                                         \
+      (void) 0
+
+#define ASSERT_EVENTUAL_CONN_COUNT(host, expect)                                   \
+   if (1) {                                                                        \
+      int64_t _start = bson_get_monotonic_time ();                                 \
+      while (true) {                                                               \
+         int32_t _got = get_current_connection_count (host);                       \
+         if (_got == expect) {                                                     \
+            break;                                                                 \
+         }                                                                         \
+         int64_t _now = bson_get_monotonic_time ();                                \
+         if (_now - _start > 5 * 1000 * 1000 /* five seconds */) {                 \
+            test_error ("Timed out waiting for expected connection count to %s:\n" \
+                        "  Expected %" PRId32 ", got %" PRId32 "\n",               \
+                        host,                                                      \
+                        expect,                                                    \
+                        _got);                                                     \
+         }                                                                         \
+      }                                                                            \
+   } else                                                                          \
+      (void) 0
+
 #define MAX_TEST_NAME_LENGTH 500
 #define MAX_TEST_CHECK_FUNCS 10