17 files changed, 436 insertions, 94 deletions

diff --git a/runtime/ChangeLog b/runtime/ChangeLog
index 8fa11695..7536cc88 100644
--- a/runtime/ChangeLog
+++ b/runtime/ChangeLog
@@ -1,3 +1,13 @@
+2008-01-14  Martin Hunt  <hunt@redhat.com>
+
+	* print.c (_stp_print_kernel_info): New function.
+	(all):  Call stp memory functions.
+
+	* alloc.c: Rewrite to track allocated memory if 
+	DEBUG_MEM is defined.
+	* counter.c, map.c, stat.c, time.c: Call stp
+	memory functions.
+	
 2007-11-14  Zhaolei  <zhaolei@cn.fujitsu.com>
 
 	From Cai Fei <caifei@cn.fujitsu.com>
diff --git a/runtime/alloc.c b/runtime/alloc.c
index 8f2a7451..71b592a7 100644
--- a/runtime/alloc.c
+++ b/runtime/alloc.c
@@ -1,6 +1,6 @@
 /* -*- linux-c -*- 
  * Memory allocation functions
- * Copyright (C) 2005, 2006, 2007 Red Hat Inc.
+ * Copyright (C) 2005-2008 Red Hat Inc.
  *
  * This file is part of systemtap, and is free software.  You can
  * redistribute it and/or modify it under the terms of the GNU General
@@ -11,55 +11,356 @@
 #ifndef _ALLOC_C_
 #define _ALLOC_C_
 
-/* counters of how much memory has been allocated */
-static int _stp_allocated_memory = 0;
 static int _stp_allocated_net_memory = 0;
-
 #define STP_ALLOC_FLAGS (GFP_KERNEL | __GFP_NORETRY | __GFP_NOWARN)
 
