Add monitor feature

src/Makefile.am

@@ -53,6 +53,7 @@ nutcracker_SOURCES =			\
 	nc_array.c nc_array.h		\
 	nc_util.c nc_util.h		\
 	nc_queue.h			\
+	nc_monitor.c nc_monitor.h    \
 	nc.c
 
 nutcracker_LDADD = $(top_builddir)/src/hashkit/libhashkit.a
@@ -81,7 +82,8 @@ test_all_SOURCES = test_all.c \
 	nc_string.c nc_string.h		\
 	nc_array.c nc_array.h		\
 	nc_util.c nc_util.h		\
-	nc_queue.h
+	nc_queue.h    \
+	nc_monitor.c nc_monitor.h
 
 test_all_LDADD = $(top_builddir)/src/hashkit/libhashkit.a
 test_all_LDADD += $(top_builddir)/src/proto/libproto.a

src/nc_client.c

@@ -18,6 +18,7 @@
 #include <nc_core.h>
 #include <nc_server.h>
 #include <nc_client.h>
+#include <nc_monitor.h>
 
 void
 client_ref(struct conn *conn, void *owner)
@@ -129,6 +130,11 @@ client_close(struct context *ctx, struct conn *conn)
 
     client_close_stats(ctx, conn->owner, conn->err, conn->eof);
 
+    /* when client close, if conn in monitor, delete it */
+    if (conn->monitor_client) {
+        del_from_monitor(conn);
+    }
+
     if (conn->sd < 0) {
         conn->unref(conn);
         conn_put(conn);

src/nc_connection.c

@@ -157,6 +157,7 @@ _conn_get(void)
     conn->done = 0;
     conn->redis = 0;
     conn->authenticated = 0;
+    conn->monitor_client = 0;
 
     ntotal_conn++;
     ncurr_conn++;

src/nc_connection.h

@@ -89,6 +89,7 @@ struct conn {
     unsigned            done:1;          /* done? aka close? */
     unsigned            redis:1;         /* redis? */
     unsigned            authenticated:1; /* authenticated? */
+    unsigned            monitor_client:1;/* monitor client? */
 };
 
 TAILQ_HEAD(conn_tqh, conn);

src/nc_core.c

@@ -21,6 +21,7 @@
 #include <nc_conf.h>
 #include <nc_server.h>
 #include <nc_proxy.h>
+#include <nc_monitor.h>
 
 static uint32_t ctx_id; /* context generation */
 
@@ -163,13 +164,15 @@ core_start(struct instance *nci)
     mbuf_init(nci);
     msg_init();
     conn_init();
+    monitor_init();
 
     ctx = core_ctx_create(nci);
     if (ctx != NULL) {
         nci->ctx = ctx;
         return ctx;
     }
 
+    monitor_deinit(ctx);
     conn_deinit();
     msg_deinit();
     mbuf_deinit();
@@ -180,6 +183,7 @@ core_start(struct instance *nci)
 void
 core_stop(struct context *ctx)
 {
+    monitor_deinit(ctx);
     conn_deinit();
     msg_deinit();
     mbuf_deinit();

src/nc_message.c

@@ -273,6 +273,7 @@ _msg_get(void)
     msg->fdone = 0;
     msg->swallow = 0;
     msg->redis = 0;
+    msg->monitor = 0;
 
     return msg;
 }
@@ -910,3 +911,28 @@ bool msg_set_placeholder_key(struct msg *r)
     return true;
 }
 
+rstatus_t
+msg_append_full(struct msg *msg, uint8_t *pos, size_t n)
+{
+    struct mbuf *mbuf = NULL;
+    size_t cidx = 0;
+    size_t mbsize = 0;
+    size_t clen = 0;
+
+    do {
+        mbuf = msg_ensure_mbuf(msg, n);
+        if (mbuf == NULL) {
+            return NC_ENOMEM;
+        }
+
+        mbsize = mbuf_size(mbuf);
+
+        clen = n > mbsize ? mbsize : n;
+        mbuf_copy(mbuf, pos+cidx, clen);
+        cidx += clen;
+        msg->mlen += (uint32_t)clen;
+        n -= clen;
+    } while(n);
+
+    return NC_OK;
+}

src/nc_message.h

