Don't print stuff, only return strings.

Author: manishrjain

z/allocator.go

@@ -17,6 +17,7 @@
 package z
 
 import (
+	"bytes"
 	"fmt"
 	"math"
 	"math/bits"
@@ -92,18 +93,21 @@ func (a *Allocator) Reset() {
 	atomic.StoreUint64(&a.compIdx, 0)
 }
 
-func PrintAllocators() {
+func Allocators() string {
 	allocsMu.Lock()
 	tags := make(map[string]uint64)
 	num := make(map[string]int)
 	for _, ac := range allocs {
 		tags[ac.Tag] += ac.Allocated()
 		num[ac.Tag] += 1
 	}
+
+	var buf bytes.Buffer
 	for tag, sz := range tags {
-		fmt.Printf("Allocator Tag: %s Num: %d Size: %s\n", tag, num[tag], humanize.IBytes(sz))
+		fmt.Fprintf(&buf, "Tag: %s Num: %d Size: %s . ", tag, num[tag], humanize.IBytes(sz))
 	}
 	allocsMu.Unlock()
+	return buf.String()
 }
 
 func (a *Allocator) String() string {

z/calloc_jemalloc.go

@@ -13,12 +13,16 @@ package z
 */
 import "C"
 import (
+	"bytes"
 	"fmt"
 	"runtime"
+	"strconv"
 	"strings"
 	"sync"
 	"sync/atomic"
 	"unsafe"
+
+	"github.com/dustin/go-humanize"
 )
 
 // The go:linkname directives provides backdoor access to private functions in
@@ -49,6 +53,11 @@ type dalloc struct {
 var dallocsMu sync.Mutex
 var dallocs map[unsafe.Pointer]*dalloc
 
+func init() {
+	// By initializing dallocs, we can start tracking allocations and deallocations via z.Calloc.
+	dallocs = make(map[unsafe.Pointer]*dalloc)
+}
+
 func Calloc(n int) []byte {
 	if n == 0 {
 		return make([]byte, 0)
@@ -123,20 +132,25 @@ func Free(b []byte) {
 	}
 }
 
-func PrintLeaks() {
+func Leaks() string {
 	if dallocs == nil {
-		fmt.Println("Leak detection disabled. Enable with 'leak' build flag.")
-		return
+		return "Leak detection disabled. Enable with 'leak' build flag."
 	}
 	dallocsMu.Lock()
 	defer dallocsMu.Unlock()
 	if len(dallocs) == 0 {
-		fmt.Println("NO leaks found.")
-		return
+		return "NO leaks found."
 	}
+	m := make(map[string]int)
 	for _, da := range dallocs {
-		fmt.Printf("LEAK: %d at file: %s %d\n", da.sz, da.f, da.no)
+		m[da.f+":"+strconv.Itoa(da.no)] += da.sz
+	}
+	var buf bytes.Buffer
+	fmt.Fprintf(&buf, "Allocations:\n")
+	for f, sz := range m {
+		fmt.Fprintf(&buf, "%s at file: %s\n", humanize.IBytes(uint64(sz)), f)
 	}
+	return buf.String()
 }
 
 // ReadMemStats populates stats with JE Malloc statistics.

z/calloc_leak.go

@@ -27,7 +27,7 @@ func CallocNoRef(n int) []byte {
 // Free does not do anything in this mode.
 func Free(b []byte) {}
 
-func PrintLeaks() {}
+func Leaks() string { return "Leaks: Using Go memory" }
 func StatsPrint() {
 	fmt.Println("Using Go memory")
 }

z/calloc_nojemalloc.go

@@ -17,6 +17,7 @@
 package z
 
 import (
+	"fmt"
 	"sync"
 	"testing"
 	"time"
@@ -98,7 +99,7 @@ func TestCalloc(t *testing.T) {
 	// _ = buf2
 	Free(buf2)
 	require.Equal(t, int64(0), NumAllocBytes())
-	PrintLeaks()
+	fmt.Println(Leaks())
 
 	// Double free would panic when debug mode is enabled in jemalloc.
 	// Free(buf2)