+#define DEBUG_MEM
+/*
+ * If DEBUG_MEM is defined (stap -DDEBUG_MEM ...) then full memory
+ * tracking is used. Each allocation is recorded and matched with 
+ * a free. Also a fence is set around the allocated memory so overflows
+ * and underflows can be detected. Errors are written to the system log
+ * with printk.
+ *
+ * NOTE: if youy system is slow or your script makes a very large number
+ * of allocations, you may get a warning in the system log:
+ * BUG: soft lockup - CPU#1 stuck for 11s! [staprun:28269]
+ * This is an expected side-effect of the overhead of tracking, especially
+ * with a simple linked list of allocations. Optimization
+ * would be nice, but DEBUG_MEM is only for testing.
+ */
+
+#ifdef DEBUG_MEM
+
+static spinlock_t _stp_mem_lock = SPIN_LOCK_UNLOCKED;
+static int _stp_allocated_memory = 0;
+
+#define MEM_MAGIC 0xc11cf77f
+#define MEM_FENCE_SIZE 32
+
+enum _stp_memtype { MEM_KMALLOC, MEM_VMALLOC, MEM_PERCPU };
+
+typedef struct {
+	enum _stp_memtype type;
+	char *alloc;
+	char *free;
+} _stp_malloc_type;
+
+static const _stp_malloc_type const _stp_malloc_types[] = {
+	{MEM_KMALLOC, "kmalloc", "kfree"},
+	{MEM_VMALLOC, "vmalloc", "vfree"},
+	{MEM_PERCPU, "alloc_percpu", "free_percpu"}
+};
+
+struct _stp_mem_entry {
+	struct list_head list;
+	int32_t magic;
+	enum _stp_memtype type;
+	size_t len;
+	void *addr;
+};
+
+#define MEM_DEBUG_SIZE (2*MEM_FENCE_SIZE+sizeof(struct _stp_mem_entry))
+
+static LIST_HEAD(_stp_mem_list);
+
+void _stp_check_mem_fence (char *addr, int size)
+{
+	char *ptr;
+	int i;
+
+	ptr = addr - MEM_FENCE_SIZE;
+	while (ptr < addr) {
+		if (*ptr != 0x55) {
+			printk("SYSTEMTAP ERROR: Memory fence corrupted before allocated memory\n");
+			printk("at addr %p. (Allocation starts at %p)", ptr, addr);
+			return;
+		}
+		ptr++;
+	}
+	ptr = addr + size;
+	while (ptr < addr + size + MEM_FENCE_SIZE) {
+		if (*ptr != 0x55) {
+			printk("SYSTEMTAP ERROR: Memory fence corrupted after allocated memory\n");
+			printk("at addr %p. (Allocation ends at %p)", ptr, addr + size - 1);
+			return;
+		}
+		ptr++;
+	}
+}
+
+void *_stp_mem_debug_setup(void *addr, size_t size, enum _stp_memtype type)
+{
+	struct list_head *p;
+	struct _stp_mem_entry *m;
+	memset(addr, 0x55, MEM_FENCE_SIZE);
+	addr += MEM_FENCE_SIZE;
+	memset(addr + size, 0x55, MEM_FENCE_SIZE);
+	p = (struct list_head *)(addr + size + MEM_FENCE_SIZE);
+	m = (struct _stp_mem_entry *)p;
+	m->magic = MEM_MAGIC;
+	m->type = type;
+	m->len = size;
+	m->addr = addr;
+	spin_lock(&_stp_mem_lock);
+	list_add(p, &_stp_mem_list); 
+	spin_unlock(&_stp_mem_lock);
+	return addr;
+}
+
+/* Percpu allocations don't have the fence. Implementing it is problematic. */
+void _stp_mem_debug_percpu(struct _stp_mem_entry *m, void *addr, size_t size)
+{
+	struct list_head *p = (struct list_head *)m;
+	m->magic = MEM_MAGIC;
+	m->type = MEM_PERCPU;
+	m->len = size;
+	m->addr = addr;
+	spin_lock(&_stp_mem_lock);
+	list_add(p, &_stp_mem_list);
+	spin_unlock(&_stp_mem_lock);	
+}
+
+void _stp_mem_debug_free(void *addr, enum _stp_memtype type)
+{
+	int found = 0;
+	struct list_head *p, *tmp;
+	struct _stp_mem_entry *m = NULL;
+
+	spin_lock(&_stp_mem_lock);
+	list_for_each_safe(p, tmp, &_stp_mem_list) {
+		m = list_entry(p, struct _stp_mem_entry, list);
+		if (m->addr == addr) {
+			list_del(p);
+			found = 1;
+			break;
+		}
+	}
+	spin_unlock(&_stp_mem_lock);
+	if (!found) {
+		printk("SYSTEMTAP ERROR: Free of unallocated memory %p type=%s\n", 
+		       addr, _stp_malloc_types[type].free);
+		return;
+	}
+	if (m->magic != MEM_MAGIC) {
+		printk("SYSTEMTAP ERROR: Memory at %p corrupted!!\n", addr);
+		return;
+	}
+	if (m->type != type) {
+		printk("SYSTEMTAP ERROR: Memory allocated with %s and freed with %s\n",
+		       _stp_malloc_types[m->type].alloc, 		       
+		       _stp_malloc_types[type].free);
+	}
+	
+	switch (m->type) {
+	case MEM_KMALLOC:
+		_stp_check_mem_fence(addr, m->len);
+		kfree(addr - MEM_FENCE_SIZE);
+		break;
+	case MEM_PERCPU:
+		free_percpu(addr);
+		kfree(p);
+		break;
+	case MEM_VMALLOC:
+		_stp_check_mem_fence(addr, m->len);
+		vfree(addr - MEM_FENCE_SIZE);		
+		break;
+	default:
+		printk("SYSTEMTAP ERROR: Attempted to free memory at addr %p len=%d with unknown allocation type.\n", addr, (int)m->len);
+	}
+
+	return;
+}
+#endif
+
 static void *_stp_kmalloc(size_t size)
 {
-	void *ret = kmalloc(size, STP_ALLOC_FLAGS);
-	if (ret)
+#ifdef DEBUG_MEM
+	void *ret = kmalloc(size + MEM_DEBUG_SIZE, STP_ALLOC_FLAGS);
+	if (likely(ret)) {
+		ret = _stp_mem_debug_setup(ret, size, MEM_KMALLOC);
 		_stp_allocated_memory += size;
+	}
 	return ret;
+#else
+	return kmalloc(size, STP_ALLOC_FLAGS);
+#endif
 }
 
 static void *_stp_kzalloc(size_t size)
 #if LINUX_VERSION_CODE < KERNEL_VERSION(2,6,15)
 {
-	void *ret = kmalloc(size, STP_ALLOC_FLAGS);
-	if (ret) {
+#ifdef DEBUG_MEM
+	void *ret = kmalloc(size + MEM_DEBUG_SIZE, STP_ALLOC_FLAGS);
+	if (likely(ret)) {
+		ret = _stp_mem_debug_setup(ret, size, MEM_KMALLOC);
 		memset (ret, 0, size);
 		_stp_allocated_memory += size;
 	}
+#else
+	void *ret = kmalloc(size, STP_ALLOC_FLAGS);
+	if (likely(ret))
+		memset (ret, 0, size);
+#endif /* DEBUG_MEM */
 	return ret;
 }
-#else
+#else /* LINUX_VERSION_CODE >= KERNEL_VERSION(2,6,15) */
 {
-	void *ret = kzalloc(size, STP_ALLOC_FLAGS);
-	if (ret)
+#ifdef DEBUG_MEM
+	void *ret = kzalloc(size + MEM_DEBUG_SIZE, STP_ALLOC_FLAGS);
+	if (likely(ret)) {
+		ret = _stp_mem_debug_setup(ret, size, MEM_KMALLOC);
 		_stp_allocated_memory += size;
+	}
 	return ret;
-}
+#else
+	return kzalloc(size, STP_ALLOC_FLAGS);
 #endif
+}
+#endif /* LINUX_VERSION_CODE < KERNEL_VERSION(2,6,15) */
 
 static void *_stp_vmalloc(unsigned long size)
 {
-	void *ret = __vmalloc(size, STP_ALLOC_FLAGS, PAGE_KERNEL);
-	if (ret)
+#ifdef DEBUG_MEM
+	void *ret = __vmalloc(size + MEM_DEBUG_SIZE, STP_ALLOC_FLAGS, PAGE_KERNEL);
+	if (likely(ret)) {
+		ret = _stp_mem_debug_setup(ret, size, MEM_VMALLOC);
 		_stp_allocated_memory += size;
+	}
 	return ret;
+#else
+	return __vmalloc(size, STP_ALLOC_FLAGS, PAGE_KERNEL);
+#endif
+
 }
 
 #if LINUX_VERSION_CODE < KERNEL_VERSION(2,6,15)