@@ -202,6 +202,7 @@ typedef enum msg_parse_result {
     ACTION( REQ_REDIS_SELECT)                  /* only during init */                               \
     ACTION( REQ_REDIS_COMMAND)                 /* Sent to random server for redis-cli completions*/ \
     ACTION( REQ_REDIS_LOLWUT)                  /* Vitally important */                              \
+    ACTION( REQ_REDIS_MONITOR)                 /* monitor */                                        \
     ACTION( RSP_REDIS_STATUS )                 /* redis response */                                 \
     ACTION( RSP_REDIS_ERROR )                                                                       \
     ACTION( RSP_REDIS_ERROR_ERR )                                                                   \
@@ -300,6 +301,7 @@ struct msg {
     unsigned             fdone:1;         /* all fragments are done? */
     unsigned             swallow:1;       /* swallow response? */
     unsigned             redis:1;         /* redis? */
+    unsigned             monitor:1;       /* monitor comamnd? */
 };
 
 TAILQ_HEAD(msg_tqh, msg);
@@ -350,4 +352,6 @@ void rsp_recv_done(struct context *ctx, struct conn *conn, struct msg *msg, stru
 struct msg *rsp_send_next(struct context *ctx, struct conn *conn);
 void rsp_send_done(struct context *ctx, struct conn *conn, struct msg *msg);
 
+rstatus_t msg_append_full(struct msg *msg, uint8_t *pos, size_t n);
+
 #endif

src/nc_monitor.c

@@ -0,0 +1,149 @@
+/*
+ * twemproxy - A fast and lightweight proxy for memcached protocol.
+ * 
+ * Copyright (C) 2021, wei huang <wei.kukey@gmail.com>
+ * All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <nc_monitor.h>
+#include <nc_rbtree.h>
+
+struct rbtree monitor_tree;
+struct rbnode monitor_sentinel_node;
+
+void monitor_init()
+{
+    rbtree_init(&monitor_tree, &monitor_sentinel_node);
+}
+
+void monitor_deinit(struct context *ctx)
+{
+    struct rbnode *node = NULL;
+
+    while ((node = rbtree_min(&monitor_tree) != NULL))
+    {
+        struct conn *c = node->data;
+        if (event_del_conn(ctx->evb, c) != NC_OK)  {
+            log_warn("event del conn c %d failed, ignored: %s",
+                 c->sd, strerror(errno));
+        }
+        c->close(ctx, c);
+
+        rbtree_delete(&monitor_tree, node);
+        nc_free(node);
+    }
+}
+
+int mointor_is_empty()
+{
+    return rbtree_is_empty(&monitor_tree);
+}
+
+rstatus_t add_to_monitor(struct conn *c)
+{
+    ASSERT(c->client);
+
+    struct rbnode *node = nc_alloc(sizeof(struct rbnode));
+    if (node == NULL)
+    {
+        return NC_ENOMEM;
+    }
+
+    c->monitor_client = 1;
+    node->key = c->sd;
+    node->data = c;
+
+    rbtree_insert(&monitor_tree, node);
+}
+
+void del_from_monitor(struct conn *c)
+{
+    ASSERT(c->client && c->monitor_client);
+
+    struct rbnode *node = rbtree_find(&monitor_tree, c->sd);
+    ASSERT(node != NULL);
+    rbtree_delete(&monitor_tree, node);
+    nc_free(node);
+}
+
+struct monitor_data
+{
+    struct msg *m;
+    struct conn *c;
+    struct context *ctx;
+    struct string *d;
+};
+
+static void monitor_callback(struct rbnode *node, void *data)
+{
+    struct monitor_data *mdata = data;
+    struct conn *req_c = node->data;
+
+    struct msg *req = req_get(req_c);
+    if (req == NULL) {
+        return;
+    }
+    struct msg *rsp = msg_get(req_c, 0, mdata->c->redis);
+    if (rsp == NULL) {
+        msg_put(req);
+        return;
+    }
+
+    req->peer = rsp;
+    rsp->peer = req;
+
+    req->done = 1;
+    rsp->done = 1;
+
+    if (msg_append_full(rsp, mdata->d->data, mdata->d->len) != NC_OK) {
+        msg_put(req);
+        msg_put(rsp);
+        return;
+    }
+    req_c->enqueue_outq(mdata->ctx, req_c, req);
+    if (event_add_out(mdata->ctx->evb, req_c) != NC_OK) {
+        req_c->err = errno;
+        req_c->dequeue_outq(mdata->ctx, req_c, req);
+        msg_put(req);
+        msg_put(rsp);
+    }
+
+    return;
+}
+
+rstatus_t make_monitor(struct context *ctx, struct conn *c, struct msg *m)
+{
+    ASSERT(c->client);
+
+    struct string monitor_message = null_string;
+    struct monitor_data mdata = {0};
+    mdata.m = m;
+    mdata.c = c;
+    mdata.d = &monitor_message;
+    mdata.ctx = ctx;
+    struct keypos *kpos = array_get(m->keys, 0);
+
+    string_printf(&monitor_message, "+%ld.%06ld [%s] command=%s key0=",
+                    m->start_ts/1000000, m->start_ts%1000000, 
+                    nc_unresolve_peer_desc(c->sd),
+                    (msg_type_string(m->type))->data);
+    string_cat_len(&monitor_message, kpos->start, kpos->end - kpos->start);
+    string_cat_len(&monitor_message, "\r\n", 2);
+
+    rbtree_inorder_traversal(monitor_tree.root, monitor_tree.sentinel, monitor_callback, &mdata);
+
+    string_deinit(&monitor_message);
+    return NC_OK;
+}

src/nc_monitor.h

@@ -0,0 +1,33 @@
+/*
+ * twemproxy - A fast and lightweight proxy for memcached protocol.
+ * 
+ * Copyright (C) 2021, wei huang <wei.kukey@gmail.com>
+ * All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#ifndef NC_MONITOR_H
+#define NC_MONITOR_H
+
+#include <nc_core.h>
+
+void monitor_init();
+void monitor_deinit(struct context *ctx);
+int mointor_is_empty();
+
+rstatus_t add_to_monitor(struct conn *c);
+void del_from_monitor(struct conn *c);
+rstatus_t make_monitor(struct context *ctx, struct conn *c, struct msg *m);
+
+#endif

src/nc_rbtree.c

@@ -342,3 +342,44 @@ rbtree_delete(struct rbtree *tree, struct rbnode *node)
 
     rbtree_black(temp);
 }
+
+struct rbnode *rbtree_find(struct rbtree *tree, int64_t key) {
+    struct rbnode **root = &tree->root;
+    struct rbnode *sentinel = tree->sentinel;
+    struct rbnode *temp, **p;
+
+    /* empty tree */
+
+    if (*root == sentinel) {
+        return NULL;
+    }
+
+    /* a binary tree find */
+    temp = *root;
+    for (;;) {
+
+        if (temp->key == key) {
+            break;
+        }
+
+        p = (key < temp->key) ? &temp->left : &temp->right;
+        if (*p == sentinel) {
+            return NULL;
+        }
+        temp = *p;
+    }
+
+    return temp;
+}
+
+void rbtree_inorder_traversal(struct rbnode *root, struct rbnode *sentinel, 
+    void (*func)(struct rbnode *, void *), void *data) {
+
+    if (root == NULL || root == sentinel) {
+        return;
+    }
+
+    func(root, data);
+    rbtree_inorder_traversal(root->left, sentinel, func, data);
+    rbtree_inorder_traversal(root->right, sentinel, func, data);
+}

src/nc_rbtree.h

@@ -23,6 +23,7 @@
 #define rbtree_is_red(_node)        ((_node)->color)
 #define rbtree_is_black(_node)      (!rbtree_is_red(_node))
 #define rbtree_copy_color(_n1, _n2) ((_n1)->color = (_n2)->color)
+#define rbtree_is_empty(_tree)      ((_tree)->root == (_tree)->sentinel)
 
 struct rbnode {
     struct rbnode *left;     /* left link */
@@ -43,5 +44,8 @@ void rbtree_init(struct rbtree *tree, struct rbnode *node);
 struct rbnode *rbtree_min(const struct rbtree *tree);
 void rbtree_insert(struct rbtree *tree, struct rbnode *node);
 void rbtree_delete(struct rbtree *tree, struct rbnode *node);
+struct rbnode *rbtree_find(struct rbtree *tree, int64_t key);
+void rbtree_inorder_traversal(struct rbnode *root, struct rbnode *sentinel, 
+                                void (*func)(struct rbnode *, void *), void *data);
 
 #endif