From 442246dde2d8842caf4c1f4d4694ef5a78a0bad8 Mon Sep 17 00:00:00 2001
From: antirez <antirez@gmail.com>
Date: Tue, 7 Feb 2012 13:05:36 +0100
Subject: [PATCH] Precision of getClientOutputBufferMemoryUsage() greatily
 improved, see issue #327 for more information.

---
 src/networking.c | 25 +++++++++++++++++--------
 src/zmalloc.c    | 14 ++++++++++++++
 src/zmalloc.h    |  4 ++++
 3 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/src/networking.c b/src/networking.c
index e92f680c..99a381ce 100644
--- a/src/networking.c
+++ b/src/networking.c
@@ -151,7 +151,7 @@ void _addReplyObjectToList(redisClient *c, robj *o) {
             listAddNodeTail(c->reply,o);
         }
     }
-    c->reply_bytes += sdslen(o->ptr);
+    c->reply_bytes += zmalloc_size(o->ptr);
     asyncCloseClientOnOutputBufferLimitReached(c);
 }
 
@@ -165,7 +165,7 @@ void _addReplySdsToList(redisClient *c, sds s) {
         return;
     }
 
-    c->reply_bytes += sdslen(s);
+    c->reply_bytes += zmalloc_size(s);
     if (listLength(c->reply) == 0) {
         listAddNodeTail(c->reply,createObject(REDIS_STRING,s));
     } else {
@@ -191,7 +191,10 @@ void _addReplyStringToList(redisClient *c, char *s, size_t len) {
     if (c->flags & REDIS_CLOSE_AFTER_REPLY) return;
 
     if (listLength(c->reply) == 0) {
-        listAddNodeTail(c->reply,createStringObject(s,len));
+        robj *o = createStringObject(s,len);
+
+        listAddNodeTail(c->reply,o);
+        c->reply_bytes += zmalloc_size(o->ptr);
     } else {
         tail = listNodeValue(listLast(c->reply));
 
@@ -199,13 +202,17 @@ void _addReplyStringToList(redisClient *c, char *s, size_t len) {
         if (tail->ptr != NULL &&
             sdslen(tail->ptr)+len <= REDIS_REPLY_CHUNK_BYTES)
         {
+            c->reply_bytes -= zmalloc_size(tail->ptr);
             tail = dupLastObjectIfNeeded(c->reply);
             tail->ptr = sdscatlen(tail->ptr,s,len);
+            c->reply_bytes += zmalloc_size(tail->ptr);
         } else {
-            listAddNodeTail(c->reply,createStringObject(s,len));
+            robj *o = createStringObject(s,len);
+
+            listAddNodeTail(c->reply,o);
+            c->reply_bytes += zmalloc_size(o->ptr);
         }
     }
-    c->reply_bytes += len;
     asyncCloseClientOnOutputBufferLimitReached(c);
 }
 
@@ -336,7 +343,7 @@ void setDeferredMultiBulkLength(redisClient *c, void *node, long length) {
 
     len = listNodeValue(ln);
     len->ptr = sdscatprintf(sdsempty(),"*%ld\r\n",length);
-    c->reply_bytes += sdslen(len->ptr);
+    c->reply_bytes += zmalloc_size(len->ptr);
     if (ln->next != NULL) {
         next = listNodeValue(ln->next);
 
@@ -638,6 +645,7 @@ void freeClientsInAsyncFreeQueue(void) {
 void sendReplyToClient(aeEventLoop *el, int fd, void *privdata, int mask) {
     redisClient *c = privdata;
     int nwritten = 0, totwritten = 0, objlen;
+    size_t objmem;
     robj *o;
     REDIS_NOTUSED(el);
     REDIS_NOTUSED(mask);
@@ -663,6 +671,7 @@ void sendReplyToClient(aeEventLoop *el, int fd, void *privdata, int mask) {
         } else {
             o = listNodeValue(listFirst(c->reply));
             objlen = sdslen(o->ptr);
+            objmem = zmalloc_size(o->ptr);
 
             if (objlen == 0) {
                 listDelNode(c->reply,listFirst(c->reply));
@@ -683,7 +692,7 @@ void sendReplyToClient(aeEventLoop *el, int fd, void *privdata, int mask) {
             if (c->sentlen == objlen) {
                 listDelNode(c->reply,listFirst(c->reply));
                 c->sentlen = 0;
-                c->reply_bytes -= objlen;
+                c->reply_bytes -= objmem;
             }
         }
         /* Note that we avoid to send more than REDIS_MAX_WRITE_PER_EVENT
@@ -1222,7 +1231,7 @@ void rewriteClientCommandArgument(redisClient *c, int i, robj *newval) {
  * the caller wishes. The main usage of this function currently is
  * enforcing the client output length limits. */
 unsigned long getClientOutputBufferMemoryUsage(redisClient *c) {
-    unsigned long list_item_size = sizeof(listNode);
+    unsigned long list_item_size = sizeof(listNode)+sizeof(robj);
 
     return c->reply_bytes + (list_item_size*listLength(c->reply));
 }
diff --git a/src/zmalloc.c b/src/zmalloc.c
index 56b9140c..89f80d83 100644
--- a/src/zmalloc.c
+++ b/src/zmalloc.c
@@ -150,6 +150,20 @@ void *zrealloc(void *ptr, size_t size) {
 #endif
 }
 
+/* Provide zmalloc_size() for systems where this function is not provided by
+ * malloc itself, given that in that case we store an header with this
+ * information as the first bytes of every allocation. */
+#ifndef HAVE_MALLOC_SIZE
+size_t zmalloc_size(void *ptr) {
+    void *realptr = (char*)ptr-PREFIX_SIZE;
+    size_t size = *((size_t*)realptr);
+    /* Assume at least that all the allocations are padded at sizeof(long) by
+     * the underlying allocator. */
+    if (size&(sizeof(long)-1)) size += sizeof(long)-(size&(sizeof(long)-1));
+    return size+PREFIX_SIZE;
+}
+#endif
+
 void zfree(void *ptr) {
 #ifndef HAVE_MALLOC_SIZE
     void *realptr;
diff --git a/src/zmalloc.h b/src/zmalloc.h
index 7ee556a3..995814c8 100644
--- a/src/zmalloc.h
+++ b/src/zmalloc.h
@@ -76,4 +76,8 @@ void zmalloc_enable_thread_safeness(void);
 float zmalloc_get_fragmentation_ratio(void);
 size_t zmalloc_get_rss(void);
 
+#ifndef HAVE_MALLOC_SIZE
+size_t zmalloc_size(void *ptr);
+#endif
+
 #endif /* __ZMALLOC_H */
-- 
2.45.2