-#define _stp_alloc_percpu(size) __alloc_percpu(size, 8)
+static void *_stp_alloc_percpu(size_t size)
+{
+#ifdef DEBUG_MEM
+	void *ret = __alloc_percpu(size, 8);
+	if (likely(ret)) {
+		struct _stp_mem_entry *m = kmalloc(sizeof(struct _stp_mem_entry), STP_ALLOC_FLAGS);
+		if (unlikely(m == NULL)) {
+			free_percpu(ret);
+			return NULL;
+		}
+		_stp_mem_debug_percpu(m, ret, size);
+		_stp_allocated_memory += size * num_online_cpus();
+	}
+	return ret;
 #else
-#define _stp_alloc_percpu(size) __alloc_percpu(size)
+	return __alloc_percpu(size, 8);
 #endif
+}
+#else
+static void *_stp_alloc_percpu(size_t size)
+{
+#ifdef DEBUG_MEM
+	void *ret = __alloc_percpu(size);
+	if (likely(ret)) {
+		struct _stp_mem_entry *m = kmalloc(sizeof(struct _stp_mem_entry), STP_ALLOC_FLAGS);
+		if (unlikely(m == NULL)) {
+			free_percpu(ret);
+			return NULL;
+		}
+		_stp_mem_debug_percpu(m, ret, size);
+		_stp_allocated_memory += size * num_online_cpus();
+	}
+	return ret;
+#else
+	return __alloc_percpu(size);
+#endif
+}
+#endif /* LINUX_VERSION_CODE */
 
 #if LINUX_VERSION_CODE < KERNEL_VERSION(2,6,12)
-#define kmalloc_node(size,flags,node) kmalloc(size,flags)
+#define _stp_kmalloc_node(size,node) _stp_kmalloc(size)
+#else
+static void *_stp_kmalloc_node(size_t size, int node)
+{
+#ifdef DEBUG_MEM
+	void *ret = kmalloc_node(size + MEM_DEBUG_SIZE, STP_ALLOC_FLAGS, node);
+	if (likely(ret)) {
+		ret = _stp_mem_debug_setup(ret, size, MEM_KMALLOC);
+		_stp_allocated_memory += size;
+	}
+	return ret;
+#else
+	return kmalloc_node(size, STP_ALLOC_FLAGS, node);
+#endif
+}
 #endif /* LINUX_VERSION_CODE */
 
+void _stp_kfree(void *addr)
+{
+#ifdef DEBUG_MEM
+	_stp_mem_debug_free(addr, MEM_KMALLOC);
+#else
+	kfree(addr);
+#endif
+}
+
+void _stp_vfree(void *addr)
+{
+#ifdef DEBUG_MEM
+	_stp_mem_debug_free(addr, MEM_VMALLOC);
+#else
+	vfree(addr);
+#endif
+}
+
+void _stp_free_percpu(void *addr)
+{
+#ifdef DEBUG_MEM
+	_stp_mem_debug_free(addr, MEM_PERCPU);
+#else
+	free_percpu(addr);
+#endif
+}
+
+void _stp_mem_debug_done(void)
+{
+#ifdef DEBUG_MEM
+	struct list_head *p, *tmp;
+	struct _stp_mem_entry *m;
+
+	spin_lock(&_stp_mem_lock);
+	list_for_each_safe(p, tmp, &_stp_mem_list) {
+		m = list_entry(p, struct _stp_mem_entry, list);
+		list_del(p);
+
+		printk("SYSTEMTAP ERROR: Memory %p len=%d allocation type: %s. Not freed.\n", 
+		       m->addr, (int)m->len, _stp_malloc_types[m->type].alloc);
+
+		if (m->magic != MEM_MAGIC) {
+			printk("SYSTEMTAP ERROR: Memory at %p len=%d corrupted!!\n", m->addr, (int)m->len);
+			/* Don't free. Too dangerous */
+			goto done;
+		}
+
+		switch (m->type) {
+		case MEM_KMALLOC:
+			_stp_check_mem_fence(m->addr, m->len);
+			kfree(m->addr - MEM_FENCE_SIZE);
+			break;
+		case MEM_PERCPU:
+			free_percpu(m->addr);
+			kfree(p);
+			break;
+		case MEM_VMALLOC:
+			_stp_check_mem_fence(m->addr, m->len);
+			vfree(m->addr - MEM_FENCE_SIZE);		
+			break;
+		default:
+			printk("SYSTEMTAP ERROR: Attempted to free memory at addr %p len=%d with unknown allocation type.\n", m->addr, (int)m->len);
+		}
+	}
+done:
+	spin_unlock(&_stp_mem_lock);
+
+	return;
+
+#endif
+}
 #endif /* _ALLOC_C_ */
