CDRIVER-4736 add an api to set custom sleep function (#1442)

Author: itrofimow

src/libmongoc/CMakeLists.txt

@@ -695,6 +695,7 @@ set (HEADERS
    ${PROJECT_SOURCE_DIR}/src/mongoc/mongoc-server-api.h
    ${PROJECT_SOURCE_DIR}/src/mongoc/mongoc-server-description.h
    ${PROJECT_SOURCE_DIR}/src/mongoc/mongoc-client-session.h
+   ${PROJECT_SOURCE_DIR}/src/mongoc/mongoc-sleep.h
    ${PROJECT_SOURCE_DIR}/src/mongoc/mongoc-socket.h
    ${PROJECT_SOURCE_DIR}/src/mongoc/mongoc-stream-tls-libressl.h
    ${PROJECT_SOURCE_DIR}/src/mongoc/mongoc-stream-tls-openssl.h

src/libmongoc/src/mongoc/mongoc-sleep.h

@@ -0,0 +1,40 @@
+#include "mongoc-prelude.h"
+
+#ifndef MONGOC_SLEEP_H
+#define MONGOC_SLEEP_H
+
+#include <bson/bson.h>
+
+#include "mongoc-macros.h"
+
+BSON_BEGIN_DECLS
+
+/**
+ * mongoc_usleep_func_t:
+ * @usec: Number of microseconds to sleep for.
+ * @user_data: User data provided to mongoc_client_set_usleep_impl().
+ */
+typedef void (*mongoc_usleep_func_t) (int64_t usec, void *user_data);
+
+/**
+ * mongoc_client_set_usleep_impl:
+ * @usleep_func: A function to perform microsecond sleep.
+ *
+ * Sets the function to be called to perform sleep during scanning.
+ * Returns the old function.
+ * If old_user_data is not NULL, *old_user_data is set to the old user_data.
+ * Not thread-safe.
+ * Providing a `usleep_func` that does not sleep (e.g. coroutine suspension) is
+ * not supported. Doing so is at the user's own risk.
+ */
+MONGOC_EXPORT (void)
+mongoc_client_set_usleep_impl (mongoc_client_t *client,
+                               mongoc_usleep_func_t usleep_func,
+                               void *user_data);
+
+MONGOC_EXPORT (void)
+mongoc_usleep_default_impl (int64_t usec, void *user_data);
+
+BSON_END_DECLS
+
+#endif /* MONGOC_SLEEP_H */

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

@@ -29,6 +29,7 @@
 #include "mongoc-crypt-private.h"
 #include "mongoc-ts-pool-private.h"
 #include "mongoc-shared-private.h"
+#include "mongoc-sleep.h"
 
 #define MONGOC_TOPOLOGY_MIN_HEARTBEAT_FREQUENCY_MS 500
 #define MONGOC_TOPOLOGY_SOCKET_CHECK_INTERVAL_MS 5000
@@ -212,6 +213,11 @@ typedef struct _mongoc_topology_t {
     * topology. This could occur if the URI is invalid.
     * An invalid topology does not monitor servers. */
    bool valid;
+
+   // `usleep_fn` and `usleep_data` may be overridden by
+   // `mongoc_client_set_usleep_impl`.
+   mongoc_usleep_func_t usleep_fn;
+   void *usleep_data;
 } mongoc_topology_t;
 
 mongoc_topology_t *

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

@@ -405,6 +405,7 @@ mongoc_topology_new (const mongoc_uri_t *uri, bool single_threaded)
 #endif
 
    topology = (mongoc_topology_t *) bson_malloc0 (sizeof *topology);
+   topology->usleep_fn = mongoc_usleep_default_impl;
    topology->session_pool =
       mongoc_server_session_pool_new_with_params (_server_session_init,
                                                   _server_session_destroy,
@@ -1263,8 +1264,7 @@ mongoc_topology_select_server_id (mongoc_topology_t *topology,
                   server_id = 0;
                   goto done;
                }
-
-               _mongoc_usleep (sleep_usec);
+               topology->usleep_fn (sleep_usec, topology->usleep_data);
             }
 
             /* takes up to connectTimeoutMS. sets "last_scan", clears "stale" */

src/libmongoc/src/mongoc/mongoc-util.c

@@ -30,6 +30,7 @@
 #include "mongoc-client-private.h" // WIRE_VERSION_* macros.
 #include "mongoc-client-session-private.h"
 #include "mongoc-trace-private.h"
+#include "mongoc-sleep.h"
 
 const bson_validate_flags_t _mongoc_default_insert_vflags =
    BSON_VALIDATE_UTF8 | BSON_VALIDATE_UTF8_ALLOW_NULL |
@@ -84,10 +85,20 @@ _mongoc_hex_md5 (const char *input)
    return bson_strdup (digest_str);
 }
 
+void
+mongoc_client_set_usleep_impl (mongoc_client_t *client,
+                               mongoc_usleep_func_t usleep_func,
+                               void *user_data)
+{
+   client->topology->usleep_fn = usleep_func;
+   client->topology->usleep_data = user_data;
+}
 
 void
-_mongoc_usleep (int64_t usec)
+mongoc_usleep_default_impl (int64_t usec, void *user_data)
 {
+   BSON_UNUSED (user_data);
+
 #ifdef _WIN32
    LARGE_INTEGER ft;
    HANDLE timer;
@@ -105,6 +116,13 @@ _mongoc_usleep (int64_t usec)
 #endif
 }
 
+void
+_mongoc_usleep (int64_t usec)
+{
+   mongoc_usleep_default_impl (usec, NULL);
+}
+
+
 int64_t
 _mongoc_get_real_time_ms (void)
 {

src/libmongoc/src/mongoc/mongoc.h

@@ -47,6 +47,7 @@
 #include "mongoc-handshake.h"
 #include "mongoc-opcode.h"
 #include "mongoc-log.h"
+#include "mongoc-sleep.h"
 #include "mongoc-socket.h"
 #include "mongoc-client-session.h"
 #include "mongoc-stream.h"

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

@@ -1,6 +1,11 @@
+#include <mongoc/mongoc.h>
 #include "mongoc/mongoc-util-private.h"
 #include "TestSuite.h"
 #include "test-libmongoc.h"
+#include <mock_server/mock-server.h>
+#include <mock_server/future.h>
+#include <mock_server/future-functions.h>
+#include <test-conveniences.h> // tmp_bson
 
 
 static void
@@ -16,8 +21,95 @@ test_mongoc_usleep_basic (void)
    ASSERT_CMPTIME ((int) duration, 200 * 1000);
 }
 
+static void
+custom_usleep_impl (int64_t usec, void *user_data)
+{
+   if (user_data) {
+      *(int64_t *) user_data = usec;
+   }
+}
+
+
+// `test_mongoc_usleep_custom` tests a custom sleep function set in
+// `mongoc_client_set_usleep_impl` is applied when topology scanning sleeps.
+static void
+test_mongoc_usleep_custom (void)
+{
+   mock_server_t *server = mock_server_new ();
+   mock_server_run (server);
+
+   mongoc_uri_t *uri = mongoc_uri_copy (mock_server_get_uri (server));
+   // Tell single-threaded clients to reconnect if an error occcurs.
+   mongoc_uri_set_option_as_bool (
+      uri, MONGOC_URI_SERVERSELECTIONTRYONCE, false);
+
+   mongoc_client_t *client = test_framework_client_new_from_uri (uri, NULL);
+   ASSERT (client);
+
+   // Bypass the five second cooldown to speed up test.
+   _mongoc_topology_bypass_cooldown (client->topology);
+   // Override `min_heartbeat_frequency_msec` to speed up test.
+   client->topology->min_heartbeat_frequency_msec = 50;
+
+   // Override the sleep.
+   int64_t last_sleep_dur = 0;
+   mongoc_client_set_usleep_impl (client, custom_usleep_impl, &last_sleep_dur);
+
+   bson_error_t error;
+   future_t *future = future_client_command_simple (client,
+                                                    "db",
+                                                    tmp_bson ("{'ping': 1}"),
+                                                    NULL /* read prefs */,
+                                                    NULL,
+                                                    &error);
+
+   // Client sends initial `isMaster`.
+   {
+      request_t *req = mock_server_receives_any_hello (server);
+      ASSERT (req);
+      // Fail the request.
+      reply_to_request_with_hang_up (req);
+      request_destroy (req);
+   }
+
+   // Client sleeps for `min_heartbeat_frequency_msec`, then sends another
+   // `isMaster`.
+   {
+      request_t *req = mock_server_receives_any_hello (server);
+      ASSERT (req);
+      reply_to_request_simple (
+         req,
+         tmp_str ("{ 'minWireVersion': %d, 'maxWireVersion' : %d, "
+                  "'isWritablePrimary': true}",
+                  WIRE_VERSION_MIN,
+                  WIRE_VERSION_MAX));
+      request_destroy (req);
+   }
+
+   // Expect custom sleep to have been called between making `isMaster` calls.
+   ASSERT_CMPINT64 (last_sleep_dur, >, 0);
+
+   // Client sends "ping".
+   {
+      request_t *req = mock_server_receives_msg (
+         server, MONGOC_MSG_NONE, tmp_bson ("{'ping': 1}"));
+      ASSERT (req);
+      reply_to_request_with_ok_and_destroy (req);
+   }
+
+   bool ok = future_wait (future);
+   ASSERT_OR_PRINT (ok, error);
+
+   future_destroy (future);
+   mongoc_client_destroy (client);
+   mongoc_uri_destroy (uri);
+   mock_server_destroy (server);
+}
+
 void
 test_usleep_install (TestSuite *suite)
 {
    TestSuite_Add (suite, "/Sleep/basic", test_mongoc_usleep_basic);
+   TestSuite_AddMockServerTest (
+      suite, "/Sleep/custom", test_mongoc_usleep_custom);
 }