adding OP_QUERY rpc exhaust support

hanumantmk · hanumantmk · commit 2451e7daebac · 2013-12-04T18:07:29.000-05:00
Provide suport for exhaust queries

Found a couple of bugs in:
* buffers didn't handle already buffered data well (negative unsigned
  ints
* gridfs files were still on possibly negative min_bytes for readv
diff --git a/doc/mongoc_gridfs_file_readv.txt b/doc/mongoc_gridfs_file_readv.txt
@@ -15,7 +15,7 @@ ssize_t
 mongoc_gridfs_file_readv (mongoc_gridfs_file_t *file,
                           struct iovec         *iov,
                           size_t                iovcnt,
-                          ssize_t               min_bytes,
+                          size_t                min_bytes,
                           bson_uint32_t         timeout_msec);
 
 ssize_t
diff --git a/mongoc/mongoc-buffer.c b/mongoc/mongoc-buffer.c
@@ -156,7 +156,7 @@ _mongoc_buffer_append_from_stream (mongoc_buffer_t *buffer,
    }
 
    buf = &buffer->data[buffer->off + buffer->len];
-   ret = mongoc_stream_read(stream, buf, size, -1, timeout_msec);
+   ret = mongoc_stream_read(stream, buf, size, size, timeout_msec);
    if (ret != size) {
       bson_set_error(error,
                      MONGOC_ERROR_STREAM,
@@ -200,11 +200,12 @@ _mongoc_buffer_fill (mongoc_buffer_t *buffer,
    BSON_ASSERT(buffer->data);
    BSON_ASSERT(buffer->datalen);
 
-   min_bytes -= buffer->len;
-   if (min_bytes <= 0) {
+   if (min_bytes <= buffer->len) {
       return buffer->len;
    }
 
+   min_bytes -= buffer->len;
+
    if (buffer->len) {
       memmove(&buffer->data[0], &buffer->data[buffer->off], buffer->len);
    }
diff --git a/mongoc/mongoc-client-private.h b/mongoc/mongoc-client-private.h
@@ -45,6 +45,7 @@ struct _mongoc_client_t
    mongoc_list_t             *conns;
    mongoc_uri_t              *uri;
    mongoc_cluster_t           cluster;
+   bson_bool_t                in_exhaust;
 
    mongoc_stream_initiator_t  initiator;
    void                      *initiator_data;
diff --git a/mongoc/mongoc-client.c b/mongoc/mongoc-client.c
@@ -404,6 +404,14 @@ _mongoc_client_sendv (mongoc_client_t              *client,
    bson_return_val_if_fail(rpcs, FALSE);
    bson_return_val_if_fail(rpcs_len, FALSE);
 
+   if (client->in_exhaust) {
+      bson_set_error(error,
+                     MONGOC_ERROR_CLIENT,
+                     MONGOC_ERROR_CLIENT_IN_EXHAUST,
+                     "A cursor derived from this client is in exhaust.");
+      RETURN(FALSE);
+   }
+
    for (i = 0; i < rpcs_len; i++) {
       rpcs[i].header.msg_len = 0;
       rpcs[i].header.request_id = ++client->request_id;
@@ -1073,6 +1081,8 @@ _mongoc_client_warm_up (mongoc_client_t *client,
       ret = _mongoc_cluster_command_early (&client->cluster, "admin", &cmd,
                                            NULL, error);
       bson_destroy (&cmd);
+   } else if (client->cluster.state == MONGOC_CLUSTER_STATE_DEAD) {
+      ret = _mongoc_cluster_reconnect(&client->cluster, error);
    }
 
    return ret;
diff --git a/mongoc/mongoc-cluster-private.h b/mongoc/mongoc-cluster-private.h
@@ -90,6 +90,7 @@ typedef struct
    bson_bool_t             requires_auth : 1;
 
    bson_int32_t            wire_version;
+   bson_bool_t             isdbgrid;
 
    mongoc_cluster_node_t   nodes[MONGOC_CLUSTER_MAX_NODES];
    mongoc_client_t        *client;
@@ -164,6 +165,15 @@ _mongoc_cluster_command_early (mongoc_cluster_t *cluster,
                                bson_error_t     *error)
    BSON_GNUC_INTERNAL;
 
+void
+_mongoc_cluster_disconnect_node (mongoc_cluster_t      *cluster,
+                                 mongoc_cluster_node_t *node)
+   BSON_GNUC_INTERNAL;
+
+bson_bool_t
+_mongoc_cluster_reconnect (mongoc_cluster_t *cluster,
+                           bson_error_t     *error)
+   BSON_GNUC_INTERNAL;
 
 BSON_END_DECLS
 
diff --git a/mongoc/mongoc-cluster.c b/mongoc/mongoc-cluster.c
@@ -458,7 +458,7 @@ _mongoc_cluster_build_basic_auth_digest (mongoc_cluster_t *cluster,
  *--------------------------------------------------------------------------
  */
 
-static void
+void
 _mongoc_cluster_disconnect_node (mongoc_cluster_t      *cluster,
                                  mongoc_cluster_node_t *node)
 {
@@ -1044,13 +1044,18 @@ _mongoc_cluster_ismaster (mongoc_cluster_t      *cluster,
    if (bson_iter_init_find (&iter, &reply, "msg") &&
        BSON_ITER_HOLDS_UTF8 (&iter) &&
        (strcmp ("isdbgrid", bson_iter_utf8 (&iter, NULL)) == 0)) {
+      /* TODO: is this sufficient to detect sharded clusters? */
+
+      cluster->isdbgrid = TRUE;
       /*
        * TODO: This is actually a sharded cluster!
        */
       if (cluster->mode != MONGOC_CLUSTER_SHARDED_CLUSTER) {
          MONGOC_INFO ("Unexpectedly connected to sharded cluster: %s",
                       node->host.host_and_port);
       }
+   } else {
+      cluster->isdbgrid = FALSE;
    }
 
    /*
@@ -1645,7 +1650,7 @@ _mongoc_cluster_reconnect_sharded_cluster (mongoc_cluster_t *cluster,
  *--------------------------------------------------------------------------
  */
 
-static bson_bool_t
+bson_bool_t
 _mongoc_cluster_reconnect (mongoc_cluster_t *cluster,
                            bson_error_t     *error)
 {
diff --git a/mongoc/mongoc-cursor-private.h b/mongoc/mongoc-cursor-private.h
@@ -60,6 +60,7 @@ struct _mongoc_cursor_t
    bson_bool_t          done         : 1;
    bson_bool_t          failed       : 1;
    bson_bool_t          end_of_event : 1;
+   bson_bool_t          in_exhaust   : 1;
 
    bson_t               query;
    bson_t               fields;
diff --git a/mongoc/mongoc-cursor.c b/mongoc/mongoc-cursor.c
@@ -70,9 +70,15 @@ _mongoc_cursor_new (mongoc_client_t           *client,
 
    ENTRY;
 
-   bson_return_val_if_fail(client, NULL);
-   bson_return_val_if_fail(db_and_collection, NULL);
-   bson_return_val_if_fail(query, NULL);
+   BSON_ASSERT(client);
+   BSON_ASSERT(db_and_collection);
+   BSON_ASSERT(query);
+
+   /* we can't have exhaust queries with limits */
+   BSON_ASSERT (!((flags & MONGOC_QUERY_EXHAUST) && limit));
+
+   /* we can't have exhaust queries with sharded clusters */
+   BSON_ASSERT (!((flags & MONGOC_QUERY_EXHAUST) && client->cluster.isdbgrid));
 
    /*
     * Cursors execute their query lazily. This sadly means that we must copy
@@ -180,7 +186,15 @@ _mongoc_cursor_destroy (mongoc_cursor_t *cursor)
 
    bson_return_if_fail(cursor);
 
-   if (cursor->rpc.reply.cursor_id) {
+   if (cursor->in_exhaust) {
+      cursor->client->in_exhaust = FALSE;
+
+      if (!cursor->done) {
+         _mongoc_cluster_disconnect_node (
+            &cursor->client->cluster,
+            &cursor->client->cluster.nodes[cursor->hint - 1]);
+      }
+   } else if (cursor->rpc.reply.cursor_id) {
       _mongoc_cursor_kill_cursor(cursor, cursor->rpc.reply.cursor_id);
    }
 
@@ -361,6 +375,11 @@ _mongoc_cursor_query (mongoc_cursor_t *cursor)
    cursor->reader = bson_reader_new_from_data(cursor->rpc.reply.documents,
                                               cursor->rpc.reply.documents_len);
 
+   if (cursor->flags & MONGOC_QUERY_EXHAUST) {
+      cursor->in_exhaust = TRUE;
+      cursor->client->in_exhaust = TRUE;
+   }
+
    cursor->done = FALSE;
    cursor->end_of_event = FALSE;
    cursor->sent = TRUE;
@@ -382,53 +401,57 @@ _mongoc_cursor_get_more (mongoc_cursor_t *cursor)
 
    ENTRY;
 
-   bson_return_val_if_fail(cursor, FALSE);
+   BSON_ASSERT(cursor);
 
-   if (!_mongoc_client_warm_up (cursor->client, &cursor->error)) {
-      cursor->failed = TRUE;
-      RETURN (FALSE);
-   }
+   if (! cursor->in_exhaust) {
+      if (!_mongoc_client_warm_up (cursor->client, &cursor->error)) {
+         cursor->failed = TRUE;
+         RETURN (FALSE);
+      }
 
-   if (!(cursor_id = cursor->rpc.reply.cursor_id)) {
-      bson_set_error(&cursor->error,
-                     MONGOC_ERROR_CURSOR,
-                     MONGOC_ERROR_CURSOR_INVALID_CURSOR,
-                     "No valid cursor was provided.");
-      goto failure;
-   }
+      if (!(cursor_id = cursor->rpc.reply.cursor_id)) {
+         bson_set_error(&cursor->error,
+                        MONGOC_ERROR_CURSOR,
+                        MONGOC_ERROR_CURSOR_INVALID_CURSOR,
+                        "No valid cursor was provided.");
+         goto failure;
+      }
+
+      rpc.get_more.msg_len = 0;
+      rpc.get_more.request_id = 0;
+      rpc.get_more.response_to = 0;
+      rpc.get_more.opcode = MONGOC_OPCODE_GET_MORE;
+      rpc.get_more.zero = 0;
+      rpc.get_more.collection = cursor->ns;
+      if ((cursor->flags & MONGOC_QUERY_TAILABLE_CURSOR)) {
+         rpc.get_more.n_return = 0;
+      } else {
+         /*
+          * TODO: We need to apply the limit to this so we don't
+          * overshoot our target.
+          */
+         rpc.get_more.n_return = cursor->batch_size;
+      }
+      rpc.get_more.cursor_id = cursor_id;
 
-   rpc.get_more.msg_len = 0;
-   rpc.get_more.request_id = 0;
-   rpc.get_more.response_to = 0;
-   rpc.get_more.opcode = MONGOC_OPCODE_GET_MORE;
-   rpc.get_more.zero = 0;
-   rpc.get_more.collection = cursor->ns;
-   if ((cursor->flags & MONGOC_QUERY_TAILABLE_CURSOR)) {
-      rpc.get_more.n_return = 0;
-   } else {
       /*
-       * TODO: We need to apply the limit to this so we don't
-       * overshoot our target.
+       * TODO: Stamp protections for disconnections.
        */
-      rpc.get_more.n_return = cursor->batch_size;
-   }
-   rpc.get_more.cursor_id = cursor_id;
 
-   /*
-    * TODO: Stamp protections for disconnections.
-    */
+      if (!_mongoc_client_sendv(cursor->client, &rpc, 1, cursor->hint,
+                                NULL, cursor->read_prefs, &cursor->error)) {
+         cursor->done = TRUE;
+         cursor->failed = TRUE;
+         RETURN(FALSE);
+      }
 
-   if (!_mongoc_client_sendv(cursor->client, &rpc, 1, cursor->hint,
-                             NULL, cursor->read_prefs, &cursor->error)) {
-      cursor->done = TRUE;
-      cursor->failed = TRUE;
-      RETURN(FALSE);
+      request_id = BSON_UINT32_FROM_LE(rpc.header.request_id);
+   } else {
+      request_id = BSON_UINT32_FROM_LE(cursor->rpc.header.request_id);
    }
 
    _mongoc_buffer_clear(&cursor->buffer, FALSE);
 
-   request_id = BSON_UINT32_FROM_LE(rpc.header.request_id);
-
    if (!_mongoc_client_recv(cursor->client,
                             &cursor->rpc,
                             &cursor->buffer,
@@ -536,7 +559,17 @@ _mongoc_cursor_next (mongoc_cursor_t  *cursor,
 
    ENTRY;
 
-   bson_return_val_if_fail(cursor, FALSE);
+   BSON_ASSERT(cursor);
+
+
+   if (cursor->client->in_exhaust && ! cursor->in_exhaust) {
+      bson_set_error(&cursor->error,
+                     MONGOC_ERROR_CLIENT,
+                     MONGOC_ERROR_CLIENT_IN_EXHAUST,
+                     "Another cursor derived from this client is in exhaust.");
+      cursor->failed = TRUE;
+      RETURN(FALSE);
+   }
 
    if (bson) {
       *bson = NULL;
@@ -568,9 +601,11 @@ _mongoc_cursor_next (mongoc_cursor_t  *cursor,
    eof = FALSE;
    b = bson_reader_read(cursor->reader, &eof);
    cursor->end_of_event = eof;
-   cursor->done = (cursor->end_of_event &&
-                   !b &&
-                   !(cursor->flags & MONGOC_QUERY_TAILABLE_CURSOR));
+
+   cursor->done = cursor->end_of_event && (
+      (cursor->in_exhaust && !cursor->rpc.reply.cursor_id) ||
+      (!b && !(cursor->flags & MONGOC_QUERY_TAILABLE_CURSOR))
+      );
 
    /*
     * Do a supplimental check to see if we had a corrupted reply in the
diff --git a/mongoc/mongoc-error.h b/mongoc/mongoc-error.h
@@ -51,6 +51,7 @@ typedef enum
    MONGOC_ERROR_CLIENT_GETNONCE,
    MONGOC_ERROR_CLIENT_AUTHENTICATE,
    MONGOC_ERROR_CLIENT_NO_ACCEPTABLE_PEER,
+   MONGOC_ERROR_CLIENT_IN_EXHAUST,
 
    MONGOC_ERROR_PROTOCOL_INVALID_REPLY,
    MONGOC_ERROR_PROTOCOL_BAD_WIRE_VERSION,
diff --git a/mongoc/mongoc-gridfs-file.c b/mongoc/mongoc-gridfs-file.c
@@ -349,7 +349,7 @@ ssize_t
 mongoc_gridfs_file_readv (mongoc_gridfs_file_t *file,
                           struct iovec         *iov,
                           size_t                iovcnt,
-                          ssize_t               min_bytes,
+                          size_t                min_bytes,
                           bson_uint32_t         timeout_msec)
 {
    bson_uint32_t bytes_read = 0;
@@ -389,7 +389,7 @@ mongoc_gridfs_file_readv (mongoc_gridfs_file_t *file,
          } else if (file->length == file->pos) {
             /* we're at the end of the file.  So we're done */
             RETURN (bytes_read);
-         } else if (min_bytes > -1 && bytes_read >= min_bytes) {
+         } else if (bytes_read >= min_bytes) {
             /* we need a new page, but we've read enough bytes to stop */
             RETURN (bytes_read);
          } else {
diff --git a/mongoc/mongoc-gridfs-file.h b/mongoc/mongoc-gridfs-file.h
@@ -75,7 +75,7 @@ ssize_t
 mongoc_gridfs_file_readv (mongoc_gridfs_file_t *file,
                           struct iovec         *iov,
                           size_t                iovcnt,
-                          ssize_t               min_bytes,
+                          size_t                min_bytes,
                           bson_uint32_t         timeout_msec);
 
 
diff --git a/mongoc/mongoc-gridfs.c b/mongoc/mongoc-gridfs.c
@@ -239,7 +239,7 @@ mongoc_gridfs_create_file_from_stream (mongoc_gridfs_t          *gridfs,
 
    for (;; ) {
       r = mongoc_stream_read (stream, iov.iov_base, MONGOC_GRIDFS_STREAM_CHUNK,
-                              -1, 0);
+                              0, 0);
 
       if (r > 0) {
          iov.iov_len = r;
diff --git a/tests/test-mongoc-client.c b/tests/test-mongoc-client.c
diff --git a/tests/test-mongoc-gridfs.c b/tests/test-mongoc-gridfs.c