diff --git a/runtime/counter.c b/runtime/counter.c
index def734f6..42555dfa 100644
--- a/runtime/counter.c
+++ b/runtime/counter.c
@@ -54,7 +54,7 @@ typedef struct _counter *Counter;
  */
 Counter _stp_counter_init (void)
 {
-	Counter cnt = alloc_percpu (struct _counter);
+	Counter cnt = _stp_alloc_percpu (struct _counter);
 #if NEED_COUNTER_LOCKS == 1
 	{
 		int i;
@@ -135,7 +135,7 @@ int64_t _stp_counter_get (Counter cnt, int clear)
  */
 void _stp_counter_free (Counter cnt)
 {
-	free_percpu (cnt);
+	_stp_free_percpu (cnt);
 }
 
 /** @} */
diff --git a/runtime/map.c b/runtime/map.c
index 2acf2fe2..70990876 100644
--- a/runtime/map.c
+++ b/runtime/map.c
@@ -205,18 +205,13 @@ static int _stp_map_init(MAP m, unsigned max_entries, int type, int key_size, in
 		
 		for (i = 0; i < max_entries; i++) {
 			if (cpu < 0)
-				tmp = kmalloc(size, STP_ALLOC_FLAGS);
+				tmp = _stp_kmalloc(size);
 			else
-				tmp = kmalloc_node(size, STP_ALLOC_FLAGS, cpu_to_node(cpu));
+				tmp = _stp_kmalloc_node(size, cpu_to_node(cpu));
 		
 			if (!tmp)
 				return -1;
 			
-			if (cpu < 0)
-				_stp_allocated_memory += size;
-			else
-				_stp_allocated_memory += size * num_online_cpus();
-				
 //			dbug ("allocated %lx\n", (long)tmp);
 			list_add((struct list_head *)tmp, &m->pool);
 			((struct map_node *)tmp)->map = m;
@@ -252,10 +247,9 @@ static PMAP _stp_pmap_new(unsigned max_entries, int type, int key_size, int data
 	if (pmap == NULL)
 		return NULL;
 
-	pmap->map = map = (MAP) alloc_percpu (struct map_root);
+	pmap->map = map = (MAP) _stp_alloc_percpu (sizeof(struct map_root));
 	if (map == NULL) 
 		goto err;
-	_stp_allocated_memory += sizeof(struct map_root) * num_online_cpus();
 
 	/* initialize the memory lists first so if allocations fail */
         /* at some point, it is easy to clean up. */
@@ -284,9 +278,9 @@ err1:
 		m = per_cpu_ptr (map, i);
 		__stp_map_del(m);
 	}
-	free_percpu(map);
+	_stp_free_percpu(map);
 err:
-	kfree(pmap);
+	_stp_kfree(pmap);
 	return NULL;
 }
 
@@ -387,13 +381,13 @@ static void __stp_map_del(MAP map)
 	/* free unused pool */
 	list_for_each_safe(p, tmp, &map->pool) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 
 	/* free used list */
 	list_for_each_safe(p, tmp, &map->head) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 }
 
@@ -409,7 +403,7 @@ void _stp_map_del(MAP map)
 
 	__stp_map_del(map);
 
-	kfree(map);
+	_stp_kfree(map);
 }
 
 void _stp_pmap_del(PMAP pmap)
@@ -423,12 +417,12 @@ void _stp_pmap_del(PMAP pmap)
 		MAP m = per_cpu_ptr (pmap->map, i);
 		__stp_map_del(m);
 	}
-	free_percpu(pmap->map);
+	_stp_free_percpu(pmap->map);
 
 	/* free agg map elements */
 	__stp_map_del(&pmap->agg);
 	
-	kfree(pmap);
+	_stp_kfree(pmap);
 }
 
 /* sort keynum values */
