diff --git a/lib/trace-cmd/include/trace-cmd-local.h b/lib/trace-cmd/include/trace-cmd-local.h
index d265e4c8..da9d8d58 100644
--- a/lib/trace-cmd/include/trace-cmd-local.h
+++ b/lib/trace-cmd/include/trace-cmd-local.h
@@ -12,6 +12,8 @@
 /* Can be overridden */
 void warning(const char *fmt, ...);
 
+void _lib_fatal(const char *fmt, ...);
+
 #define __packed __attribute__((packed))
 
 /* trace.dat file format version */
diff --git a/lib/trace-cmd/trace-input.c b/lib/trace-cmd/trace-input.c
index 2093a3dc..3687de09 100644
--- a/lib/trace-cmd/trace-input.c
+++ b/lib/trace-cmd/trace-input.c
@@ -16,6 +16,7 @@
 #include <linux/time64.h>
 
 #include "trace-write-local.h"
+#include "trace-cmd-local.h"
 #include "trace-local.h"
 #include "kbuffer.h"
 #include "list.h"
@@ -1087,8 +1088,10 @@ static void __free_page(struct tracecmd_input *handle, struct page *page)
 	struct page **pages;
 	int index;
 
-	if (!page->ref_count)
-		die("Page ref count is zero!\n");
+	if (!page->ref_count) {
+		_lib_fatal("Page ref count is zero!\n");
+		return;
+	}
 
 	page->ref_count--;
 	if (page->ref_count)
@@ -1146,16 +1149,20 @@ void tracecmd_free_record(struct tep_record *record)
 	if (!record)
 		return;
 
-	if (!record->ref_count)
-		die("record ref count is zero!");
+	if (!record->ref_count) {
+		_lib_fatal("record ref count is zero!");
+		return;
+	}
 
 	record->ref_count--;
 
 	if (record->ref_count)
 		return;
 
-	if (record->locked)
-		die("freeing record when it is locked!");
+	if (record->locked) {
+		_lib_fatal("freeing record when it is locked!");
+		return;
+	}
 
 	record->data = NULL;
 
@@ -1318,7 +1325,7 @@ static int get_page(struct tracecmd_input *handle, int cpu,
 
 	if (offset & (handle->page_size - 1)) {
 		errno = -EINVAL;
-		die("bad page offset %llx", offset);
+		_lib_fatal("bad page offset %llx", offset);
 		return -1;
 	}
 
@@ -1326,7 +1333,7 @@ static int get_page(struct tracecmd_input *handle, int cpu,
 	    offset > handle->cpu_data[cpu].file_offset +
 	    handle->cpu_data[cpu].file_size) {
 		errno = -EINVAL;
-		die("bad page offset %llx", offset);
+		_lib_fatal("bad page offset %llx", offset);
 		return -1;
 	}
 
@@ -1891,8 +1898,10 @@ tracecmd_peek_data(struct tracecmd_input *handle, int cpu)
 	if (handle->cpu_data[cpu].next) {
 
 		record = handle->cpu_data[cpu].next;
-		if (!record->data)
-			die("Something freed the record");
+		if (!record->data) {
+			_lib_fatal("Something freed the record");
+			return NULL;
+		}
 
 		if (handle->cpu_data[cpu].timestamp == record->ts)
 			return record;
diff --git a/lib/trace-cmd/trace-util.c b/lib/trace-cmd/trace-util.c
index 538adbc2..e8ed22fa 100644
--- a/lib/trace-cmd/trace-util.c
+++ b/lib/trace-cmd/trace-util.c
@@ -353,12 +353,12 @@ trace_load_plugins(struct tep_handle *tep, int flags)
 	return list;
 }
 
-void __noreturn __vdie(const char *fmt, va_list ap)
+int __vlib_fatal(const char *fmt, va_list ap)
 {
 	int ret = errno;
 
 	if (errno)
-		perror("trace-cmd");
+		perror("libtracecmd");
 	else
 		ret = -1;
 
@@ -366,36 +366,32 @@ void __noreturn __vdie(const char *fmt, va_list ap)
 	vfprintf(stderr, fmt, ap);
 
 	fprintf(stderr, "\n");
-	exit(ret);
+	return ret;
 }
 
-void __noreturn __die(const char *fmt, ...)
+#ifdef LIB_DIE
+
+void __noreturn _lib_fatal(const char *fmt, ...)
 {
 	va_list ap;
+	int ret;
 
 	va_start(ap, fmt);
-	__vdie(fmt, ap);
+	ret = __vlib_fatal(fmt, ap);
 	va_end(ap);
+	exit(ret);
 }
 
-void __weak __noreturn die(const char *fmt, ...)
+#else
+void _lib_fatal(const char *fmt, ...)
 {
 	va_list ap;
 
 	va_start(ap, fmt);
-	__vdie(fmt, ap);
+	__vlib_fatal(fmt, ap);
 	va_end(ap);
 }
-
-void __weak *malloc_or_die(unsigned int size)
-{
-	void *data;
-
-	data = malloc(size);
-	if (!data)
-		die("malloc");
-	return data;
-}
+#endif
 
 #define LOG_BUF_SIZE 1024
 static void __plog(const char *prefix, const char *fmt, va_list ap, FILE *fp)
@@ -546,8 +542,10 @@ int tracecmd_count_cpus(void)
 	pbuf = buf;
 
 	fp = fopen("/proc/cpuinfo", "r");
-	if (!fp)
-		die("Can not read cpuinfo");
+	if (!fp) {
+		_lib_fatal("Can not read cpuinfo");
+		return 0;
+	}
 
 	while ((r = getline(&pbuf, pn, fp)) >= 0) {
 		char *p;
