fix small error and memory leaks in SORT

[redis.git] / zmalloc.c

diff --git a/zmalloc.c b/zmalloc.c
index 922856590e38e306a618d4b54cf7e5a22324251f..b1e96fc3c42149f367ed6cea6993dad93fb019f3 100644 (file)
--- a/zmalloc.c
+++ b/zmalloc.c
@@ -1,6 +1,6 @@
 /* zmalloc - total amount of allocated memory aware version of malloc()
  *
- * Copyright (c) 2006-2009, Salvatore Sanfilippo <antirez at gmail dot com>
+ * Copyright (c) 2009-2010, Salvatore Sanfilippo <antirez at gmail dot com>
  * All rights reserved.
  *
  * Redistribution and use in source and binary forms, with or without
@@ -28,45 +28,108 @@
  * POSSIBILITY OF SUCH DAMAGE.
  */
 
+#include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <pthread.h>
+#include "config.h"
+
+#if defined(__sun)
+#define PREFIX_SIZE sizeof(long long)
+#else
+#define PREFIX_SIZE sizeof(size_t)
+#endif
+
+#define increment_used_memory(_n) do { \
+    if (zmalloc_thread_safe) { \
+        pthread_mutex_lock(&used_memory_mutex);  \
+        used_memory += _n; \
+        pthread_mutex_unlock(&used_memory_mutex); \
+    } else { \
+        used_memory += _n; \
+    } \
+} while(0)
+
+#define decrement_used_memory(_n) do { \
+    if (zmalloc_thread_safe) { \
+        pthread_mutex_lock(&used_memory_mutex);  \
+        used_memory -= _n; \
+        pthread_mutex_unlock(&used_memory_mutex); \
+    } else { \
+        used_memory -= _n; \
+    } \
+} while(0)
 
 static size_t used_memory = 0;
+static int zmalloc_thread_safe = 0;
+pthread_mutex_t used_memory_mutex = PTHREAD_MUTEX_INITIALIZER;
+
+static void zmalloc_oom(size_t size) {
+    fprintf(stderr, "zmalloc: Out of memory trying to allocate %zu bytes\n",
+        size);
+    fflush(stderr);
+    abort();
+}
 
 void *zmalloc(size_t size) {
-    void *ptr = malloc(size+sizeof(size_t));
+    void *ptr = malloc(size+PREFIX_SIZE);
 
+    if (!ptr) zmalloc_oom(size);
+#ifdef HAVE_MALLOC_SIZE
+    increment_used_memory(redis_malloc_size(ptr));
+    return ptr;
+#else
     *((size_t*)ptr) = size;
-    used_memory += size+sizeof(size_t);
-    return (char*)ptr+sizeof(size_t);
+    increment_used_memory(size+PREFIX_SIZE);
+    return (char*)ptr+PREFIX_SIZE;
+#endif
 }
 
 void *zrealloc(void *ptr, size_t size) {
+#ifndef HAVE_MALLOC_SIZE
     void *realptr;
+#endif
     size_t oldsize;
     void *newptr;
 
     if (ptr == NULL) return zmalloc(size);
-    realptr = (char*)ptr-sizeof(size_t);
+#ifdef HAVE_MALLOC_SIZE
+    oldsize = redis_malloc_size(ptr);
+    newptr = realloc(ptr,size);
+    if (!newptr) zmalloc_oom(size);
+
+    decrement_used_memory(oldsize);
+    increment_used_memory(redis_malloc_size(newptr));
+    return newptr;
+#else
+    realptr = (char*)ptr-PREFIX_SIZE;
     oldsize = *((size_t*)realptr);
-    newptr = realloc(realptr,size+sizeof(size_t));
-    if (!newptr) return NULL;
+    newptr = realloc(realptr,size+PREFIX_SIZE);
+    if (!newptr) zmalloc_oom(size);
 
     *((size_t*)newptr) = size;
-    used_memory -= oldsize;
-    used_memory += size;
-    return (char*)newptr+sizeof(size_t);
+    decrement_used_memory(oldsize);
+    increment_used_memory(size);
+    return (char*)newptr+PREFIX_SIZE;
+#endif
 }
 
 void zfree(void *ptr) {
+#ifndef HAVE_MALLOC_SIZE
     void *realptr;
     size_t oldsize;
+#endif
 
     if (ptr == NULL) return;
-    realptr = (char*)ptr-sizeof(size_t);
+#ifdef HAVE_MALLOC_SIZE
+    decrement_used_memory(redis_malloc_size(ptr));
+    free(ptr);
+#else
+    realptr = (char*)ptr-PREFIX_SIZE;
     oldsize = *((size_t*)realptr);
-    used_memory -= oldsize+sizeof(size_t);
+    decrement_used_memory(oldsize+PREFIX_SIZE);
     free(realptr);
+#endif
 }
 
 char *zstrdup(const char *s) {
@@ -78,5 +141,14 @@ char *zstrdup(const char *s) {
 }
 
 size_t zmalloc_used_memory(void) {
-    return used_memory;
+    size_t um;
+
+    if (zmalloc_thread_safe) pthread_mutex_lock(&used_memory_mutex);
+    um = used_memory;
+    if (zmalloc_thread_safe) pthread_mutex_unlock(&used_memory_mutex);
+    return um;
+}
+
+void zmalloc_enable_thread_safeness(void) {
+    zmalloc_thread_safe = 1;
 }