diff --git a/runtime/print.c b/runtime/print.c
index a451f622..0442ba09 100644
--- a/runtime/print.c
+++ b/runtime/print.c
@@ -1,6 +1,6 @@
 /* -*- linux-c -*- 
  * Print Functions
- * Copyright (C) 2007 Red Hat Inc.
+ * Copyright (C) 2007-2008 Red Hat Inc.
  *
  * This file is part of systemtap, and is free software.  You can
  * redistribute it and/or modify it under the terms of the GNU General
@@ -50,26 +50,25 @@ void *Stp_lbuf = NULL;
 /* create percpu print and io buffers */
 int _stp_print_init (void)
 {
-	Stp_pbuf = alloc_percpu(_stp_pbuf);
+	Stp_pbuf = _stp_alloc_percpu(sizeof(_stp_pbuf));
 	if (unlikely(Stp_pbuf == 0))
 		return -1;
 
 	/* now initialize IO buffer used in io.c */
-	Stp_lbuf = alloc_percpu(_stp_lbuf);
+	Stp_lbuf = _stp_alloc_percpu(sizeof(_stp_lbuf));
 	if (unlikely(Stp_lbuf == 0)) {
-		free_percpu(Stp_pbuf);
+		_stp_free_percpu(Stp_pbuf);
 		return -1;
 	}
-	_stp_allocated_memory += (sizeof(_stp_pbuf)+sizeof(_stp_lbuf)) * num_online_cpus();
 	return 0;
 }
 
 void _stp_print_cleanup (void)
 {
 	if (Stp_pbuf)
-		free_percpu(Stp_pbuf);
+		_stp_free_percpu(Stp_pbuf);
 	if (Stp_lbuf)
-		free_percpu(Stp_lbuf);
+		_stp_free_percpu(Stp_lbuf);
 }
 
 #define __DEF_EXPORT_FN(fn, postfix) fn ## _ ## postfix
@@ -276,5 +275,31 @@ static char *next_fmt(char *fmt, int *num)
 	return f;
 }
 
+void _stp_print_kernel_info(char *vstr, int ctx, int num_probes)
+{
+#ifdef DEBUG_MEM
+	printk(KERN_DEBUG "%s: systemtap: %s, base: %p, memory: %lu+%lu+%u+%u+%u data+text+ctx+net+alloc, probes: %d\n",
+	       THIS_MODULE->name,
+	       vstr, 
+	       THIS_MODULE->module_core,  
+	       (unsigned long) (THIS_MODULE->core_size - THIS_MODULE->core_text_size),
+               (unsigned long) THIS_MODULE->core_text_size,
+	       ctx,
+	       _stp_allocated_net_memory,
+	       _stp_allocated_memory - _stp_allocated_net_memory,
+		num_probes);
+#else
+	printk(KERN_DEBUG "%s: systemtap: %s, base: %p, memory: %lu+%lu+%u+%u data+text+ctx+net, probes: %d\n",
+	       THIS_MODULE->name,
+	       vstr, 
+	       THIS_MODULE->module_core,  
+	       (unsigned long) (THIS_MODULE->core_size - THIS_MODULE->core_text_size),
+               (unsigned long) THIS_MODULE->core_text_size,
+	       ctx,
+	       _stp_allocated_net_memory,
+	       num_probes);
+#endif
+}
+
 /** @} */
 #endif /* _PRINT_C_ */
diff --git a/runtime/stack-i386.c b/runtime/stack-i386.c
index a80de048..3675fc97 100644
--- a/runtime/stack-i386.c
+++ b/runtime/stack-i386.c
@@ -1,6 +1,6 @@
 /* -*- linux-c -*-
  * i386 stack tracing functions
- * Copyright (C) 2005, 2006, 2007 Red Hat Inc.
+ * Copyright (C) 2005-2008 Red Hat Inc.
  *
  * This file is part of systemtap, and is free software.  You can
  * redistribute it and/or modify it under the terms of the GNU General
@@ -37,7 +37,7 @@ static void __stp_stack_print (struct pt_regs *regs, int verbose, int levels)
 #else
 	while (_stp_valid_stack_ptr(context, (unsigned long)stack)) {
 		addr = *stack++;
-		_stp_func_print(addr, verbose);
+		_stp_func_print(addr, verbose, 1);
 	}
 #endif
 }
diff --git a/runtime/stack-x86_64.c b/runtime/stack-x86_64.c
index 16e3bf7c..186b2ad4 100644
--- a/runtime/stack-x86_64.c
+++ b/runtime/stack-x86_64.c
@@ -15,6 +15,6 @@ static void __stp_stack_print (struct pt_regs *regs, int verbose, int levels)
 
 	while ((long)stack & (THREAD_SIZE-1)) {
 		addr = *stack++;
-		_stp_func_print(addr, verbose);
+		_stp_func_print(addr, verbose, 1);
 	}
 }
diff --git a/runtime/stat.c b/runtime/stat.c
index 3df1d063..298a6653 100644
--- a/runtime/stat.c
+++ b/runtime/stat.c
@@ -52,7 +52,7 @@
 /** Stat struct for stat.c. Maps do not need this */
 struct _Stat {
 	struct _Hist hist;
-	/* per-cpu data. allocated with alloc_percpu() */
+	/* per-cpu data. allocated with _stp_alloc_percpu() */
 	stat *sd;
 	/* aggregated data */   
 	stat *agg;  
@@ -130,9 +130,9 @@ Stat _stp_stat_init (int type, ...)
 	return st;
 
 exit2:
-	kfree (sd);
+	_stp_kfree (sd);
 exit1:
-	kfree (st);
+	_stp_kfree (st);
 	return NULL;
 }
 
