/* -*- mode: C; c-basic-offset: 4 -*- */
#ident "Copyright (c) 2007-2010 Tokutek Inc.  All rights reserved."

#include <toku_portability.h>
#include "memory.h"
#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <malloc.h>
#include "toku_assert.h"
#include "toku_pthread.h"

static malloc_fun_t  t_malloc  = 0;
static malloc_fun_t  t_xmalloc = 0;
static free_fun_t    t_free    = 0;
static realloc_fun_t t_realloc = 0;
static realloc_fun_t t_xrealloc = 0;

static MEMORY_STATUS_S status;
static toku_pthread_mutex_t memory_status_mutex = TOKU_PTHREAD_MUTEX_INITIALIZER;

static void 
memory_status_lock(void) {
    int r = toku_pthread_mutex_lock(&memory_status_mutex); assert_zero(r);
}

static void
memory_status_unlock(void) {
    int r = toku_pthread_mutex_unlock(&memory_status_mutex); assert_zero(r);
}

void 
toku_memory_get_status(MEMORY_STATUS s) {
    memory_status_lock();
    *s = status;
    memory_status_unlock();
}

// max_in_use may be slightly off because use of max_in_use is not thread-safe.
// It is not worth the overhead to make it completely accurate.
static inline void 
set_max(uint64_t sum_used, uint64_t sum_freed) {
    uint64_t in_use = (sum_used - sum_freed);
    if ((!(in_use & 0x8000000000000000)) // if wrap due to another thread, ignore bogus "negative" value
	&& (in_use > status.max_in_use)) {
	status.max_in_use = in_use;
    }
}

void *toku_malloc(size_t size) {
    void *p = t_malloc ? t_malloc(size) : os_malloc(size);
    if (p) {
        size_t used = malloc_usable_size(p);
        memory_status_lock();
        status.malloc_count += 1;
        status.requested += size;
        status.used += used;
        set_max(status.used, status.freed);
        memory_status_unlock();
    } else {
        memory_status_lock();
        status.malloc_fail += 1;
        memory_status_unlock();
    }
    return p;
}

void *
toku_calloc(size_t nmemb, size_t size) {
    size_t newsize = nmemb * size;
    void *p = toku_malloc(newsize);
    if (p) memset(p, 0, newsize);
    return p;
}

void *
toku_realloc(void *p, size_t size) {
    size_t used_orig = p ? malloc_usable_size(p) : 0;
    void *q = t_realloc ? t_realloc(p, size) : os_realloc(p, size);
    if (q) {
	size_t used = malloc_usable_size(q);
        memory_status_lock();
	status.realloc_count += 1L;
	status.requested += size;
	status.used += used;
	status.freed += used_orig;
	set_max(status.used, status.freed);
        memory_status_unlock();
    } else {
        memory_status_lock();
	status.realloc_fail += 1;
        memory_status_unlock();
    }
    return q;
}

void *
toku_memdup(const void *v, size_t len) {
    void *p = toku_malloc(len);
    if (p) memcpy(p, v,len);
    return p;
}

char *
toku_strdup(const char *s) {
    return toku_memdup(s, strlen(s)+1);
}

void
toku_free(void *p) {
    if (p) {
	size_t used = malloc_usable_size(p);
        memory_status_lock();
	status.free_count += 1L;
	status.freed += used;
        memory_status_unlock();
	if (t_free)
	    t_free(p);
	else
	    os_free(p);
    }
}

void
toku_free_n(void* p, size_t size __attribute__((unused))) {
    toku_free(p);
}

void *
toku_xmalloc(size_t size) {
    void *p = t_xmalloc ? t_xmalloc(size) : os_malloc(size);
    if (p == NULL)  // avoid function call in common case
        resource_assert(p);
    size_t used = malloc_usable_size(p);
    memory_status_lock();
    status.malloc_count += 1L;
    status.requested += size;
    status.used += used;
    set_max(status.used, status.freed);
    memory_status_unlock();
    return p;
}

void *
toku_xcalloc(size_t nmemb, size_t size) {
    size_t newsize = nmemb * size;
    void *vp = toku_xmalloc(newsize);
    if (vp) memset(vp, 0, newsize);
    return vp;
}

void *
toku_xrealloc(void *v, size_t size) {
    size_t used_orig = v ? malloc_usable_size(v) : 0;
    void *p = t_xrealloc ? t_xrealloc(v, size) : os_realloc(v, size);
    if (p == 0)  // avoid function call in common case
        resource_assert(p);
    size_t used = malloc_usable_size(p);
    memory_status_lock();
    status.realloc_count += 1L;
    status.requested += size;
    status.used += used;
    status.freed += used_orig;
    set_max(status.used, status.freed);
    memory_status_unlock();
    return p;
}

size_t 
toku_malloc_usable_size(void *p) {
    return malloc_usable_size(p);
}

void *
toku_xmemdup (const void *v, size_t len) {
    void *p = toku_xmalloc(len);
    memcpy(p, v, len);
    return p;
}

char *
toku_xstrdup (const char *s) {
    return toku_xmemdup(s, strlen(s)+1);
}

void
toku_set_func_malloc(malloc_fun_t f) {
    t_malloc = f;
    t_xmalloc = f;
}

void
toku_set_func_xmalloc_only(malloc_fun_t f) {
    t_xmalloc = f;
}

void
toku_set_func_malloc_only(malloc_fun_t f) {
    t_malloc = f;
}

void
toku_set_func_realloc(realloc_fun_t f) {
    t_realloc = f;
    t_xrealloc = f;
}

void
toku_set_func_xrealloc_only(realloc_fun_t f) {
    t_xrealloc = f;
}

void
toku_set_func_realloc_only(realloc_fun_t f) {
    t_realloc = f;

}

void
toku_set_func_free(free_fun_t f) {
    t_free = f;
}