CCBC-1646: Server groups for replica reads

Change-Id: Ia3aa07ce4bb8c0fc81aecd38b999fe16a0574fb8
Reviewed-on: https://review.couchbase.org/c/libcouchbase/+/216990
Tested-by: Build Bot <build@couchbase.com>
Reviewed-by: Jared Casey <jared.casey@couchbase.com>

Author: avsej

include/libcouchbase/cntl.h

@@ -1326,11 +1326,27 @@ typedef enum {
  */
 #define LCB_CNTL_ENABLE_OP_METRICS 0x67
 
+/**
+ * Select server group to use for replica APIs.
+ *
+ * For some use-cases it might be necessary to restrict list of the nodes,
+ * that are used in replica read APIs to single server group to optimize
+ * network costs.
+ *
+ * Use `preferred_server_group` in the connection string.
+ *
+ * @committed
+ * @cntl_arg_both{const char*}
+ *
+ * @see https://docs.couchbase.com/server/current/manage/manage-groups/manage-groups.html
+ */
+#define LCB_CNTL_PREFERRED_SERVER_GROUP 0x68
+
 /**
  * This is not a command, but rather an indicator of the last item.
  * @internal
  */
-#define LCB_CNTL__MAX 0x68
+#define LCB_CNTL__MAX 0x69
 /**@}*/
 
 #ifdef __cplusplus

include/libcouchbase/couchbase.h

@@ -793,6 +793,11 @@ typedef enum {
     LCB_REPLICA_MODE__MAX
 } lcb_REPLICA_MODE;
 
+typedef enum {
+    LCB_REPLICA_READ_PREFERENCE_NONE = 0x00,
+    LCB_REPLICA_READ_PREFERENCE_SELECTED_SERVER_GROUP = 0x01,
+} lcb_REPLICA_READ_PREFERENCE;
+
 typedef struct lcb_RESPGETREPLICA_ lcb_RESPGETREPLICA;
 
 LIBCOUCHBASE_API lcb_STATUS lcb_respgetreplica_status(const lcb_RESPGETREPLICA *resp);
@@ -817,6 +822,8 @@ LIBCOUCHBASE_API lcb_STATUS lcb_cmdgetreplica_collection(lcb_CMDGETREPLICA *cmd,
                                                          const char *collection, size_t collection_len);
 LIBCOUCHBASE_API lcb_STATUS lcb_cmdgetreplica_key(lcb_CMDGETREPLICA *cmd, const char *key, size_t key_len);
 LIBCOUCHBASE_API lcb_STATUS lcb_cmdgetreplica_timeout(lcb_CMDGETREPLICA *cmd, uint32_t timeout);
+LIBCOUCHBASE_API lcb_STATUS lcb_cmdgetreplica_read_preference(lcb_CMDGETREPLICA *cmd,
+                                                              lcb_REPLICA_READ_PREFERENCE preference);
 /**
  * @internal Internal: This should never be used and is not supported.
  */

include/libcouchbase/vbucket.h

@@ -110,6 +110,7 @@ typedef struct {
     char *alt_hostname;         /**< selected alternative hostname for the node */
     lcbvb_SERVICES alt_svc;     /**< selected alternative plain services */
     lcbvb_SERVICES alt_svc_ssl; /**< selected alternative SSL Services */
+    char *server_group;         /**< the group name or NULL */
 } lcbvb_SERVER;
 
 /**@volatile. ABI/API compatibility not guaranteed between versions */

src/capi/cmd_get_replica.hh

@@ -75,6 +75,17 @@ struct lcb_CMDGETREPLICA_ {
         return LCB_SUCCESS;
     }
 
+    lcb_STATUS read_preference(lcb_REPLICA_READ_PREFERENCE preference)
+    {
+        read_preference_ = preference;
+        return LCB_SUCCESS;
+    }
+
+    lcb_REPLICA_READ_PREFERENCE read_preference() const
+    {
+        return read_preference_;
+    }
+
     lcb_STATUS collection(lcb::collection_qualifier collection)
     {
         collection_ = std::move(collection);
@@ -194,6 +205,7 @@ struct lcb_CMDGETREPLICA_ {
     int select_index_{0};
     std::string impostor_{};
     std::vector<std::string> extra_privileges_{};
+    lcb_REPLICA_READ_PREFERENCE read_preference_{LCB_REPLICA_READ_PREFERENCE_NONE};
 };
 
 struct lcb_RESPGETREPLICA_ {

src/cntl.cc

@@ -747,6 +747,22 @@ HANDLER(network_handler)
     return LCB_SUCCESS;
 }
 
+HANDLER(preferred_server_group_handler)
+{
+    if (mode == LCB_CNTL_SET) {
+        const char *val = reinterpret_cast<const char *>(arg);
+        free(LCBT_SETTING(instance, preferred_server_group));
+        LCBT_SETTING(instance, network) = nullptr;
+        if (val) {
+            LCBT_SETTING(instance, preferred_server_group) = lcb_strdup(val);
+        }
+    } else {
+        *(const char **)arg = LCBT_SETTING(instance, preferred_server_group);
+    }
+    (void)cmd;
+    return LCB_SUCCESS;
+}
+
 HANDLER(durable_write_handler){RETURN_GET_SET(int, LCBT_SETTING(instance, enable_durable_write))}
 
 HANDLER(unordered_execution_handler)
@@ -860,6 +876,7 @@ static ctl_handler handlers[] = {
     enable_errmap_handler,                /* LCB_CNTL_ENABLE_ERRMAP */
     timeout_common,                       /* LCB_CNTL_OP_METRICS_FLUSH_INTERVAL */
     enable_op_metrics_handler,            /* LCB_CNTL_ENABLE_OP_METRICS */
+    preferred_server_group_handler,       /* LCB_CNTL_PREFERRED_SERVER_GROUP */
     nullptr
 };
 /* clang-format on */
@@ -1101,6 +1118,7 @@ static cntl_OPCODESTRS stropcode_map[] = {
     {"enable_errmap", LCB_CNTL_ENABLE_ERRMAP, convert_intbool},
     {"operation_metrics_flush_interval", LCB_CNTL_OP_METRICS_FLUSH_INTERVAL, convert_timevalue},
     {"enable_operation_metrics", LCB_CNTL_ENABLE_OP_METRICS, convert_intbool},
+    {"preferred_server_group", LCB_CNTL_PREFERRED_SERVER_GROUP, convert_passthru},
     {nullptr, -1}};
 
 #define CNTL_NUM_HANDLERS (sizeof(handlers) / sizeof(handlers[0]))

src/operations/get_replica.cc

@@ -24,6 +24,9 @@
 #include "capi/cmd_get.hh"
 #include "capi/cmd_get_replica.hh"
 
+#include <random>
+#include <algorithm>
+
 LIBCOUCHBASE_API lcb_STATUS lcb_respgetreplica_status(const lcb_RESPGETREPLICA *resp)
 {
     return resp->ctx.rc;
@@ -139,6 +142,12 @@ LIBCOUCHBASE_API lcb_STATUS lcb_cmdgetreplica_key(lcb_CMDGETREPLICA *cmd, const
     return cmd->key(std::string(key, key_len));
 }
 
+LIBCOUCHBASE_API lcb_STATUS lcb_cmdgetreplica_read_preference(lcb_CMDGETREPLICA *cmd,
+                                                              lcb_REPLICA_READ_PREFERENCE preference)
+{
+    return cmd->read_preference(preference);
+}
+
 LIBCOUCHBASE_API lcb_STATUS lcb_cmdgetreplica_on_behalf_of(lcb_CMDGETREPLICA *cmd, const char *data, size_t data_len)
 {
     return cmd->on_behalf_of(std::string(data, data_len));
@@ -243,6 +252,90 @@ RGetCookie::RGetCookie(void *cookie_, lcb_INSTANCE *instance_, get_replica_mode
 {
 }
 
+struct readable_node {
+    bool is_replica;
+    std::size_t node_index;    // index of the node in configuration
+    std::size_t replica_index; // zero-based index of the replica
+};
+
+template <typename COMMAND>
+static std::vector<readable_node> select_effective_node_indexes(lcb_INSTANCE *instance, COMMAND cmd)
+{
+    mc_CMDQUEUE *cq = &instance->cmdq;
+    int vbid;
+    int active_index;
+    lcb_KEYBUF keybuf{LCB_KV_COPY, {cmd->key().c_str(), cmd->key().size()}};
+    mcreq_map_key(cq, &keybuf, MCREQ_PKT_BASESIZE, &vbid, &active_index);
+
+    if (active_index < 0) {
+        return {};
+    }
+
+    std::vector<readable_node> effective_nodes{};
+
+    if (cmd->mode() == get_replica_mode::select) {
+        int replica_index = cmd->selected_replica_index();
+        int node_index = lcbvb_vbreplica(cq->config, vbid, replica_index);
+        if (node_index >= 0) {
+            effective_nodes.push_back(
+                {true, static_cast<std::size_t>(node_index), static_cast<std::size_t>(replica_index)});
+        }
+        return effective_nodes;
+    }
+
+    std::string preferred_server_group{};
+    bool use_preferred_server_group{cmd->read_preference() == LCB_REPLICA_READ_PREFERENCE_SELECTED_SERVER_GROUP};
+
+    if (use_preferred_server_group) {
+        if (instance->settings->preferred_server_group == nullptr) {
+            return {};
+        } else {
+            preferred_server_group = instance->settings->preferred_server_group;
+        }
+    }
+
+    /* Make sure they're all online */
+    for (std::size_t ii = 0; ii < LCBT_NREPLICAS(instance); ii++) {
+        int node_index = lcbvb_vbreplica(cq->config, vbid, ii);
+        if (node_index < 0) {
+            if (cmd->mode() == get_replica_mode::all) {
+                return {};
+            }
+            continue;
+        }
+        if (use_preferred_server_group) {
+            const char *server_group = cq->config->servers[node_index].server_group;
+            if (server_group != nullptr && preferred_server_group == server_group) {
+                effective_nodes.push_back({true, static_cast<std::size_t>(node_index), ii});
+            }
+        } else {
+            effective_nodes.push_back({true, static_cast<std::size_t>(node_index), ii});
+        }
+    }
+
+    if (cmd->mode() == get_replica_mode::any) {
+        if (effective_nodes.empty()) {
+            return {};
+        }
+        static std::random_device rd;
+        std::mt19937 g(rd());
+        std::shuffle(effective_nodes.begin(), effective_nodes.end(), g);
+        return {effective_nodes.front()};
+    }
+
+    /* read from active */
+    if (use_preferred_server_group) {
+        const char *active_server_group = cq->config->servers[active_index].server_group;
+        if (active_server_group != nullptr && preferred_server_group == active_server_group) {
+            effective_nodes.push_back({false, static_cast<std::size_t>(active_index), 0xff});
+        }
+    } else {
+        effective_nodes.push_back({false, static_cast<std::size_t>(active_index), 0xff});
+    }
+
+    return effective_nodes;
+}
+
 static lcb_STATUS get_replica_validate(lcb_INSTANCE *instance, const lcb_CMDGETREPLICA *cmd)
 {
     if (cmd->key().empty()) {
@@ -283,6 +376,13 @@ static lcb_STATUS get_replica_validate(lcb_INSTANCE *instance, const lcb_CMDGETR
             break;
 
         case get_replica_mode::all:
+            if (cmd->read_preference() == LCB_REPLICA_READ_PREFERENCE_SELECTED_SERVER_GROUP) {
+                auto effective_nodes = select_effective_node_indexes(instance, cmd);
+                if (effective_nodes.empty()) {
+                    return LCB_ERR_DOCUMENT_UNRETRIEVABLE;
+                }
+            }
+
             r0 = 0;
             r1 = LCBT_NREPLICAS(instance);
             /* Make sure they're all online */
@@ -309,7 +409,6 @@ static lcb_STATUS get_replica_schedule(lcb_INSTANCE *instance, std::shared_ptr<l
      */
     mc_CMDQUEUE *cq = &instance->cmdq;
     int vbid, ixtmp;
-    protocol_binary_request_header req{};
     unsigned r0 = 0, r1 = 0;
 
     lcb_KEYBUF keybuf{LCB_KV_COPY, {cmd->key().c_str(), cmd->key().size()}};
@@ -355,6 +454,11 @@ static lcb_STATUS get_replica_schedule(lcb_INSTANCE *instance, std::shared_ptr<l
         return LCB_ERR_NO_MATCHING_SERVER;
     }
 
+    const auto effective_nodes = select_effective_node_indexes(instance, cmd);
+    if (effective_nodes.empty()) {
+        return LCB_ERR_DOCUMENT_UNRETRIEVABLE;
+    }
+
     std::vector<std::uint8_t> framing_extras;
     if (cmd->want_impersonation()) {
         lcb_STATUS err = lcb::flexible_framing_extras::encode_impersonate_user(cmd->impostor(), framing_extras);
@@ -376,68 +480,67 @@ static lcb_STATUS get_replica_schedule(lcb_INSTANCE *instance, std::shared_ptr<l
         rck->start + cmd->timeout_or_default_in_nanoseconds(LCB_US2NS(LCBT_SETTING(instance, operation_timeout)));
 
     /* Initialize the packet */
+    protocol_binary_request_header req{};
     req.request.magic = framing_extras.empty() ? PROTOCOL_BINARY_REQ : PROTOCOL_BINARY_AREQ;
-    req.request.opcode = PROTOCOL_BINARY_CMD_GET_REPLICA;
     req.request.datatype = PROTOCOL_BINARY_RAW_BYTES;
     req.request.vbucket = htons(static_cast<std::uint16_t>(vbid));
     req.request.cas = 0;
     req.request.extlen = 0;
 
     auto ffextlen = static_cast<std::uint8_t>(framing_extras.size());
 
-    rck->r_cur = r0;
-    do {
-        int curix;
-        mc_PIPELINE *pl;
-        mc_PACKET *pkt;
-
-        curix = lcbvb_vbreplica(cq->config, vbid, r0);
-        /* XXX: this is always expected to be in range. For the FIRST mode
-         * it will seek to the first valid index (checked above), and for the
-         * ALL mode, it will fail if not all replicas are already online
-         * (also checked above) */
-        pl = cq->pipelines[curix];
-        pkt = mcreq_allocate_packet(pl);
-        if (!pkt) {
-            delete rck;
-            return LCB_ERR_NO_MEMORY;
-        }
+    for (auto node : effective_nodes) {
+        if (node.is_replica) {
+            req.request.opcode = PROTOCOL_BINARY_CMD_GET_REPLICA;
+            rck->r_cur = node.replica_index;
+            mc_PIPELINE *pl;
+            mc_PACKET *pkt;
+
+            /* XXX: this is always expected to be in range. For the FIRST mode
+             * it will seek to the first valid index (checked above), and for the
+             * ALL mode, it will fail if not all replicas are already online
+             * (also checked above) */
+            pl = cq->pipelines[node.node_index];
+            pkt = mcreq_allocate_packet(pl);
+            if (!pkt) {
+                delete rck;
+                return LCB_ERR_NO_MEMORY;
+            }
 
-        pkt->u_rdata.exdata = rck;
-        pkt->flags |= MCREQ_F_REQEXT;
-
-        mcreq_reserve_key(pl, pkt, sizeof(req.bytes) + ffextlen, &keybuf, cmd->collection().collection_id());
-        size_t nkey = pkt->kh_span.size - MCREQ_PKT_BASESIZE + pkt->extlen;
-        req.request.keylen = htons((uint16_t)nkey);
-        req.request.bodylen = htonl((uint32_t)nkey + framing_extras.size());
-        req.request.opaque = pkt->opaque;
-        rck->remaining++;
-        mcreq_write_hdr(pkt, &req);
-        if (!framing_extras.empty()) {
-            memcpy(SPAN_BUFFER(&pkt->kh_span) + sizeof(req.bytes), framing_extras.data(), framing_extras.size());
-        }
-        mcreq_sched_add(pl, pkt);
-    } while (++r0 < r1);
-
-    if (cmd->need_get_active()) {
-        req.request.opcode = PROTOCOL_BINARY_CMD_GET;
-        mc_PIPELINE *pl;
-        mc_PACKET *pkt;
-        lcb_STATUS err = mcreq_basic_packet(cq, &keybuf, cmd->collection().collection_id(), &req, 0, ffextlen, &pkt,
-                                            &pl, MCREQ_BASICPACKET_F_FALLBACKOK);
-        if (err != LCB_SUCCESS) {
-            delete rck;
-            return err;
-        }
-        req.request.opaque = pkt->opaque;
-        pkt->u_rdata.exdata = rck;
-        pkt->flags |= MCREQ_F_REQEXT;
-        rck->remaining++;
-        mcreq_write_hdr(pkt, &req);
-        if (!framing_extras.empty()) {
-            memcpy(SPAN_BUFFER(&pkt->kh_span) + sizeof(req.bytes), framing_extras.data(), framing_extras.size());
+            pkt->u_rdata.exdata = rck;
+            pkt->flags |= MCREQ_F_REQEXT;
+
+            mcreq_reserve_key(pl, pkt, sizeof(req.bytes) + ffextlen, &keybuf, cmd->collection().collection_id());
+            size_t nkey = pkt->kh_span.size - MCREQ_PKT_BASESIZE + pkt->extlen;
+            req.request.keylen = htons((uint16_t)nkey);
+            req.request.bodylen = htonl((uint32_t)nkey + framing_extras.size());
+            req.request.opaque = pkt->opaque;
+            rck->remaining++;
+            mcreq_write_hdr(pkt, &req);
+            if (!framing_extras.empty()) {
+                memcpy(SPAN_BUFFER(&pkt->kh_span) + sizeof(req.bytes), framing_extras.data(), framing_extras.size());
+            }
+            mcreq_sched_add(pl, pkt);
+        } else {
+            req.request.opcode = PROTOCOL_BINARY_CMD_GET;
+            mc_PIPELINE *pl;
+            mc_PACKET *pkt;
+            lcb_STATUS err = mcreq_basic_packet(cq, &keybuf, cmd->collection().collection_id(), &req, 0, ffextlen, &pkt,
+                                                &pl, MCREQ_BASICPACKET_F_FALLBACKOK);
+            if (err != LCB_SUCCESS) {
+                delete rck;
+                return err;
+            }
+            req.request.opaque = pkt->opaque;
+            pkt->u_rdata.exdata = rck;
+            pkt->flags |= MCREQ_F_REQEXT;
+            rck->remaining++;
+            mcreq_write_hdr(pkt, &req);
+            if (!framing_extras.empty()) {
+                memcpy(SPAN_BUFFER(&pkt->kh_span) + sizeof(req.bytes), framing_extras.data(), framing_extras.size());
+            }
+            mcreq_sched_add(pl, pkt);
         }
-        mcreq_sched_add(pl, pkt);
     }
 
     MAYBE_SCHEDLEAVE(instance)