@@ -144,9 +144,9 @@ exit1:
 void _stp_stat_del (Stat st)
 {
 	if (st) {
-		free_percpu (st->sd);
-		kfree (st->agg);
-		kfree (st);
+		_stp_free_percpu (st->sd);
+		_stp_kfree (st->agg);
+		_stp_kfree (st);
 	}
 }
 	
diff --git a/runtime/sym.c b/runtime/sym.c
index c40f48db..56c93064 100644
--- a/runtime/sym.c
+++ b/runtime/sym.c
@@ -193,7 +193,7 @@ void _stp_symbol_print (unsigned long address)
 	_stp_printf ("%p", (int64_t)address);
 
 	if (name) {		
-		if (modname)
+		if (modname && *modname)
 			_stp_printf (" : %s+%#lx/%#lx [%s]", name, offset, size, modname);
 		else
 			_stp_printf (" : %s+%#lx/%#lx", name, offset, size);
@@ -202,22 +202,28 @@ void _stp_symbol_print (unsigned long address)
 
 /* Like _stp_symbol_print, except only print if the address is a valid function address */
 
-void _stp_func_print (unsigned long address, int verbose)
+void _stp_func_print (unsigned long address, int verbose, int exact)
 { 
 	char *modname;
         const char *name;
         unsigned long offset, size;
+	char *exstr;
+
+	if (exact)
+		exstr = "";
+	else
+		exstr = " (inexact)";
 
         name = _stp_kallsyms_lookup(address, &size, &offset, &modname, NULL);
 
 	if (name) {
 		if (verbose) {
-			if (modname)
-				_stp_printf (" %p : %s+%#lx/%#lx [%s]\n", 
-					     (int64_t)address, name, offset, size, modname);
+			if (modname && *modname)
+				_stp_printf (" %p : %s+%#lx/%#lx [%s]%s\n", 
+					     (int64_t)address, name, offset, size, modname, exstr);
 			else
-				_stp_printf (" %p : %s+%#lx/%#lx\n", 
-					     (int64_t)address, name, offset, size);
+				_stp_printf (" %p : %s+%#lx/%#lx%s\n", 
+					     (int64_t)address, name, offset, size, exstr);
 		} else 
 			_stp_printf ("%p ", (int64_t)address);
 	}
diff --git a/runtime/time.c b/runtime/time.c
index 688427f4..52a2edbb 100644
--- a/runtime/time.c
+++ b/runtime/time.c
@@ -219,7 +219,7 @@ _stp_kill_time(void)
         }
 #endif
 
-        free_percpu(stp_time);
+        _stp_free_percpu(stp_time);
     }
 }
 
@@ -229,10 +229,9 @@ _stp_init_time(void)
 {
     int ret = 0;
 
-    stp_time = alloc_percpu(stp_time_t);
+    stp_time = _stp_alloc_percpu(sizeof(stp_time_t));
     if (unlikely(stp_time == 0))
 	    return -1;
-    _stp_allocated_memory += sizeof(stp_time_t) * num_online_cpus();
     
     stp_timer_reregister = 1;
     ret = on_each_cpu(__stp_init_time, NULL, 0, 1);
diff --git a/runtime/transport/ChangeLog b/runtime/transport/ChangeLog
index 263a6bf8..26268e66 100644
--- a/runtime/transport/ChangeLog
+++ b/runtime/transport/ChangeLog
@@ -1,3 +1,9 @@
+2008-01-15  Martin Hunt  <hunt@redhat.com>
+
+	Support for DEBUG_MEM
+	* transport.c (_stp_transport): Call stp_mem_debug_done();
+        * (*.c):  Call stp malloc and free functions.
+
 2007-11-09  Masami Hiramatsu  <mhiramat@redhat.com>
 
 	PR3858
diff --git a/runtime/transport/control.c b/runtime/transport/control.c
index 7846572a..3cfeca1e 100644
--- a/runtime/transport/control.c
+++ b/runtime/transport/control.c
@@ -1,7 +1,7 @@
 /* -*- linux-c -*-
  *
  * debugfs control channel
- * Copyright (C) 2007 Red Hat Inc.
+ * Copyright (C) 2007, 2008 Red Hat Inc.
  *
  * This file is part of systemtap, and is free software.  You can
  * redistribute it and/or modify it under the terms of the GNU General
@@ -422,7 +422,7 @@ static int _stp_register_ctl_channel (void)
 
 	/* allocate buffers */
 	for (i = 0; i < STP_DEFAULT_BUFFERS; i++) {
-		p = (struct list_head *)kmalloc(sizeof(struct _stp_buffer),STP_ALLOC_FLAGS);
+		p = (struct list_head *)_stp_kmalloc(sizeof(struct _stp_buffer));
 		// printk("allocated buffer at %lx\n", (long)p);
 		if (!p)
 			goto err0;
@@ -448,7 +448,7 @@ err0:
 
 	list_for_each_safe(p, tmp, &_stp_pool_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 	errk ("Error creating systemtap debugfs entries.\n");
 	return -1;
@@ -464,15 +464,15 @@ static void _stp_unregister_ctl_channel (void)
 	/* free memory pools */
 	list_for_each_safe(p, tmp, &_stp_pool_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 	list_for_each_safe(p, tmp, &_stp_sym_ready_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 	list_for_each_safe(p, tmp, &_stp_ctl_ready_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 }
 
diff --git a/runtime/transport/procfs.c b/runtime/transport/procfs.c
index 069e379e..5d40cae4 100644
--- a/runtime/transport/procfs.c
+++ b/runtime/transport/procfs.c
@@ -1,7 +1,7 @@
 /* -*- linux-c -*-
  *
  * /proc transport and control
- * Copyright (C) 2005-2007 Red Hat Inc.
+ * Copyright (C) 2005-2008 Red Hat Inc.
  *
  * This file is part of systemtap, and is free software.  You can
  * redistribute it and/or modify it under the terms of the GNU General
@@ -485,7 +485,7 @@ static int _stp_set_buffers(int num)
 	
 	if (num > _stp_current_buffers) {
 		for (i = 0; i < num - _stp_current_buffers; i++) {
-			p = (struct list_head *)kmalloc(sizeof(struct _stp_buffer),STP_ALLOC_FLAGS);
+			p = (struct list_head *)_stp_kmalloc(sizeof(struct _stp_buffer));
 			if (!p)	{
 				_stp_current_buffers += i;
 				goto err;
@@ -501,7 +501,7 @@ static int _stp_set_buffers(int num)
 			p = _stp_pool_q.next;
 			list_del(p);
 			spin_unlock_irqrestore(&_stp_pool_lock, flags);
-			kfree(p);
+			_stp_kfree(p);
 		}
 	}
 	_stp_current_buffers = num;
@@ -542,7 +542,7 @@ static int _stp_register_ctl_channel (void)
 
 	/* allocate buffers */
 	for (i = 0; i < STP_DEFAULT_BUFFERS; i++) {
-		p = (struct list_head *)kmalloc(sizeof(struct _stp_buffer),STP_ALLOC_FLAGS);
+		p = (struct list_head *)_stp_kmalloc(sizeof(struct _stp_buffer));
 		// printk("allocated buffer at %lx\n", (long)p);
 		if (!p)
 			goto err0;
@@ -593,7 +593,7 @@ err2:
 err1:
 #ifdef STP_BULKMODE
 	for (de = _stp_proc_root->subdir; de; de = de->next)
-		kfree (de->data);
+		_stp_kfree (de->data);
 	for_each_cpu(j) {
 		if (j == i)
 			break;
@@ -607,7 +607,7 @@ err1:
 err0:
 	list_for_each_safe(p, tmp, &_stp_pool_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 
 	errk ("Error creating systemtap /proc entries.\n");
@@ -624,7 +624,7 @@ static void _stp_unregister_ctl_channel (void)
 	struct proc_dir_entry *de;
 	kbug("unregistering procfs\n");
 	for (de = _stp_proc_root->subdir; de; de = de->next)
-		kfree (de->data);
+		_stp_kfree (de->data);
 
 	for_each_cpu(i) {
 		sprintf(buf, "%d", i);
@@ -640,15 +640,15 @@ static void _stp_unregister_ctl_channel (void)
 	/* free memory pools */
 	list_for_each_safe(p, tmp, &_stp_pool_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 	list_for_each_safe(p, tmp, &_stp_sym_ready_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 	list_for_each_safe(p, tmp, &_stp_ctl_ready_q) {
 		list_del(p);
-		kfree(p);
+		_stp_kfree(p);
 	}
 }
 
diff --git a/runtime/transport/relayfs.c b/runtime/transport/relayfs.c
index 2c7ca754..375c8e59 100644
--- a/runtime/transport/relayfs.c
+++ b/runtime/transport/relayfs.c
@@ -118,7 +118,7 @@ err:
 		_stp_remove_relay_dir(utt->dir);
 	if (utt->utt_tree_root)
 		_stp_remove_relay_root(utt->utt_tree_root);
-	kfree(utt);
+	_stp_kfree(utt);
 	return NULL;
 }
 
@@ -173,7 +173,7 @@ int utt_trace_remove(struct utt_trace *utt)
 			_stp_remove_relay_dir(utt->dir);
 		if (utt->utt_tree_root)
 			_stp_remove_relay_root(utt->utt_tree_root);
-		kfree(utt);
+		_stp_kfree(utt);
 	}
 	return 0;
 }
diff --git a/runtime/transport/symbols.c b/runtime/transport/symbols.c
index 1d930183..a3ab5807 100644
--- a/runtime/transport/symbols.c
+++ b/runtime/transport/symbols.c
@@ -90,10 +90,10 @@ static struct _stp_module * _stp_alloc_module(unsigned num, unsigned datasize)
 bad:
 	if (mod) {
 		if (mod->allocated && mod->symbols)
-			vfree(mod->symbols);
+			_stp_vfree(mod->symbols);
 		else
-			kfree(mod->symbols);
-		kfree(mod); 
+			_stp_kfree(mod->symbols);
+		_stp_kfree(mod); 
 	}
 	return NULL;
 }
@@ -109,19 +109,19 @@ static void _stp_free_module(struct _stp_module *mod)
 	/* free symbol memory */
 	if (mod->num_symbols) {
 		if (mod->allocated & 1)
-			vfree(mod->symbols);
+			_stp_vfree(mod->symbols);
 		else
-			kfree(mod->symbols);
+			_stp_kfree(mod->symbols);
 		if (mod->allocated & 2)
-			vfree(mod->symbol_data);
+			_stp_vfree(mod->symbol_data);
 		else
-			kfree(mod->symbol_data);
+			_stp_kfree(mod->symbol_data);
 	}
 	if (mod->sections)
-		kfree(mod->sections);
+		_stp_kfree(mod->sections);
 
 	/* free module memory */
-	kfree(mod);
+	_stp_kfree(mod);
 }
 
 /* Delete a module and free its memory. */
@@ -455,7 +455,7 @@ static int _stp_do_module(const char __user *buf, int count)
 		return -EFAULT;
 	}
 	if (copy_from_user ((char *)tmpmod.sections, buf+sizeof(tmpmod), count-sizeof(tmpmod))) {
-		kfree(tmpmod.sections);
+		_stp_kfree(tmpmod.sections);
 		return -EFAULT;
 	}
 	for (i = 0; i < tmpmod.num_sections; i++) {
@@ -472,7 +472,7 @@ static int _stp_do_module(const char __user *buf, int count)
 	/* load symbols from tmpmod.module to mod */
 	mod = _stp_load_module_symbols(&tmpmod);	
 	if (mod == NULL) {
-		kfree(tmpmod.sections);
+		_stp_kfree(tmpmod.sections);
 		return 0;
 	}
 
diff --git a/runtime/transport/transport.c b/runtime/transport/transport.c
index 2904cf2f..0b18b275 100644
--- a/runtime/transport/transport.c
+++ b/runtime/transport/transport.c
@@ -2,7 +2,7 @@
  * transport.c - stp transport functions
  *
  * Copyright (C) IBM Corporation, 2005
- * Copyright (C) Red Hat Inc, 2005-2007
+ * Copyright (C) Red Hat Inc, 2005-2008
  * Copyright (C) Intel Corporation, 2006
  *
  * This file is part of systemtap, and is free software.  You can
@@ -215,6 +215,7 @@ void _stp_transport_close()
 	_stp_free_modules();
 	_stp_kill_time();
 	_stp_print_cleanup(); 	/* free print buffers */
+	_stp_mem_debug_done();
 	kbug("---- CLOSED ----\n");
 }
 
diff --git a/runtime/transport/utt.c b/runtime/transport/utt.c
index ac2e1513..182c1178 100644
--- a/runtime/transport/utt.c
+++ b/runtime/transport/utt.c
@@ -88,7 +88,7 @@ void utt_trace_cleanup(struct utt_trace *utt)
 	if (utt->dropped_file)
 		debugfs_remove(utt->dropped_file);
 	utt_remove_tree(utt);
-	kfree(utt);
+	_stp_kfree(utt);
 }
 
 int utt_trace_remove(struct utt_trace *utt)
@@ -207,7 +207,7 @@ struct utt_trace *utt_trace_setup(struct utt_trace_setup *utts)
 		goto err;
 
 	ret = -ENOMEM;
-	utt = kzalloc(sizeof(*utt), GFP_KERNEL);
+	utt = _stp_kzalloc(sizeof(*utt));
 	if (!utt)
 		goto err;
 
@@ -251,7 +251,7 @@ err:
 			debugfs_remove(utt->dropped_file);
 		if (utt->rchan)
 			relay_close(utt->rchan);
-		kfree(utt);
+		_stp_kfree(utt);
 	}
 	if (dir)
 		utt_remove_tree(utt);